## matern.jl : implementation of Mat\'ern covariance function?

## Matern ##
struct Matern{T} <: IsotropicCovarianceStructure{T}
    λ::T
    ν::T
    σ::T
    p::T

    function Matern{T}(λ::T, ν::T, σ::T, p::T) where T
        λ > 0 || throw(DomainError(λ, "correlation length λ of Mat\u00E9rn covariance cannot be negative or zero"))
        ν > 0 || throw(DomainError(ν, "smoothness ν of Mat\u00E9rn covariance cannot be negative or zero"))
        σ > 0 || throw(DomainError(σ, "marginal standard deviation σ of Mat\u00E9rn covariance cannot be negative or zero"))
        p >= 1 || throw(DomainError(p, "in p-norm, p must be greater than or equal to 1"))
    	isinf(p) && throw(DomainError(p, "in p-norm, p cannot be infinity"))

        new{T}(λ, ν, σ, p)
    end
end

"""
    Matern(λ, ν, [σ = 1], [p = 2])

Mat\u00E9rn covariance structure with correlation length `λ`, smoothness `ν`, (optional) marginal standard deviation `σ` and (optional) `p`-norm, defined as

``C(x, y) = σ \\displaystyle\\frac{2^{1 - ν}}{Γ(ν)} \\left(\\frac{\sqrt{2\nu}ρ}{λ}\\right)^ν K_ν\\left(\\frac{\sqrt{2\nu}ρ}{λ}\\right)``

with ``ρ = ||x - y||_p``.

# Examples
```jldoctest
julia> Matern(0.1, 1.0)
Matérn (λ=0.1, ν=1.0, σ=1.0, p=2.0)

julia> Matern(1, 1, σ=2.0)
Matérn (λ=1.0, ν=1.0, σ=2.0, p=2.0)

```
See also: [`Exponential`](@ref), [`Linear`](@ref), [`Spherical`](@ref), [`Whittle`](@ref), [`Gaussian`](@ref), [`SquaredExponential`](@ref)
"""
Matern(λ::Real, ν::Real; σ::Real=1.0, p::Real=2) =
    Matern{promote_type(typeof(λ),typeof(ν),typeof(σ),typeof(p))}(promote(λ, ν, σ, p)...)

# TODO : Evaluation of the besselk function is extremely expensive; maybe try to catch special cases for ν = 1/2, 3/2 etc.
# evaluate Matern covariance
function apply(m::Matern, x::Real)
    if iszero(x)
        float(one(x))
    else
        2^(1 - m.ν) / gamma(m.ν) * (sqrt(2 * m.ν) * x / m.λ)^m.ν * besselk(m.ν, sqrt(2 * m.ν) * x / m.λ)
    end
end

# short name
shortname(::Matern) = "Mat\u00E9rn"


# Computing deviation in a variety of ways

# squared L2 distance
function sqL2dist{T<:Number}(a::AbstractArray{T}, b::AbstractArray{T})
    n = length(a)
    length(b) == n || throw(DimensionMismatch("Input dimension mismatch"))
    r = 0.0
    for i = 1:n
        @inbounds r += abs2(a[i] - b[i])
    end
    return r
end

# L2 distance
L2dist{T<:Number}(a::AbstractArray{T}, b::AbstractArray{T}) = sqrt(sqL2dist(a, b))

# L1 distance
function L1dist{T<:Number}(a::AbstractArray{T}, b::AbstractArray{T})
    n = length(a)
    length(b) == n || throw(DimensionMismatch("Input dimension mismatch"))
    r = 0.0
    for i = 1:n
        @inbounds r += abs(a[i] - b[i])
    end
    return r
end

# Linf distance
function Linfdist{T<:Number}(a::AbstractArray{T}, b::AbstractArray{T})
    n = length(a)
    length(b) == n || throw(DimensionMismatch("Input dimension mismatch"))
    r = 0.0
    for i = 1:n
        @inbounds v = abs(a[i] - b[i])
        if r < v
            r = v
        end
    end
    return r
end

# Generalized KL-divergence
function gkldiv{T<:FloatingPoint}(a::AbstractArray{T}, b::AbstractArray{T})
    n = length(a)
    r = 0.0
    for i = 1:n
        @inbounds ai = a[i]
        @inbounds bi = b[i]
        if ai > 0
            r += (ai * log(ai / bi) - ai + bi)
        else
            r += bi
        end
    end
    return r::Float64
end

# MeanAD: mean absolute deviation
meanad{T<:Number}(a::AbstractArray{T}, b::AbstractArray{T}) = L1dist(a, b) / length(a)

# MaxAD: maximum absolute deviation
maxad{T<:Number}(a::AbstractArray{T}, b::AbstractArray{T}) = Linfdist(a, b)

# MSD: mean squared deviation
msd{T<:Number}(a::AbstractArray{T}, b::AbstractArray{T}) = sqL2dist(a, b) / length(a)

# RMSD: root mean squared deviation
function rmsd{T<:Number}(a::AbstractArray{T}, b::AbstractArray{T}; normalize::Bool=false)
	v = sqrt(msd(a, b))
	if normalize
		amin, amax = extrema(a)
		v /= (amax - amin)
	end
	return v
end

# PSNR: peak signal-to-noise ratio
function psnr{T<:Real}(a::AbstractArray{T}, b::AbstractArray{T}, maxv::Real)
    20. * log10(maxv) - 10. * log10(msd(a, b))
end


"""
    Difference array

# Brief:
    For any array A with N numbers, we can create difference array D with N numbers also.
    Every k-th element of D is equal to difference between A[k] and A[k-1]
        D[k] = A[k-1] - A[k]

# Complexity of some operations
    - add some value x to the m consecutive elements of array - O(1)
    - print array after any numbers of changes                - O(N)

# Functions
    - create_diff_arr(original::Array{T}) - Create difference array for array 'original'
    - calculate_arr(diff_arr::Array{T}) - Create a original array from the given difference array
    - add_to_arr(diff_arr::Array{T}, l::Int, r::Int, x::Number) - Add x to all elements with index from [l, r]

# Contributed by: [Nikola Mircic](https://github.com/Nikola-Mircic)
"""
module DifferenceArray
# Create difference array for array 'original'
# Parameters:
#   - original - an array which is used for calculating the difference array
function create_diff_arr(original::Array{T}) where {T<:Number}
    n = length(original)

    diff_arr = Array(original)

    for i in 2:n
        diff_arr[i] = original[i] - original[i-1]
    end

    return diff_arr
end

# Create a original array from the given difference array
# Parameters:
#   - diff_arr - an difference array which is used for calculating the original array
function calculate_arr(diff_arr::Array{T}) where {T<:Number}
    n = length(diff_arr)

    arr = Array(diff_arr)

    for i in 2:n
        arr[i] = diff_arr[i] + arr[i-1]
    end

    return arr
end

# Add x to all elements with index from [l, r]
# Parameters:
#   - dif_arr - a difference array of the array you want to change
#   - l - leftmost index of the affected range
#   - r - rightmost index of the affected range
#   - x - a value to be added to all elements from a given range
function add_to_arr(
    diff_arr::Array{T},
    l::Int,
    r::Int,
    x::Number,
) where {T<:Number}
    diff_arr[l] += x
    if r < length(diff_arr)
        diff_arr[r+1] -= x
    end
end
end


# ------------------------------------------------------------------------------------------
# # An Intro to DifferentialEquations.jl
#
# ## Basic Introduction Via Ordinary Differential Equations
#
# This notebook will get you started with DifferentialEquations.jl by introducing you to the
# functionality for solving ordinary differential equations (ODEs). The corresponding
# documentation page is the [ODE
# tutorial](http://docs.juliadiffeq.org/latest/tutorials/ode_example.html). While some of
# the syntax may be different for other types of equations, the same general principles hold
# in each case. Our goal is to give a gentle and thorough introduction that highlights these
# principles in a way that will help you generalize what you have learned.
#
# ### Background
#
# If you are new to the study of differential equations, it can be helpful to do a quick
# background read on [the definition of ordinary differential
# equations](https://en.wikipedia.org/wiki/Ordinary_differential_equation). We define an
# ordinary differential equation as an equation which describes the way that a variable $u$
# changes, that is
#
# $$ u' = f(u,p,t) $$
#
# where $p$ are the parameters of the model, $t$ is the time variable, and $f$ is the
# nonlinear model of how $u$ changes. The initial value problem also includes the
# information about the starting value:
#
# $$ u(t_0) = u_0 $$
#
# Together, if you know the starting value and you know how the value will change with time,
# then you know what the value will be at any time point in the future. This is the
# intuitive definition of a differential equation.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ### First Model: Exponential Growth
#
# Our first model will be the canonical exponential growth model. This model says that the
# rate of change is proportional to the current value, and is this:
#
# $$ u' = au $$
#
# where we have a starting value $u(0)=u_0$. Let's say we put 1 dollar into Bitcoin which is
# increasing at a rate of $98\%$ per year. Then calling now $t=0$ and measuring time in
# years, our model is:
#
# $$ u' = 0.98u $$
#
# and $u(0) = 1.0$. We encode this into Julia by noticing that, in this setup, we match the
# general form when
#
# $$ f(u,p,t) = 0.98u $$
#
# with $ u_0 = 1.0 $. If we want to solve this model on a time span from `t=0.0` to `t=1.0`,
# then we define an `ODEProblem` by specifying this function `f`, this initial condition
# `u0`, and this time span as follows:
# ------------------------------------------------------------------------------------------

using DifferentialEquations
f(u,p,t) = 0.98u
u0 = 1.0
tspan = (0.0,1.0)
prob = ODEProblem(f,u0,tspan)

# ------------------------------------------------------------------------------------------
# To solve our `ODEProblem` we use the command `solve`.
# ------------------------------------------------------------------------------------------

sol = solve(prob)

# ------------------------------------------------------------------------------------------
# and that's it: we have succesfully solved our first ODE!
#
# #### Analyzing the Solution
#
# Of course, the solution type is not interesting in and of itself. We want to understand
# the solution! The documentation page which explains in detail the functions for analyzing
# the solution is the [Solution
# Handling](http://docs.juliadiffeq.org/latest/basics/solution.html) page. Here we will
# describe some of the basics. You can plot the solution using the plot recipe provided by
# [Plots.jl](http://docs.juliaplots.org/latest/):
# ------------------------------------------------------------------------------------------

using Plots; gr()
plot(sol)

# ------------------------------------------------------------------------------------------
# From the picture we see that the solution is an exponential curve, which matches our
# intuition. As a plot recipe, we can annotate the result using any of the [Plots.jl
# attributes](http://docs.juliaplots.org/latest/attributes/). For example:
# ------------------------------------------------------------------------------------------

plot(sol,linewidth=5,title="Solution to the linear ODE with a thick line",
     xaxis="Time (t)",yaxis="u(t) (in μm)",label="My Thick Line!") # legend=false

# ------------------------------------------------------------------------------------------
# Using the mutating `plot!` command we can add other pieces to our plot. For this ODE we
# know that the true solution is $u(t) = u_0 exp(at)$, so let's add some of the true
# solution to our plot:
# ------------------------------------------------------------------------------------------

plot!(sol.t, t->1.0*exp(0.98t),lw=3,ls=:dash,label="True Solution!")

# ------------------------------------------------------------------------------------------
# In the previous command I demonstrated `sol.t`, which grabs the array of time points that
# the solution was saved at:
# ------------------------------------------------------------------------------------------

sol.t

# ------------------------------------------------------------------------------------------
# We can get the array of solution values using `sol.u`:
# ------------------------------------------------------------------------------------------

sol.u

# ------------------------------------------------------------------------------------------
# `sol.u[i]` is the value of the solution at time `sol.t[i]`. We can compute arrays of
# functions of the solution values using standard comprehensions, like:
# ------------------------------------------------------------------------------------------

[t+u for (u,t) in tuples(sol)]

# ------------------------------------------------------------------------------------------
# However, one interesting feature is that, by default, the solution is a continuous
# function. If we check the print out again:
# ------------------------------------------------------------------------------------------

sol

# ------------------------------------------------------------------------------------------
# you see that it says that the solution has a 4th order interpolation, meaning that it is a
# continuous function of 4th order accuracy. We can call the solution as a function of time
# `sol(t)`. For example, to get the value at `t=0.45`, we can use the command:
# ------------------------------------------------------------------------------------------

sol(0.45)

# ------------------------------------------------------------------------------------------
# #### Controlling the Solver
#
# DifferentialEquations.jl has a common set of solver controls among its algorithms which
# can be found [at the Common Solver
# Options](http://docs.juliadiffeq.org/latest/basics/common_solver_opts.html) page. We will
# detail some of the most widely used options.
#
# The most useful options are the tolerances `abstol` and `reltol`. These tell the internal
# adaptive time stepping engine how precise of a solution you want. Generally, `reltol` is
# the relative accuracy while `abstol` is the accuracy when `u` is near zero. These
# tolerances are local tolerances and thus are not global guarantees. However, a good rule
# of thumb is that the total solution accuracy is 1-2 digits less than the relative
# tolerances. Thus for the defaults `abstol=1e-6` and `reltol=1e-3`, you can expect a global
# accuracy of about 1-2 digits. If we want to get around 6 digits of accuracy, we can use
# the commands:
# ------------------------------------------------------------------------------------------

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

# ------------------------------------------------------------------------------------------
# Now we can see no visible difference against the true solution:
# ------------------------------------------------------------------------------------------

plot(sol)
plot!(sol.t, t->1.0*exp(0.98t),lw=3,ls=:dash,label="True Solution!")

# ------------------------------------------------------------------------------------------
# Notice that by decreasing the tolerance, the number of steps the solver had to take was
# `9` instead of the previous `5`. There is a trade off between accuracy and speed, and it
# is up to you to determine what is the right balance for your problem.
#
# Another common option is to use `saveat` to make the solver save at specific time points.
# For example, if we want the solution at an even grid of `t=0.1k` for integers `k`, we
# would use the command:
# ------------------------------------------------------------------------------------------

sol = solve(prob,saveat=0.1)

# ------------------------------------------------------------------------------------------
# Notice that when `saveat` is used the continuous output variables are no longer saved and
# thus `sol(t)`, the interpolation, is only first order. We can save at an uneven grid of
# points by passing a collection of values to `saveat`. For example:
# ------------------------------------------------------------------------------------------

sol = solve(prob,saveat=[0.2,0.7,0.9])

# ------------------------------------------------------------------------------------------
# By default it always saves the first and last values, but we can turn this off as well:
# ------------------------------------------------------------------------------------------

sol = solve(prob,saveat=[0.2,0.7,0.9],save_start = false, save_end = false)

# ------------------------------------------------------------------------------------------
# If we need to reduce the amount of saving, we can also turn off the continuous output
# directly via `dense=false`:
# ------------------------------------------------------------------------------------------

sol = solve(prob,dense=false)

# ------------------------------------------------------------------------------------------
# and to turn off all intermediate saving we can use `save_everystep=false`:
# ------------------------------------------------------------------------------------------

sol = solve(prob,save_everystep=false)

# ------------------------------------------------------------------------------------------
# More advanced saving behaviors, such as saving functionals of the solution, are handled
# via the `SavingCallback` in the [Callback Library](http://docs.juliadiffeq.org/latest/feat
# ures/callback_library.html#SavingCallback-1) which will be addressed later in the
# tutorial.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# #### Choosing Solver Algorithms
#
# There is no best algorithm for numerically solving a differential equation. When you call
# `solve(prob)`, DifferentialEquations.jl makes a guess at a good algorithm for your
# problem, given the properties that you ask for (the tolerances, the saving information,
# etc.). However, in many cases you may want more direct control. A later notebook will help
# introduce the various *algorithms* in DifferentialEquations.jl, but for now let's
# introduce the *syntax*.
#
# The most crucial determining factor in choosing a numerical method is the stiffness of the
# model. Stiffness is roughly characterized by a Jacobian `f` with large eigenvalues. That's
# quite mathematical, and we can think of it more intuitively: if you have big numbers in
# `f` (like parameters of order `1e5`), then it's probably stiff. Or, as the creator of the
# MATLAB ODE Suite, Lawrence Shampine, likes to define it, if the standard algorithms are
# slow, then it's stiff. We will go into more depth about diagnosing stiffness in a later
# tutorial, but for now note that if you believe your model may be stiff, you can hint this
# to the algorithm chooser via `alg_hints = [:stiff]`.
# ------------------------------------------------------------------------------------------

sol = solve(prob,alg_hints=[:stiff])

# ------------------------------------------------------------------------------------------
# Stiff algorithms have to solve implicit equations and linear systems at each step so they
# should only be used when required.
#
# If we want to choose an algorithm directly, you can pass the algorithm type after the
# problem as `solve(prob,alg)`. For example, let's solve this problem using the `Tsit5()`
# algorithm, and just for show let's change the relative tolerance to `1e-6` at the same
# time:
# ------------------------------------------------------------------------------------------

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

# ------------------------------------------------------------------------------------------
# ### Systems of ODEs: The Lorenz Equation
#
# Now let's move to a system of ODEs. The [Lorenz
# equation](https://en.wikipedia.org/wiki/Lorenz_system) is the famous "butterfly attractor"
# that spawned chaos theory. It is defined by the system of ODEs:
#
# $$ \frac{dx}{dt} = \sigma (y - x) $$
# $$ \frac{dy}{dt} = x (\rho - z) -y $$
# $$ \frac{dz}{dt} = xy - \beta z $$
#
# To define a system of differential equations in DifferentialEquations.jl, we define our
# `f` as a vector function with a vector initial condition. Thus, for the vector `u =
# [x,y,z]'`, we have the derivative function:
# ------------------------------------------------------------------------------------------

function lorenz!(du,u,p,t)
    σ,ρ,β = p
    du[1] = σ*(u[2]-u[1])
    du[2] = u[1]*(ρ-u[3]) - u[2]
    du[3] = u[1]*u[2] - β*u[3]
end

# ------------------------------------------------------------------------------------------
# Notice here we used the in-place format which writes the output to the preallocated vector
# `du`. For systems of equations the in-place format is faster. We use the initial condition
# $u_0 = [1.0,0.0,0.0]$ as follows:
# ------------------------------------------------------------------------------------------

u0 = [1.0,0.0,0.0]

# ------------------------------------------------------------------------------------------
# Lastly, for this model we made use of the parameters `p`. We need to set this value in the
# `ODEProblem` as well. For our model we want to solve using the parameters $\sigma = 10$,
# $\rho = 28$, and $\beta = 8/3$, and thus we build the parameter collection:
# ------------------------------------------------------------------------------------------

p = (10,28,8/3) # we could also make this an array, or any other type!

# ------------------------------------------------------------------------------------------
# Now we generate the `ODEProblem` type. In this case, since we have parameters, we add the
# parameter values to the end of the constructor call. Let's solve this on a time span of
# `t=0` to `t=100`:
# ------------------------------------------------------------------------------------------

tspan = (0.0,100.0)
prob = ODEProblem(lorenz!,u0,tspan,p)

# ------------------------------------------------------------------------------------------
# Now, just as before, we solve the problem:
# ------------------------------------------------------------------------------------------

sol = solve(prob)

# ------------------------------------------------------------------------------------------
# The same solution handling features apply to this case. Thus `sol.t` stores the time
# points and `sol.u` is an array storing the solution at the corresponding time points.
#
# However, there are a few extra features which are good to know when dealing with systems
# of equations. First of all, `sol` also acts like an array. `sol[i]` returns the solution
# at the `i`th time point.
# ------------------------------------------------------------------------------------------

sol.t[10],sol[10]

# ------------------------------------------------------------------------------------------
# Additionally, the solution acts like a matrix where `sol[j,i]` is the value of the `j`th
# variable at time `i`:
# ------------------------------------------------------------------------------------------

sol[2,10]

# ------------------------------------------------------------------------------------------
# We can get a real matrix by performing a conversion:
# ------------------------------------------------------------------------------------------

A = convert(Array,sol)

# ------------------------------------------------------------------------------------------
# This is the same as sol, i.e. `sol[i,j] = A[i,j]`, but now it's a true matrix. Plotting
# will by default show the time series for each variable:
# ------------------------------------------------------------------------------------------

plot(sol)

# ------------------------------------------------------------------------------------------
# If we instead want to plot values against each other, we can use the `vars` command. Let's
# plot variable `1` against variable `2` against variable `3`:
# ------------------------------------------------------------------------------------------

plot(sol,vars=(1,2,3))

# ------------------------------------------------------------------------------------------
# This is the classic Lorenz attractor plot, where the `x` axis is `u[1]`, the `y` axis is
# `u[2]`, and the `z` axis is `u[3]`. Note that the plot recipe by default uses the
# interpolation, but we can turn this off:
# ------------------------------------------------------------------------------------------

plot(sol,vars=(1,2,3),denseplot=false)

# ------------------------------------------------------------------------------------------
# Yikes! This shows how calculating the continuous solution has saved a lot of computational
# effort by computing only a sparse solution and filling in the values! Note that in vars,
# `0=time`, and thus we can plot the time series of a single component like:
# ------------------------------------------------------------------------------------------

plot(sol,vars=(0,2))

# ------------------------------------------------------------------------------------------
# ### A DSL for Parameterized Functions
#
# In many cases you may be defining a lot of functions with parameters. There exists the
# domain-specific language (DSL) defined by the `@ode_def` macro for helping with this
# common problem. For example, we can define the Lotka-Volterra equation:
#
# $$ \frac{dx}{dt} = ax - bxy $$
# $$ \frac{dy}{dt} = -cy + dxy $$
#
# as follows:
# ------------------------------------------------------------------------------------------

function lotka_volterra!(du,u,p,t)
  du[1] = p[1]*u[1] - p[2]*u[1]*u[2]
  du[2] = -p[3]*u[2] + p[4]*u[1]*u[2]
end

# ------------------------------------------------------------------------------------------
# However, that can be hard to follow since there's a lot of "programming" getting in the
# way. Instead, you can use the `@ode_def` macro:
# ------------------------------------------------------------------------------------------

lv! = @ode_def LotkaVolterra begin
  dx = a*x - b*x*y
  dy = -c*y + d*x*y
end a b c d

# ------------------------------------------------------------------------------------------
# We can then use the result just like an ODE function from before:
# ------------------------------------------------------------------------------------------

u0 = [1.0,1.0]
p = (1.5,1.0,3.0,1.0)
tspan = (0.0,10.0)
prob = ODEProblem(lv!,u0,tspan,p)
sol = solve(prob)
plot(sol)

# ------------------------------------------------------------------------------------------
# Not only is the DSL convenient syntax, but it does some magic behind the scenes. For
# example, further parts of the tutorial will describe how solvers for stiff differential
# equations have to make use of the Jacobian in calculations. Here, the DSL uses symbolic
# differentiation to automatically derive that function:
# ------------------------------------------------------------------------------------------

lv!.Jex

# ------------------------------------------------------------------------------------------
# The DSL can derive many other functions; this ability is used to speed up the solvers. An
# extension to DifferentialEquations.jl, [Latexify.jl](https://korsbo.github.io/Latexify.jl/
# latest/tutorials/parameterizedfunctions.html), allows you to extract these pieces as LaTeX
# expressions.
# ------------------------------------------------------------------------------------------

# ------------------------------------------------------------------------------------------
# ## Internal Types
#
# The last basic user-interface feature to explore is the choice of types.
# DifferentialEquations.jl respects your input types to determine the internal types that
# are used. Thus since in the previous cases, when we used `Float64` values for the initial
# condition, this meant that the internal values would be solved using `Float64`. We made
# sure that time was specified via `Float64` values, meaning that time steps would utilize
# 64-bit floats as well. But, by simply changing these types we can change what is used
# internally.
#
# As a quick example, let's say we want to solve an ODE defined by a matrix. To do this, we
# can simply use a matrix as input.
# ------------------------------------------------------------------------------------------

A  = [1. 0  0 -5
      4 -2  4 -3
     -4  0  0  1
      5 -2  2  3]
u0 = rand(4,2)
tspan = (0.0,1.0)
f(u,p,t) = A*u
prob = ODEProblem(f,u0,tspan)
sol = solve(prob)

# ------------------------------------------------------------------------------------------
# There is no real difference from what we did before, but now in this case `u0` is a `4x2`
# matrix. Because of that, the solution at each time point is matrix:
# ------------------------------------------------------------------------------------------

sol[3]

# ------------------------------------------------------------------------------------------
# In DifferentialEquations.jl, you can use any type that defines `+`, `-`, `*`, `/`, and has
# an appropriate `norm`. For example, if we want arbitrary precision floating point numbers,
# we can change the input to be a matrix of `BigFloat`:
# ------------------------------------------------------------------------------------------

big_u0 = big.(u0)

# ------------------------------------------------------------------------------------------
# and we can solve the `ODEProblem` with arbitrary precision numbers by using that initial
# condition:
# ------------------------------------------------------------------------------------------

prob = ODEProblem(f,big_u0,tspan)
sol = solve(prob)

sol[1,3]

# ------------------------------------------------------------------------------------------
# To really make use of this, we would want to change `abstol` and `reltol` to be small!
# Notice that the type for "time" is different than the type for the dependent variables,
# and this can be used to optimize the algorithm via keeping multiple precisions. We can
# convert time to be arbitrary precision as well by defining our time span with `BigFloat`
# variables:
#
# ------------------------------------------------------------------------------------------

prob = ODEProblem(f,big_u0,big.(tspan))
sol = solve(prob)

# ------------------------------------------------------------------------------------------
# Let's end by showing a more complicated use of types. For small arrays, it's usually
# faster to do operations on static arrays via the package
# [StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl). The syntax is similar
# to that of normal arrays, but for these special arrays we utilize the `@SMatrix` macro to
# indicate we want to create a static array.
# ------------------------------------------------------------------------------------------

using StaticArrays
A  = @SMatrix [ 1.0  0.0 0.0 -5.0
                4.0 -2.0 4.0 -3.0
               -4.0  0.0 0.0  1.0
                5.0 -2.0 2.0  3.0]
u0 = @SMatrix rand(4,2)
tspan = (0.0,1.0)
f(u,p,t) = A*u
prob = ODEProblem(f,u0,tspan)
sol = solve(prob)

sol[3]

# ------------------------------------------------------------------------------------------
# ## Conclusion
#
# These are the basic controls in DifferentialEquations.jl. All equations are defined via a
# problem type, and the `solve` command is used with an algorithm choice (or the default) to
# get a solution. Every solution acts the same, like an array `sol[i]` with `sol.t[i]`, and
# also like a continuous function `sol(t)` with a nice plot command `plot(sol)`. The Common
# Solver Options can be used to control the solver for any equation type. Lastly, the types
# used in the numerical solving are determined by the input types, and this can be used to
# solve with arbitrary precision and add additional optimizations (this can be used to solve
# via GPUs for example!). While this was shown on ODEs, these techniques generalize to other
# types of equations as well.
# ------------------------------------------------------------------------------------------


import Pkg; Pkg.add(Pkg.PackageSpec(url="https://github.com/JuliaComputing/JuliaAcademyData.jl"))
using JuliaAcademyData; activate("Foundations of machine learning")

# # Minimizing functions - how a computer learns

#-

# In the previous notebooks, we've seen that by changing **parameters** in a function, we could find a "best fit" of that function to some data. We use a **loss function** to quantify the "goodness" of a set of parameters and we look for those parameters that **optimize**, in fact **minimize**, the loss function. If we teach a machine how to minimize the loss function on its own, we say that machine is able to **learn** how to model our data.
#
# In the last notebook, we threw around terms like **derivative**, **gradient**, and **gradient descent** to give you a rough sense of how we minimize a function on a computer. In this notebook, we will step through these concepts more carefully, with the aim of being able to implement them using Julia.

#-

# ## Minimizing a 1D function using calculus

#-

# Let's draw the 1D loss function $L_1$ from a previous notebook again:

σ(x) = 1 / (1 + exp(-x))
f(x, w) = σ(w * x)

x1 = 2
y1 = 0.8

L1(w) = (y1 - f(x1, w))^2

#-

using Plots; gr()

#-

plot(L1, -2, 1.5, xlabel="w", ylabel="L1(w)", leg=false)

# By eye, we can see that the minimum is around $w=0.6$. But how can we get the computer to work this out on its own?

#-

# In the previous notebook, we thought of this function plot as a hill, viewed from the side. We could find the minimum by making the hill sticky, and letting a ball roll down it. The ball will find and settle in the minimum of the function.  Now let's see how to teach a computer to do this.
#
# We need to find the downhill direction along the hill, which is related to its *slope* (how steep it is). Calculus provides us with tools to calculate that slope!

#-

# Namely, the slope of a curve $L_1(w)$ at $w$ is given by its **derivative** $L_1'(w)$; geometrically, this is the slope of the **tangent line** to the curve at that point, i.e. the straight line which touches the curve at that point.

#-

# Calculus provides us with some rules to calculate an analytical formula for the derivative, and we will see later how to apply these rules, albeit indirectly, for machine learning.
# To gain understanding, however, we will see here how to get the computer to help us out by calculating the derivatives numerically instead!

#-

# ## Approximating derivatives

#-

# Let's recall that the derivative $L_1'(w)$ of a function is defined as
#
# $$L_1'(w) \simeq \frac{L_1(w+h) - L_1(w)}{h},$$
#
# for a small step size $h$. (Strictly speaking, we must take the limit when $h$ tends to $0$ to obtain the exact value of the derivative.)

#-

# #### Exercise 1
#
# Write a function to calculate the derivative of a function at a given point. Note that in Julia, we can easily pass functions as arguments to other functions!
# The function should take the function, the point $w$ at which to calculate the derivative, and the value of $h$, which should have a default value of 0.001.
#
# *Note*: A function input argument can have a *default value* if we set the input argument equal to that default value when we define the function. For example,
#
# ```julia
# f(x, a = 3) = a * x^2
# ```
#
# The function `f` will square the input we give it and multiply by `a`. However,  if we choose to call `f(x)` *without* passing it an `a` input, it will assume `a` is `3` and return `3*x^2`.

#-

# #### Solution
#
# We could define the `derivative` function as

function derivative(f, w, h=0.001)
    return ( f(w + h) - f(w) ) / h
end
## Or the short form with:
## derivative(f, w, h=0.001) = ( f(w + h) - f(w) ) / h

# **Tests**

@assert derivative(x->x^2, 0.2) ≈ 0.401    # check default h works
@assert derivative(x->x^2, 0.2, 0.5) ≈ .9  # check we can specify our own h
@assert derivative(x->x^3, 0.2) ≈ 0.120601
@assert derivative(x->x^3, -0.2) ≈ 0.119401

# #### Exercise 2
#
# Write an interactive visualization of the tangent line to the graph of $L_1$, so that we can visualize the tangent at any point on $L_1$. Include the current value of the derivative in the title.
#
# *Hint*: Recall that the straight line through the point $(x_0, y_0)$ with slope $m$ is given by
#
# $$ \frac{y - y_0}{x - x_0} = m,$$
#
# so
#
# $$y = y_0 + m*(x - x_0).$$

#-

# #### Solution

wmin = -2
wmax = 1

w = 0.5 # Try manipulating w between wmin and wmax (-2 to 1)
plot(L1, wmin, wmax, label="L_1")

x0 = w
y0 = L1(w)

scatter!([x0], [y0], label="")

m = derivative(L1, w)

plot!(x->y0 + m*(x - x0), ls=:dash, label="tangent line")

ylims!(-0.1, 0.6)
xlims!(wmin, wmax)

title!("Derivative = $m")
xlabel!("w")
ylabel!("L1(w)")

# #### Exercise 3
#
# What is the value of the derivative (slope of the tangent line) at a minimum? Can this happen anywhere else?

#-

# #### Solution

#-

# The derivative is `0` at a minimum. This happens at any minimum, maximum or "point of inflection" of a smooth function.

#-

# ## Minimization by gradient descent

#-

# The tangent line at a point is a good approximation of the function near that point. Therefore the derivative tells us in which direction, and how fast, a function grows or shrinks when we move a small distance away from that point.
#
# As we saw in the previous notebook, we can think of the function as being a hill, and having a ball that we want to move down the hill. Gravity will automatically pull the ball in the direction *down* the hill; we want to emulate this in the computer!

#-

# #### Exercise 4
#
# When the derivative $L_1'(w)$ is positive, that means that $L_1$ increases from left to right at point $w$.
#
# If the derivative $L_1'(w)$ is positive (> 0), in which direction should we move $w$ to *decrease* $L_1$?

#-

# #### Solution
#
# We should move $w$ left, i.e. decrease it.

#-

# #### Exercise 5
#
# If the derivative $L_1'(w)$ is negative (< 0), in which direction should we move $w$ to *decrease* $L_1$?

#-

# We can use this information to tell the computer which way to take a step. This constitutes the numerical algorithm called **gradient descent**; it is called this since we are descending (moving downwards) along the graph of the function by using information about its gradient (slope).

#-

# #### Exercise 6
#
# Implement gradient descent by following this prescription for an **iterative (repetitive) algorithm**:
#
# 1. Start at an initial guess $w_0$ for $w$.
#
# 2. At each step, calculate the derivative, $L_1'(w_n)$ at the current value of $w_n$ using the function that you created above.
#
# 3. Modify the value of $w$ by a small multiple (for example, $\eta=0.01$) of the value of the derivative you just created, via $w_{n+1} = w_n - \eta L_1'(w_n)$.
#
# For this problem, start with $w_0 = -2.0$. Repeat steps 2 and 3 a total of `2000` times.
#
# Package this code into a function called `gradient_descent` that takes two inputs, a function and a range for values of $w$, and returns the final value of $w$ and $L_1(w)$.
#
# Using `L1` and `-2:0.01:1.5` as your inputs to `gradient_descent`, for what value of $w$ is $L_1$ at a minimum?

#-

# #### Solution

function gradient_descent(f, wrange)
    derivative(f, w, h=0.001) = ( f(w + h) - f(w) ) / h
    w = -2.0
    for i in 1:2000
        fderiv = derivative(f, w)
        step = .01 * fderiv
        w = w - step
    end
    return w, f(w)
end

gradient_descent(L1, -2:0.01:1.5)

# **Tests**

@assert 0.6383 < gradient_descent(L1, -2:0.01:1.5)[1] < 0.6385
@assert 0.00032 < gradient_descent(L1, -2:0.01:1.5)[2] < 0.00033

# #### Exercise 7
#
# Modify your code for gradient descent to return the result once it has found an answer within some *tolerance*, rather than taking a set number of steps. The new prescription for this algorithm is:
#
# 1. Start at an initial guess $w_0$ for $w$.
#
# 2. At each step, calculate the derivative, $L_1'(w_n)$ at the current value of $w_n$, using the function that you created above.
#
# 3. Modify the value of $w$ by a small multiple (for example, $\eta=0.01$) of the value of the derivative you just created, via $w_{n+1} = w_n - \eta L_1'(w_n)$.
#
# 4. Check how different $w_{n+1}$ is from $w_n$. If you're satisfied that $L_1(w_{n+1})$ is minimized, return $w_{n+1}$ and $L_1(w_{n+1})$. Otherwise, go to step (2) and continue.
#
# Edit `gradient_descent` so that it takes three inputs: a function, a range for values of $w$, and a tolerance that tells you how close $w_{n+1}$ must be to $w_n$ before you can stop iterating.
#
# Using `L1`, `-2:0.01:1.5`, and `.000001` as your inputs to `gradient_descent`, for what value of $w$ is $L_1$ at a minimum?

#-

# #### Solution
#
#
# Your code should look something like this
#
# $L_1$ is at a minimum when `gradient_descent(L1, -2:0.01:1.5, .000001)[1]`, or $w$, is around `0.69216`.

function gradient_descent(f, wrange, tol)
    w = -2.0
    while true
        fderiv = derivative(f, w)
        step = .01 * fderiv
        w = w - step
        if abs(step) < tol
           return w, f(w)
        end
    end
end

gradient_descent(L1, -2:0.01:1.5, .000001)

# **Tests**

@assert 0.69215 < gradient_descent(L1, -2:0.01:1.5, .000001)[1] < 0.69217

# After you've completed this step, you will have found a good approximation to the value $w^*$ where $L_1$ reaches its minimum, and the minimum value $L_1(w^*)$, which in this case is (almost) $0$.

#-

# #### Exercise 8
#
# Alter the function `gradient_descent` so that it stores the results `(w, L1(w))` at each step of the algorithm as an array and returns this array. How many steps does the algorithm take for input parameters `(L1, -2:0.01:1.5, .000001)` before terminating? You should count your starting $w_0$ as your first step.

#-

# #### Solution

#-

# Your altered code should look like this, getting the number of steps with `length`:

function gradient_descent(f, wrange, tol)
    w = -2.0
    store_w = [(w, f(w))]
    while true
        fderiv = derivative(f, w)
        step = .01 * fderiv
        w = w - step
        push!(store_w, (w, f(w)))
        if abs(step) < tol
           return store_w
        end
    end
end

length(gradient_descent(L1, -2:0.01:1.5, .000001))

# **Tests**

@assert length(gradient_descent(L1, -2:0.01:1.5, .000001)) == 4255

# #### Exercise 9
#
# Overlay the steps taken in the last exercise with a plot of $L_1(w)$ vs. $w$.
#
# Where does our algorithm take the largest steps? (Where does the ball move fastest down the hill?)
#
# A) Between w = -2:-1

# B) Between w = -1:0

# C) Between w = 0:.6
#
# *Hint*: It may be easier to see what's going on if you only plot, for example, every 15th step.

#-

# #### Solution
#
# B) The following code generates the desired figure:

wsteps = gradient_descent(L1, -2:0.01:1.5, .000001)
fig1 = plot(L1, -2, 1.5, xlabel="w", ylabel="C1(w)", leg=false)
for step in wsteps[1:15:end]
   scatter!([step[1]], [step[2]])
end
fig1

# ## Functions of 2 variables and their derivatives

#-

# So far, we have been looking at the minimizing a loss function $L_1(w)$ that depends on a single parameter, $w$. Now let's turn to the cost function $L_2(w, b)$ from a previous notebook, that is a function of *two* parameters, $w$ and $b$, and try to minimize it.
#
# As we've seen, we get a **surface**, instead of a curve, when we graph $L_2$ as a function of both of its parameters.

#-

# **Exercise 10**
#
# Draw a surface plot of $L_2$, given by
#
# $$L_{2}(w, b) = \sum_i(y_i - g(x_i, w, b))^2$$
#
# using the `surface` function from `Plots.jl`. For this plot, use the values of `xs` and `ys` from notebook 5:
#
# We can get a nice interactive 3D plot by using the Plotly backend of `Plots.jl` by executing the command
#
#     plotly()

xs = [2, -3, -1, 1]
ys = [0.8, 0.3, 0.4, 0.4]

# #### Solution

#-

# The following code will generate the desired graphic:

plotly()
σ(x) = 1 / (1 + exp(-x))
g(x, w, b) = σ(w*x) + b

xs = [2, -3, -1, 1]
ys = [0.8, 0.3, 0.4, 0.4]

L2(w, b) = sum(abs2, ys .- g.(xs, w, b))

surface(-2:0.02:2, -2:0.02:2, L2, alpha=0.6, xlabel = "w", ylabel = "b", zlabel = "L2(w, b)")

# ### Finding the minimum

#-

# We can just about see, by rotating the graph, that $L_2$ has a single minimum. We want to find the values of $w$ and $b$ where this minimum value is reached.

#-

# Following what we did for the function $L_1$ above, we expect that we will need to calculate derivatives of $L_2$. Since the function is more complicated, though, the derivatives are too!
#
# It turns out that the right concept is that of the [**gradient**](https://en.wikipedia.org/wiki/Gradient) of $L_2$, denoted $\nabla L_2(w, b)$. This is a **vector** consisting of $2$ numbers if there are $2$ parameters [or, in general, $n$ numbers if there are $n$ parameters].
#
# The numbers that form the gradient $\nabla L_2(w, b)$ are called the **partial derivatives** of $L_2$ with respect to $w$ and $b$,  written as
#
# $$\frac{\partial L_2}{\partial w} \quad \text{and} \quad \frac{\partial L_2}{\partial b}.$$
#
# Although this notation might look complicated, all it means is that we calculate derivatives just like we did before, except that we fix the value of the other variable.
# For example, to calculate $\frac{\partial L_2}{\partial w}$, the partial derivative of $L_2$ with respect to $w$, we fix the value of $b$ and think of the resulting function as a function of the single variable $w$; then we use the formula for derivatives of functions of a single variable.
#
# [Note that $\frac{\partial L_2}{\partial w}$ is itself a function of $w$ and $b$; we could write $\frac{\partial L_2}{\partial w}(w, b)$.]

#-

# #### Exercise 11
#
# Write functions that will allow you to calculate the partial derivatives of a function $f$ of two variables.
#
# In particular, declare functions called `partial_w` and `partial_b`. Each should take four inputs - a function $f$ of two variables, the first input argument to $f$, the second input argument to $f$, and a step size `h` with default value `0.001`. `partial_w` should return the partial derivative of $f$ with respect to its first input argument and `partial_b` should return the partial derivative of $f$ with respect to its second input argument.

#-

# #### Solution
#
# The partial derivatives are

partial_w(f, w, b, h = 0.001) = ( f(w + h, b) - f(w, b) ) / h
partial_b(f, w, b, h = 0.001) = ( f(w, b + h) - f(w, b) ) / h

# **Tests**

## Test that default parameters work:
@assert partial_w((x, y) -> 3 * x * y, 0.3, 0.3) ≈ 0.9
@assert partial_b((x, y) -> 3 * x * y, 0.3, 0.3) ≈ 0.9
## Test that you can specify your own step sizes:
@assert partial_w((x, y) -> 3 * x^2 * y^3, 2.3, 10.2, 1) ≈ 17828.2944
@assert partial_b((x, y) -> 3 * x^2 * y^3, 2.3, 10.2, 1) ≈ 5454.8364
## Test that x and y are treated separately:
@assert partial_w((x, y) -> 3 * x * y^2, 0.3, 0.2) ≈ 0.12
@assert partial_b((x, y) -> 3 * x * y^2, 0.3, 0.2) ≈ 0.3609
@assert partial_w((x, y) -> 3 * x^2 * y, 0.3, 0.2, 0.1) ≈ 0.42
@assert partial_b((x, y) -> 3 * x^2 * y, 0.3, 0.2, 0.1) ≈ 0.27

# #### Exercise 12
#
# Use `partial_b` from the last exercise to find the partial derivative of $L_2$ with respect to $w$ at b = 0.3, $\frac{\partial L_2}{\partial w}|_{b = 0.3}$ for `w = -2:0.01:1`

#-

# #### Solution

wrange = -2:0.01:1
partialw = partial_w.(L2, wrange, 0.3 )

# **Tests**

@assert isapprox(partialw[3], -0.212573, rtol=2e-6)

# #### Exercise 13
#
# Plot the cross section of the surface of $L_2(w, b)$ at $b = 0.3$. Make this plot interactive to show that the function `partial_w` gives the slope of the tangent to this cross section for any point `w` in the range `-2:0.01:1`.
#
# For what value of $w$ in this range is the slope of the cross section closest to -1?

#-

# #### Solution

#-

# The slope is closest to -1.0 at w = -0.55. You can see this by running the following code:

L2_fixedb = L2.(wrange, 0.3)

w = 0.0 # Try manipulating w to pick a value within `wrange`
plot(wrange, L2_fixedb)
scatter!([w], [L2(w, 0.3)])
slope = partial_w(L2, w, 0.3)
tangent_points = slope .* wrange .- slope * w .+ L2.(w, 0.3)
plot!(wrange, tangent_points)
ylims!(0.3, 2.0)
title!("Slope = $slope @ $(w)")

# ## ***Optional**: Functions with $n$ inputs

#-

# If a function $f$ takes $n$ input arguments, we can write them as $p_1, \ldots, p_n$, where $p_i$ means the "$i$th parameter". In Julia, we can wrap them up into a single **vector**. Now we can calculate the partial derivative $\frac{\partial L_2}{\partial p_i}$ with respect to the $i$th variable.

#-

# #### Exercise 14
#
# For the next exercise, you will need to use the splat command, `...`. You can use this command to "open up" a collection and pass all the elements of that collection as inputs to a function.
#
# For example, say you have an array, `numbers`,
#
# ```julia
# numbers = [4, 3, 2]
# ```
#
# and you want to use `numbers` to create a $4\times3\times3$ randomly populated array via `rand`. `rand(numbers)` will not do what you want. You could index into `numbers` to grab the values you want and pass them to `rand`, as in
#
# ```julia
# rand(numbers[1], numbers[2], numbers[3])
# ```
#
# or you could use a splat:
#
# ```julia
# rand(numbers...)
# ```
#
# Use `...` to pass the contents of `inputs`
#
# ```julia
# inputs = [30, 12, "cats"]
# ```
#
# to the function `dreams`
#
# ```julia
# dreams(i, j, perfect_mammal) = "I wish I had $(i + j) $perfect_mammal."
# ```

#-

# #### Solution

inputs = [30, 12, "cats"]
dreams(i, j, perfect_mammal) = "I wish I had $(i + j) $perfect_mammal."
dreams(inputs...)

# **Tests**

@assert dreams(inputs...) == "I wish I had 42 cats."

# #### Exercise 15:
#
# Write a function, `partial`, to calculate the $i$th partial derivative of a function. This function should have four inputs
#
# * a function, $f$, for which you want to compute the partial derivative
# * an array, *p*, specifying the values of all input arguments to $f$ at the point where you want $\frac{\partial f}{\partial p_i}$ computed
# * the index, $i$, of the variable with respect to which you want to calculate the partial derivative of $f$
# * a step size with default value 0.001
#
# Hint: you will need to `copy` and modify `p` within `partial`.

#-

# #### Solution

function partial(f, p, i, h = 0.001)
    p_next = copy(p)
    p_next[i] += h
    return ( f(p_next...) - f(p...) ) / h
end

# **Tests**

@assert partial((x, y, z) -> x^2 * y * z, [1.0, 2.0, 1.0], 1, 0.0001) ≈ 4.0002
@assert partial((x, y, z) -> x^2 * y * z, [1.0, 2.0, 1.0], 1, 1) ≈ 6

@assert partial((w, x, y, z) -> w * x^2 * y^3 * z^4, [1.1, 2.2, 3.3, 4.4], 1) ≈ 65192.53776074038
@assert partial((w, x, y, z) -> w * x^2 * y^3 * z^4, [1.1, 2.2, 3.3, 4.4], 2) ≈ 65207.35424659506
@assert partial((w, x, y, z) -> w * x^2 * y^3 * z^4, [1.1, 2.2, 3.3, 4.4], 3) ≈ 65212.29507069802
@assert partial((w, x, y, z) -> w * x^2 * y^3 * z^4, [1.1, 2.2, 3.3, 4.4], 4) ≈ 65214.7658571339

# ## Gradient descent in 2 dimensions

#-

# It turns out that the gradient vector of the function $L_2(w, b)$ gives the direction in the plane $(w, b)$ in which the function $L_2$ **increases fastest**.
#
# In other words, if you start at the position $(w_0, b_0)$ and take a small step of length $\eta$ to new values $(w', b')$, the value of the function will change to a new value $L_2(w', b')$. How do we get to the minimum of $L_2(w, b)$ as fast as possible? We want to step in the direction where $L_2$ decreases fastest!
#
# In multivariable calculus courses, it is shown that $L_2$ will *increase* fastest if you take a step **in the direction of the gradient $\nabla L_2(w, b)$**! To decrease $L_2$ the fastest, we should take a step in the *opposite* direction, $-\nabla L_2(w, b)$.
#
# Let's now generalize the gradient descent algorithm that we wrote previously to work our 2-dimensional function.

#-

# #### Exercise 16
#
# Extend your 1D implementation of `gradient_descent` so that it will minimize the function $L_2$.
#
# Requirements:
#
# * Your new method for `gradient_descent` will take four input arguments: the function $f$ for which you seek the minimum, the range of values for $f$'s first input argument that you will consider, the range of values for $f$'s second input argument that you will consider, and a tolerance that will specify the maximum allowable step size, $\sum_i \eta \frac{\partial f}{\partial p_i}$
#
# * Use $\eta = .01$. For example, for a function $f(w, b)$, update $w$ such that $w_{n+1} = w_n - 0.01 * \frac{\partial f}{\partial w_n}$
#
# * Seed `gradient_descent` with the starting coordinates [-2.0, -2.0], i.e. $w_0 = -2.0$ and $b_0 = -2.0$.
#
# * Return all steps (their coordinates) taken during gradient descent and the values of the loss function at these coordinates.
#
# Once you have done this, execute
#
# ```julia
# gradient_descent(L2, -2:0.02:2, -2:0.02:2, .001)
# ```
#
# How many steps were taken by gradient descent?
#
# Hint: Do not count your starting coordinates `[-2.0, -2.0]` as a step.

#-

# #### Solution
#
#
# The new code for `gradient_descent` should look something like

function gradient_descent(f, wrange, brange, tol)
    w, b = -2.0, -2
    store_steps = [[w, b]]
    store_C2_vals = [f(w, b)]
    while true
        # Calculate partial derivatives
        fpartial1 = partial(f, [w, b], 1)
        fpartial2 = partial(f, [w, b], 2)
        # Update parameters
        w = w - .01 * fpartial1
        b = b - .01 * fpartial2
        # Store parameters and resulting function
        push!(store_steps, [w, b])
        push!(store_C2_vals, f(w, b))

        if abs(w + b) < tol
           return store_steps, store_C2_vals
        end
    end
end

# **Tests**

steps, values = gradient_descent(L2, -2:0.02:2, -2:0.02:2, .001 )
@assert length(values) == 396

# #### Exercise 17
#
# Use the `surface` and `scatter!` commands to illustrate the path taken by `gradient_descent` from [-2.0, -2.0] to wherever the algorithm terminates.
#
# Where do the scattered points representing the steps of gradient descent appear the most dense?
#
# A) Near the starting point at [-2.0, -2.0]

# B) Near the point [-1.8, 0] where $C_2$ appears nearly flat

# C) Near the minimum of $C_2$

#-

# #### Solution
#
# B) The following code will plot the function `C2` with an overlay of the steps taken in gradient descent
#
# You can change the frequency with which steps are plotted by changing the iterator in the `for` loop from `1:length(values)` to `1:n:length(values)` for varying values of `n`. This will allow you to see more clearly where steps are densely or sparsely packed.

fig = surface(-2:0.02:2, -2:0.02:2, L2, alpha=0.6, xlabel = "w", ylabel = "b", zlabel = "C2(w, b)")
for i in 1:length(values)
   scatter!([steps[i][1]], [steps[i][2]], [values[i]])
end
fig

# ## What we have learnt

#-

# To recap, in this notebook, we have seen that the computer can calculate approximate derivatives, and that we can use those inside a relatively simple algorithm, gradient descent, to minimize functions of 1 and 2 variables!
#
# Now that we've used gradient descent to **learn** the best values of `b` and `w` for a cost function, we can imagine how we could do our classification of apples vs. bananas by allowing the machine to choose our parameters! To do this, we need to define a loss function for apple vs. banana classification, using images that have already been categorised as apples or bananas.

Courses/Foundations of machine learning/0900.What-is-learning-quiz.jl


# # Investment Science - Module 1
#
# \tableofcontents
#
# ## Problem 4
# The IRR is generally calculated using an iterative procedure.
# Suppose that we define $f(\lambda) = -a_{1} + a_{2}\lambda^{2} + \ldots + a_{n}\lambda^{n}$,
# where all $a_{i}$'s are positive and $n > 1$.
# Here is an iterative technique that generates a sequence
# $\lambda_{0}, \lambda_{1}, \lambda_{2}, \ldots \lambda_{k}, \ldots$
# of estimates that converges to the root $\overline{\lambda} > 0$, solving $f(\overline{\lambda}) = 0$.
# Start with any $\lambda_{0} > 0$ close to the solution.
# Assuming $\lambda_{k}$ has been calculated, evaluate
#
# $$
# f'(\lambda_{k}) = a_{1} + 2a_{2}\lambda_{k} + 3a_{3}\lambda_{k}^{2} + \ldots + na_{n}\lambda_{k}^{n-1}
# $$
#
# and define
#
# $$
# \lambda_{k+1} = \lambda_{k} - \frac{f(\lambda_{k})}{f'(\lambda_{k})}
# $$
#
# This is Newton's method.
# It is based on approximating the function $f$ by a line tangent to its graph at $\lambda_{k}$.
# Try the procedure on $f(\lambda) = -1 + \lambda + \lambda^{2}$.
# Start with $\lambda_{0} = 1$ and compute four additional estimates.
#
# ### Solution
#
# We know
#
# $$
# \begin{aligned}
# \lambda_{0} &= 1 \\
# \lambda_{k+1} &= \lambda_{k} - \frac{f(\lambda_{k})}{f'(\lambda_{k})} \\
# f(\lambda) &= -1 + \lambda + \lambda^{2} \\
# f'(\lambda) &= 1 + 2\lambda \\
# \end{aligned}
# $$
#
# Using all these equations we can compute four additional estimates.

## define function f
f(λ) = -1 + λ + λ^2
## define derivative of f
ḟ(λ) = 1 + 2*λ
## define function to calculate the next λ
next(λ) = λ - f(λ)/ḟ(λ)
## create a list of λs starting with 1
λs = [1.0]
## calculate next four
for n ∈ 1:4
    push!(λs, next(λs[n]))
end
## display results
@show λs

# ## Problem 5
#
# Suppose that you have the opportunity to plant trees that alter can be sold for lumber. This project requires an initial outlay of money in order to purchase and plant the seedlings. No other cash flow occurs until trees are harvested. However, you have a choice as to when to harvest. If you harvest after 1 year, you get your return quickly; but if you wait, the trees will have additional growth and the revenue generated from the sale of the trees will be greater. Assume that the case flow streams associated with these alternatives are:
#
# - Wait $1$ year: $(-1, 2)$
# - Wait $2$ year: $(-1, ,0, 3)$
# - Wait $3$ year: $(-1, 0, 0, 4)$
# - $\vdots$
# - Wait $n$ year: $(-1, 0, 0, \ldots, n+1)$
#
# The prevailing interest rate is 10%. When is the best time to cut the trees?
#
# ### Solution
#
# Let’s first define a function for calculating the Net Present Value, say on a cash flow stream that looks like $(x_{0}, \ldots, x_{n})$ where $r$ is the prevailing interest rate and $k$ is the number of years waited to harvest:
#
# $$
# NPV = -x_{0} + \frac{x_{k}}{(1+r)^{k}}
# $$
#
# We can create a function to calculate the **_NPV_** given a _cash flow stream_ and _interest rate_. Let's create a table of the first several values to determine the optimal time to harvest.

## define function to calculate net present value
function net_present_value(investment, revenue, year, rate)
    -investment + revenue/(1+rate)^year
end
## define empty array for values
values = []
## populate some values
for i ∈ 1:20
    append!(values, net_present_value(1, i+1, i, 0.10))
end
## display some values
@show values

# To make reading easier, we can create a table or `dataframe`.

using DataFrames
df = DataFrame(Dict(:year => 1:20, :revenue => values))
show(df)

# A quick glance reveals the maximum revenue.

maximum(values)
#+
using PyPlot;
figure(figsize=(8, 6));
plot(df.year, df.revenue);

xlabel("Year", fontsize=14);
ylabel("Revenue", fontsize=14);

plt.savefig(joinpath(@OUTPUT, "invsci-01-time-to-harvest.svg")); # hide

# \figalt{Time to Harvest}{invsci-01-time-to-harvest.svg}
#
# We can see that the graph levels off at the top, and upon inspection of the table we confirm that **years 9 and 10** will yield optimal revenue.


# Question 1
import Pkg; Pkg.add("Distributions"); Pkg.add("QuantEcon")
using Distributions, QuantEcon

function profit_max_q(a, c, mu, sigma, method, n)
        b = rand(LogNormal(mu, sigma), n)
        if method == "mc"
            q = Float64((a-c)/(2*mean(b)))
            return q
        elseif method == "quad"
            nodes_weights = qnwlogn(n, mu, sigma) #http://quantecon.github.io/QuantEcon.jl/latest/api/QuantEcon.html
            nodes = nodes_weights[1]
            weights = nodes_weights[2]
            mean_b = sum(nodes .* weights)
            q = Float64((a-c)/(2*mean_b))
        end
        return q
    end

profit_max_q(3,1,1,1,"mc",10)

# we dont have to worry about code_warntype
@code_warntype profit_max_q(3,1,1,1,"mc",10)

# Question 2
using Distributions
function estimate_pi(lower,upper,num_draws)
    n_land_in_circle = 0
    x = rand(Uniform(lower, upper), num_draws)
    y = rand(Uniform(lower, upper), num_draws)
        for i in 1:num_draws
        r_sq = x[i]^2 + y[i]^2
            if r_sq < 1
                n_land_in_circle += 1
            end
        end
    return n_land_in_circle / num_draws * 4
 end

estimate_pi(0,1,1000000)


"""
    jacobi(x, p::Integer, α, β)

Evaluate the Legendre polynomial with parameters `α`, `β` of degree `p` at `x`
using the three term recursion [Karniadakis and Sherwin, Spectral/hp Element
Methods for CFD, Appendix A].
"""
function jacobi(x, p::Integer, α, β)
    T = typeof( (2+α+β)*x / 2 )
    a = one(T)
    b = ((2+α+β)*x + α - β) / 2

    if p <= 0
        return a
    elseif p == 1
        return b
    end

    for n in 2:p
        a1 = 2n*(n+α+β)*(2n-2+α+β)
        a2 = (2n-1+α+β)*(α+β)*(α-β)
        a3 = (2n-2+α+β)*(2n-1+α+β)*(2n+α+β)
        a4 = 2*(n-1+α)*(n-1+β)*(2n+α+β)
        a, b = b, ( (a2+a3*x)*b - a4*a ) / a1
    end

    b
end

"""
    jacobi_and_derivative(x, p::Integer, α, β)

Evaluate the Jacobi polynomial with parameters `α`, `β` of degree `p` and its
derivative at `x` using the three term recursion [Karniadakis and Sherwin,
Spectral/hp Element Methods for CFD, Appendix A].
"""
function jacobi_and_derivative(x, p::Integer, α, β)
    T = typeof( (2+α+β)*x / 2 )
    # Coefficients for the polynomial...
    a = one(T)
    b = ((2+α+β)*x + α - β) / 2
    # ... and for the derivative. The derivative is computed via eq. (A.1.8)
    # of Karniadakis and Sherwin as Jacobi polynomial with parameters α+1, β+1.
    aa = one(T)
    bb = ((4+α+β)*x + α - β) / 2

    if p <= 0
        return a, zero(T)
    elseif p == 1
        return b, (2+α+β)/2
    end

    for n in 2:p
        a1 = 2n*(n+α+β)*(2n-2+α+β)
        a2 = (2n-1+α+β)*(α+β)*(α-β)
        a3 = (2n-2+α+β)*(2n-1+α+β)*(2n+α+β)
        a4 = 2*(n-1+α)*(n-1+β)*(2n+α+β)
        a, b = b, ( (a2+a3*x)*b - a4*a ) / a1

        b1 = 2n*(n+2+α+β)*(2n+α+β)
        b2 = (2n+1+α+β)*(α+β+2)*(α-β)
        b3 = (2n+α+β)*(2n+1+α+β)*(2n+2+α+β)
        b4 = 2*(n+α)*(n+β)*(2n+2+α+β)
        aa, bb = bb, ( (b2+b3*x)*bb - b4*aa ) / b1
    end

    b, (α+β+p+1)*aa/2
end

"""
    gauss_jacobi_nodes_and_weights(p, α, β, T=Float64::Type, tol=4*eps(T), maxit=100)

Compute the Gauss-Jacobi nodes and weights for polynomials of degree `p` with
parameters `α`, `β` using the scalar type `T`, tolerance `tol` and maximal number
of Newton iterations `maxit` [Karniadakis and Sherwin, Spectral/hp Element
Methods for CFD, Appendix B].
"""
function gauss_jacobi_nodes_and_weights(p, α, β, T=Float64::Type, tol=4*eps(T), maxit=1000)
    T = promote_type(typeof(α), typeof(β), T)
    nodes = Vector{T}(undef, p+1)
    weights = Vector{T}(undef, p+1)
    for j in 0:p
        x = -cospi(T(2j+1)/(2p+2))
        if j > 0
            x = (x + nodes[j]) / 2
        end
        for k in 1:maxit
            s = zero(T)
            for l in 1:j
                s += 1 / (x - nodes[l])
            end
            pol, der = jacobi_and_derivative(x, p+1, α, β)
            Δ = -pol / (der - s*pol)
            x = x + Δ
            abs(Δ) <= tol*abs(x) && break
        end
        pol, der = jacobi_and_derivative(x, p+1, α, β)
        nodes[j+1] = x
        weights[j+1] = 2^(α+β+1) * gamma(α+p+2) * gamma(β+p+2) /
                        (gamma(p+2) * gamma(α+β+p+2) * (1-x^2) * der^2)
    end

    return nodes, weights
end

"""
    jacobi_vandermonde(nodes, α, β)

Computes the Vandermonde matrix with respect to the Jacobi polynomials with
parameters `α`, `β` and the nodal basis on `nodes`.
The Vandermonde matrix `V` is the transformation matrix from the modal Jacobi
basis to the nodal Lagrange basis associated with `nodes`.
"""
function jacobi_vandermonde(nodes::AbstractVector, α, β)
    T = eltype(nodes)
    pp1 = length(nodes)
    V = Array{T}(undef, pp1, pp1)
    for j in 1:pp1, (i,x) in enumerate(nodes)
      V[i, j] = jacobi(x, j-1, α, β)
    end
    V
end

jacobi_vandermonde(basis::NodalBasis{Line}, α, β) = jacobi_vandermonde(basis.nodes, α, β)


function sum_n_below(factor, max)
    # Return the sum of all the natural numbers of factor below max
    max_nb = div(max-1,factor)
    ((factor / 2 * max_nb * (max_nb + 1)))
end

function n_below(n)
    # Return all the natural number of 3 and 5 multipliers
    return sum_n_below(3, n) + sum_n_below(5, n) - sum_n_below(3*5, n)
end

@assert n_below(10) == 23
println(n_below(1000))


"""
    newton(f::Function, x0::Real; niters=10^2, tol=1e-8)

Compute roots of scalar `f` by Newton's method

"""

function newton(f::Function, x0::Real; niters=10^2, tol=1e-8)
    x = xnew = x0
    for it in 1:niters
        df = ForwardDiff.derivative(f, x)
        xnew = x - f(x)/df
        abs(x - xnew) < tol && return xnew
        x, xnew = xnew, x
    end
    xnew
end


"""
Computes entropy of a sequence.
"""
function entropy(X :: AbstractVector)
  counts = Dict{eltype(X), Int}()
  n = length(X)
  for i = 1:length(X)
    counts[X[i]] = get(counts, X[i], 0) + 1
  end
  sum(kv -> log2(n / kv[2]) * kv[2], counts) / n
end

"""
Computes conditional entropy `H(Y | X)`.
"""
function conditional_entropy(X :: AbstractVector, Y :: AbstractVector)
  cnts_for_pref = Dict{eltype(X), Dict{eltype(Y), Int}}()
  n = length(X)
  for i = 1:length(X)
    if !haskey(cnts_for_pref, X[i])
      cnts_for_pref[X[i]] = Dict{eltype(Y), Int}()
    end
    counts = cnts_for_pref[X[i]]
    counts[Y[i]] = get(counts, Y[i], 0) + 1
  end

  # this helper function computes "unnormalized entropy" for given counts dict
  # i.e. entropy multiplied by sum of counts
  function _prefix_helper(counts :: Dict{eltype(Y), Int})
    counts_sum = sum(kv -> kv[2], counts)
    sum(kv -> log2(counts_sum / kv[2]) * kv[2], counts)
  end

  sum(kv -> _prefix_helper(kv[2]), cnts_for_pref) / n
end


#simple DFT function
function DFT(x)
    N = length(x)

    # We want two vectors here for real space (n) and frequency space (k)
    n = 0:N-1
    k = n'
    transform_matrix = exp.(-2im*pi*n*k/N)
    return transform_matrix*x

end

# Implementing the Cooley-Tukey Algorithm
function cooley_tukey(x)
    N = length(x)

    if (N > 2)
        x_odd = cooley_tukey(x[1:2:N])
        x_even = cooley_tukey(x[2:2:N])
    else
        x_odd = x[1]
        x_even = x[2]
    end
    n = 0:N-1
    half = div(N,2)
    factor = exp.(-2im*pi*n/N)
    return vcat(x_odd + x_even .* factor[1:half],
                x_odd - x_even .* factor[1:half])

end

function bitreverse(a::Array)
    # First, we need to find the necessary number of bits
    digits = convert(Int,ceil(log2(length(a))))

    indices = [i for i = 0:length(a)-1]

    bit_indices = []
    for i = 1:length(indices)
        push!(bit_indices, bits(indices[i]))
    end

    # Now stripping the unnecessary numbers
    for i = 1:length(bit_indices)
        bit_indices[i] = bit_indices[i][end-digits:end]
    end

    # Flipping the bits
    for i =1:length(bit_indices)
        bit_indices[i] = reverse(bit_indices[i])
    end

    #replacing indices
    for i = 1:length(indices)
        indices[i] = 0
        for j = 1:digits
            indices[i] += 2^(j-1) * parse(string(bit_indices[i][end-j]))
        end
       indices[i] += 1
    end

    b = [float(i) for i = 1:length(a)]
    for i = 1:length(indices)
        b[i] = a[indices[i]]
    end

    return b
end

function iterative_cooley_tukey(x)
    N = length(x)
    logN = convert(Int,ceil(log2(length(x))))
    bnum = div(N,2)
    stride = 0;

    x = bitreverse(x)

    z = [Complex(x[i]) for i = 1:length(x)]
    for i = 1:logN
       stride = div(N, bnum)
       for j = 0:bnum-1
           start_index = j*stride + 1
           y = butterfly(z[start_index:start_index + stride - 1])
           for k = 1:length(y)
               z[start_index+k-1] = y[k]
           end
       end
       bnum = div(bnum,2)
    end

    return z
end

function butterfly(x)
    N = length(x)
    half = div(N,2)
    n = [i for i = 0:N-1]
    half = div(N,2)
    factor = exp.(-2im*pi*n/N)

    y = [0 + 0.0im for i = 1:length(x)]

    for i = 1:half
        y[i] = x[i] + x[half+i]*factor[i]
        y[half+i] = x[i] - x[half+i]*factor[i]
    end

    return y
end

function approx(x, y)
    val = true
    for i = 1:length(x)
        if (abs(x[i]) - abs(y[i]) > 1e-5)
            val = false
        end
    end
    println(val)
end

function main()
    x = rand(128)
    y = cooley_tukey(x)
    z = iterative_cooley_tukey(x)
    w = fft(x)
    approx(y, w)
    approx(z, w)
end

main()


"""
barychebinit(n)

Get Chebyshev nodes of first kind and corresponding barycentric Lagrange interpolation weights.
Reference: Berrut, J.P. and Trefethen, L.N., 2004. Barycentric lagrange interpolation. SIAM review, 46(3), pp.501-517.

# Arguments
- `n`: order of the Chebyshev interpolation

# Returns
- Chebyshev nodes
- Barycentric Lagrange interpolation weights
"""
function barychebinit(n)
    x = zeros(Float64, n)
    w = similar(x)
    for i in 1:n
        c = (2i - 1)π / (2n)
        x[n - i + 1] = cos(c)
        w[n - i + 1] = (-1)^(i - 1) * sin(c)
    end
    return x, w
end

"""
function barycheb(n, x, f, wc, xc)

Barycentric Lagrange interpolation at Chebyshev nodes
Reference: Berrut, J.P. and Trefethen, L.N., 2004. Barycentric lagrange interpolation. SIAM review, 46(3), pp.501-517.

# Arguments
- `n`: order of the Chebyshev interpolation
- `x`: coordinate to interpolate
- `f`: array of size n, function at the Chebyshev nodes
- `wc`: array of size n, Barycentric Lagrange interpolation weights
- `xc`: array of size n, coordinates of Chebyshev nodes

# Returns
- Interpolation result
"""
function barycheb(n, x, f, wc, xc)
    for j in 1:n
        if x == xc[j]
            return f[j]
        end
    end

    num, den = 0.0, 0.0
    for j in 1:n
        q = wc[j] / (x - xc[j])
        num += q * f[j]
        den += q
    end
    return num / den
end

function barycheb2(n, x, f, wc, xc)
    for j in 1:n
        if x == xc[j]
            return f[j, :]
        end
    end

    den = 0.0
    num = zeros(eltype(f), size(f)[2])
    for j in 1:n
        q = wc[j] / (x - xc[j])
        num += q .* f[j, :]
        den += q
    end
    return num ./ den
end


using Polynomials: Polynomial,fit,coeffs
using Plots
using Distributions: Normal
using Random
using Statistics: mean, std

p = Polynomial([1,2,3])
@show coeffs(p)

errstd = 3
errNormal = Normal(0,errstd)

x = 0:10
y = 2 .+ 3 .* x .+rand(errNormal,11)

scatter(x,y,yerror=errstd,label=false)

fit_line = fit(x,y,1)
@show coeffs(fit_line)
plot!(x,fit_line.(x),label=false)

b_list = []
m_list = []
N = 1000
for i in 1:N
    y = 2 .+ 3 .* x .+rand(errNormal,11)
    fit_c = coeffs(fit(x,y,1))
    push!(b_list,fit_c[1])
    push!(m_list,fit_c[2])
end

@show mean(b_list)
@show mean(m_list)
@show std(b_list)
@show std(m_list)

histogram(b_list,bins=30)
histogram(m_list,bins=30)


#=~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~=#
# Problem set 7 solutions
# Written by Tyler Ransom
# Commented by Giuseppe Grasso
#=~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~=#

using Random
using LinearAlgebra
using Statistics
using Distributions
using Optim
using DataFrames
using CSV
using HTTP
using GLM
using FreqTables
using ForwardDiff
using LineSearches
using SMM
cd("/Users/peppegrass/Documents/GitHub/fall-2020/ProblemSets/PS6-ccp/") # GG: sets working directory
pwd() ## GG: prints working directory
readdir() # GG: equivalent to -ls- to see elements of working directory

function allwrap()

    #=
    GG: In this problem set, we will practice estimating models by
    Generalized Method of Moments (GMM) and Simulated Method of Moments (SMM).
    =#

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 1 - GMM for simple regression model
    # data & model are SAME AS PS2, QUESTION 2
    #
    # 1. Estimate the linear regression model from Question 2 of Problem Set 2 by GMM.
    # Write down the moment function as in slide #8 of the Lecture 9 slide deck and use Optim for estimation.
    # Use the N × N Identity matrix as your weighting matrix. Check your answer using the closed-form matrix formula for the OLS estimator.
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    url = "https://raw.githubusercontent.com/OU-PhD-Econometrics/fall-2020/master/ProblemSets/PS1-julia-intro/nlsw88.csv"
    df = CSV.read(HTTP.get(url).body)
    X = [ones(size(df,1),1) df.age df.race.==1 df.collgrad.==1]
    y = log.(df.wage)

    function ols_gmm(β, X, y) # GG: GMM function
        K = size(X,2) # GG: don't really needed for GMM, unless for computing RMSE in step below
        N = size(y,1) # GG: don't really needed for GMM, unless for computing RMSE in step below
        ŷ = X*β
        g = y .- ŷ # GG: force y and y_hat to be as close as possible
        J = g'*I*g # GG: GMM objective function
        return J
    end

    # we can also estimate the RMSE of the regression by ADDING ANOTHER MOMENT (and assuming homoskedasticity)
    function ols_gmm_with_σ(θ, X, y)
        K = size(X,2)
        N = size(y,1)
        β = θ[1:end-1]
        σ = θ[end]
        g = y .- X*β
        g = vcat(g,( (N-1)/(N-K) )*var(g) .- σ^2) # this builds on the assumption that MSE: SSR/(N-K) = σ² # GG: performing DOFs corrections
        J = g'*I*g
        return J
    end

    β_hat_gmm = optimize(b -> ols_gmm(b, X, y), rand(size(X,2)), LBFGS(), Optim.Options(g_tol=1e-6, iterations=100_000, show_trace=true))
    println(β_hat_gmm.minimizer)
    println("ans: ",X\y) # GG: X\y gives the least squares; cool uh!?
    β_hat_gmm = optimize(b -> ols_gmm_with_σ(b, X, y), rand(size(X,2)+1), LBFGS(), Optim.Options(g_tol=1e-6, iterations=100_000, show_trace=true))
    println(β_hat_gmm.minimizer)
    println("ans: ",vcat(X\y,sqrt( sum( (y .- X*β_hat_gmm.minimizer[1:end-1]).^2 )/( size(y,1)-size(X,2) ) )))

    # GG: of course ols_gmm(b, X, y) and ols_gmm_with_σ(b, X, y) yield same results but latter with MSE

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 2 - GMM for multinomial logit model
    # data & model are SAME AS PS2, QUESTION 5
    #
    # 2. Estimate the multinomial logit model from Question 5 of Problem Set 2 by the following means:
    #   (a) Maximum likelihood (i.e. re-run your code [or mine] from Question 5 of Problem Set 2)
    #   (b) GMM with the MLE estimates as starting values. Your g object should be a vector of dimension N ×J where
    #   N is the number of rows of the X matrix and J is the dimension of the choice set. Each element, g should equal d − P,
    #   where d and P are “stacked” vectors of dimension N×J
    #   (c) GMM with random starting values
    # Compare your estimates from part (b) and (c). Is the objective function globally concave?
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    df = dropmissing(df, :occupation)
    df[df.occupation.==8 ,:occupation] .= 7
    df[df.occupation.==9 ,:occupation] .= 7
    df[df.occupation.==10,:occupation] .= 7
    df[df.occupation.==11,:occupation] .= 7
    df[df.occupation.==12,:occupation] .= 7
    df[df.occupation.==13,:occupation] .= 7

    #df = @transform(df, white = race.==1)
    df[:,:white] = df.race .== 1
    X = [ones(size(df,1),1) df.age df.white df.collgrad]
    y = df.occupation

    # get starting values for mlogit from series of binary logits
    svals = zeros(size(X,2),7)
    for j=1:7
        tempname = Symbol(string("occ",j))
        df[:,tempname] = df.occupation.==j
    end
    # GG: this is to derive some "reasonable starting values" for a multinomial logit model
    # GG: what do we mean for "reasonable starting values". You don't wanna have neither a vector of 0 nor a vector of random numbers
    # GG: So what starting values for a Multinomial logit. Estimate a bunch of binary logits and then...
    svals[:,1] = coef(lm(@formula(occ1 ~ age + white + collgrad), df))
    svals[:,2] = coef(lm(@formula(occ2 ~ age + white + collgrad), df))
    svals[:,3] = coef(lm(@formula(occ3 ~ age + white + collgrad), df))
    svals[:,4] = coef(lm(@formula(occ4 ~ age + white + collgrad), df))
    svals[:,5] = coef(lm(@formula(occ5 ~ age + white + collgrad), df))
    svals[:,6] = coef(lm(@formula(occ6 ~ age + white + collgrad), df))
    svals[:,7] = coef(lm(@formula(occ7 ~ age + white + collgrad), df))
    svals = svals[:,1:6] .- svals[:,7] # GG: ... impose the identification normalization (i.e. you subtract for the first 6 estimates the one of the 7th)
    svals = svals[:] # GG: then vectorize it and use it as a starting value

    function mlogit_mle(α, X, y)

        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j
        end
        bigα = [reshape(α,K,J-1) zeros(K)]

        P = exp.(X*bigα) ./ sum.(eachrow(exp.(X*bigα)))

        loglike = -sum( bigY.*log.(P) ) # GG: MLE objective function

        return loglike
    end

    α_zero = zeros(6*size(X,2))
    α_rand = rand(6*size(X,2))
    # α_true = [-.5120867,-.6950393,.3165616,-1.646016,-2.300697,-.7656569]
    α_true = [.1910213,-.0335262,.5963968,.4165052,-.1698368,-.0359784,1.30684,-.430997,.6894727,-.0104578,.5231634,-1.492475,-2.26748,-.0053001,1.391402,-.9849661,-1.398468,-.0142969,-.0176531,-1.495123,.2454891,-.0067267,-.5382892,-3.78975]
    α_start = α_true.*rand(size(α_true))
    println(size(α_true))
    α_hat_optim = optimize(a -> mlogit_mle(a, X, y), svals, LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
    α_hat_mle = α_hat_optim.minimizer
    println("mlogit MLE estimates: ",α_hat_mle)

    function mlogit_gmm(α, X, y)
        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j
        end
        bigα = [reshape(α,K,J-1) zeros(K)]

        P = exp.(X*bigα) ./ sum.(eachrow(exp.(X*bigα)))

        g = zeros((J-1)*K)
        for j=1:(J-1)
            for k = 1:K
                g[(j-1)*K+k] = mean( (bigY[:,j] .- P[:,j]) .* X[:,k] ) # GG: GMM (exactly-identified) objective function
                # GG: (j-1)*K moment conditions; (j-1) equations and K betas; same number of elements as number non-zero elements in bigαlpha
            end
        end
        J = N*g⋅g #GG: exactly identified; no weighting matrix
        return J
    end

    function mlogit_gmm_overid(α, X, y)
        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j
        end
        bigα = [reshape(α,K,J-1) zeros(K)]

        P = exp.(X*bigα) ./ sum.(eachrow(exp.(X*bigα))) # GG: same as before P = exp.(X*bigα) ./ sum.(eachrow(exp.(X*bigα)))

        g = bigY[:] .- P[:] # GG: now in over-identified case; we vectorize big Y and P with the [:] operator and subtract the two
        # GG: this will be equivalent to NONLINEAR LEAST SQUARES

        J = g'*I*g # GG: now we have weighting matrix
        return J
    end
    #α_zero = zeros(6*size(X,2))
    #α_rand = rand(6*size(X,2))
    α_true = [.1910213,-.0335262,.5963968,.4165052,-.1698368,-.0359784,1.30684,-.430997,.6894727,-.0104578,.5231634,-1.492475,-2.26748,-.0053001,1.391402,-.9849661,-1.398468,-.0142969,-.0176531,-1.495123,.2454891,-.0067267,-.5382892,-3.78975]
    #α_start = α_true.*rand(size(α_true))
    #println("starting values: ",α_rand)

    # GG: here we're playing around with different optimization algorithms, starting values, and later using autodiff

    # LBFGS
    α_hat_optim = optimize(a -> mlogit_gmm_overid(a, X, y), α_true .+ .0001*rand(size(α_true)), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=1_000, show_trace=true, show_every=100))
    α_hat_gmm = α_hat_optim.minimizer
    println("LBFGS GMM minimum: ",α_hat_optim.minimum)
    println("LBFGS GMM estimates: ",α_hat_gmm)

    # use autodiff with LBFGS
    td = TwiceDifferentiable(b -> mlogit_gmm_overid(b, X, y), svals; autodiff = :forward)
    α_hat_optim_ad = optimize(td, α_true .+ .0001*rand(size(α_true)), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=10_000, show_trace=true, show_every=100))
    α_hat_gmm_ad = α_hat_optim_ad.minimizer
    println("LBFGS GMM autodiff minimum: ",α_hat_optim_ad.minimum)
    println("LBFGS GMM autodiff estimates: ",α_hat_gmm_ad)

    # use autodiff with Newton
    α_hat_optim_ad_newton = optimize(td, α_true .+ .0001*rand(size(α_true)), Newton(linesearch = BackTracking()), Optim.Options(g_tol = 1e-5, iterations=10_000, show_trace=true, show_every=100))
    α_hat_gmm_ad_newton = α_hat_optim_ad_newton.minimizer
    println("Newton GMM autodiff minimum: ",α_hat_optim_ad_newton.minimum)
    println("Newton GMM autodiff estimates: ",α_hat_gmm_ad_newton)

    # autodiff from bad starting values
    α_hat_optim_ad_fr_bad = optimize(td, α_rand.*sign.(α_true), LBFGS(), Optim.Options(g_tol = 1e-10, f_tol = 1e-10, x_tol = 1e-10, iterations=10_000, show_trace=true, show_every=50))
    α_hat_gmm_ad_fr_bad = α_hat_optim_ad_fr_bad.minimizer
    println("LBFGS GMM autodiff minimum: ",α_hat_optim_ad_fr_bad.minimum)
    println("LBFGS GMM autodiff estimates: ",α_hat_gmm_ad_fr_bad)

    # GG: putting different answers in a DataFrame and comparing them
    # GG: found out close answers; the one using bad starting value didn't get close to other solutions; some times GMM tolerance values needed to be played around with
    compare = DataFrame(mle = α_hat_mle, gmm = α_hat_gmm, gmm_ad = α_hat_gmm_ad, gmm_ad_newton = α_hat_gmm_ad_newton, gmm_ad_fr_bad = α_hat_gmm_ad_fr_bad)
    println(compare)
    # GG: objective function globally concave, but couldn't get to global maximum from random starting values, although pretty close

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 3 - SIMULATE DATA from a logit model
    #
    # 3. SIMULATE A DATASET FROM A MULTINOMIAL LOGIT MODEL, and then estimate its parameter values and verify that the estimates are close to the
    # parameter values you set. That is, for a given sample size N, choice set dimension J and parameter vector β, write a function that outputs
    # data X and Y. I will let you choose N, J, β and the number of covariates in X (K), but J should be larger than 2 and K should be larger than 1.
    # If you haven’t done this before, you may want to follow these steps:
    #   (a) Generate X using a random number generator—rand() or randn().
    #   (b) Set values for β such that conformability with X and J is satisﬁed
    #   (c) Generate the N×J matrix of choice probabilities P
    #   (d) Draw the preference shocks ε as a N×1 vector of U[0,1] random numbers
    #   (e) Generate Y as follows:
    #     • Initialize Y as an N ×1 vector of 0s
    #     • Update Y i = ∑ j=1 J 1 P ik > εi  ∑k=j J  [see text of problem set]
    #   (f) An alternative way to generate choices would be to draw a N × J matrix of ε’s from a T1EV distribution.
    #   This distribution is already deﬁned in the Distributions package. Then Y i = argmax j X i β j + ε i j .
    #   I’ll show you an example of how to do that in the solutions code for this problem set.
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    function sim_logit(N=100_000,J=4) # putting "=" after the arguments here sets default values for these arguments
        # generate X, β and P
        X = hcat(ones(N),randn(N),2 .+ 2 .* randn(N))
        if J==4
            β = hcat([1, -1, 0.5],[-2, 0.5, 0.3],[0, -0.5, 2],zeros(3)) # GG: numbers that square nicely with 4
        else
            β = 2 .* rand(size(X,2),J) .- 1
        end
        P = exp.(X*β) ./ sum.(eachrow(exp.(X*β))) # GG: usual one-liner to generate multinomial ogit choice probabilities; eachrow is summation
        # draw choices
        draw = rand(N)
        Y = zeros(N) # Initializing Y as an Nx1 vector of 0s
        for j=1:J
            Ytemp = sum(P[:,j:J];dims=2) .> draw # GG: then updating it (summing across j) based on whether choice probability is above or below draw [see formula of problem set]
            Y += Ytemp # GG: for each iteration of the loop adding Ytemp to Y
        end
        return Y,X
    end
    ySim,XSim = sim_logit() # GG: it will use "default values" for N and J set in line 244 through s"et to the default operator =""
    α_hat_sim_optim = optimize(a -> mlogit_mle(a, XSim, ySim), rand(9), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
    α_hat_sim_mle = α_hat_sim_optim.minimizer
    println("simulated mlogit MLE estimates: ",α_hat_sim_mle) # GG: and here are our simulated estimates; values pretty close because sample size big and SEs low

    # now with drawing from Gumbel instead (#GG: i.e T1EV)
    # GG: python does not have the nice Distributions package as in Julia
    function sim_logit_w_gumbel(N=100_000,J=4) # putting = here sets default values for these arguments
        # generate X and β (No P needed)
        X = hcat(ones(N),randn(N),2 .+ 2 .* randn(N))
        if J==4
            β = hcat([1, -1, 0.5],[-2, 0.5, 0.3],[0, -0.5, 2],zeros(3))
        else
            β = 2 .* rand(size(X,2),J) .- 1
        end
        # draw choices
        ϵ = rand(Gumbel(0,1),N,J) # GG: using Julia's Distributions package; Gumbel(0,1) is "standard" T1EV
        Y = argmax.(eachrow(X*β .+ ϵ)) # GG: getting argmax of eachrow; eachrow allows avoiding to have to write a loop; tells Julia to do argmax computation row-wise and find column j maximizer
        # GG: that's just another way of deriving the Y; the loop in the previous function is just as fine
        return Y,X
    end
    ySim,XSim = sim_logit_w_gumbel()
    α_hat_sim_optim = optimize(a -> mlogit_mle(a, XSim, ySim), rand(9), LBFGS(), Optim.Options(g_tol = 1e-5, iterations=100_000, show_trace=true, show_every=50))
    α_hat_sim_mle = α_hat_sim_optim.minimizer
    println("simulated mlogit MLE estimates: ",α_hat_sim_mle)

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 4 - run SMM.parallelNormal()
    #
    # 4. Use SMM.jl to run the example code on slide #21 of the Lecture 9 slide deck.
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # GG: SMM.jl has a nice count-down that Optim does not have
    MA = SMM.parallelNormal()
    dc = SMM.history(MA.chains[1])
    dc = dc[dc[:accepted].==true, :]
    println(describe(dc)) # GG: can verify you get the mean of the MVN distribution

    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # Question 5 - estimate mlogit by SMM
    # GG: Wanna use SIMULATED METHOD OF MOMENTS (using code of Q3) to ESTIMATE RESULT OF Q2
    #
    # 5. Use your code from Question 3 to estimate the multinomial logit model from Question 2 using SMM and the code example from slide #18 of
    # the Lecture 9 slide deck.
    #:::::::::::::::::::::::::::::::::::::::::::::::::::
    # GG: X and y of Q2
    X = [ones(size(df,1),1) df.age df.white df.collgrad]
    y = df.occupation
    # GG: chunk below was commented out by Tyler beacuse the set of moment conditions used were just not working out for him
    #function mlogit_smm(α, X, y, D)
    #    K = size(X,2)
    #    J = length(unique(y))
    #    N = length(y)
    #    bigY = zeros(N,J)
    #    bigỸ = zeros(N,J)
    #    for j=1:J
    #        bigY[:,j] = y.==j
    #    end
    #    bigα = [reshape(α,K,J-1) zeros(K)]

    #    T = promote_type(eltype(α),eltype(X))
    #    gmodel = zeros(T,(J-1)*K,D)
    #    gdata  = zeros(T,(J-1)*K)

    #    # data moments
    #    for j=1:(J-1)
    #        gdata[(j-1)*K+1] = mean( (bigY[:,j]) )
    #        for k = 2:K
    #            gdata[(j-1)*K+k] = cov( bigY[:,j], X[:,k] )
    #        end
    #    end

    #    Random.seed!(1234)
    #    # simulated model moments
    #    for d=1:D
    #        # draw choices
    #        ε = rand(Gumbel(0,1),N,J)
    #        ỹ = argmax.(eachrow(X*bigα .+ ε))
    #        for j=1:J
    #            bigỸ[:,j] = ỹ.==j
    #        end
    #        # data moments
    #        for j=1:(J-1)
    #            gmodel[(j-1)*K+1,d] = mean( (bigỸ[:,j]) )
    #            for k = 2:K
    #                gmodel[(j-1)*K+k,d] = cov( bigỸ[:,j], X[:,k] )
    #            end
    #        end
    #    end

    #    # minimize squared difference between data and moments
    #    g = mean(gmodel .- gdata; dims=2)
    #    J = N*g⋅g
    #    return J
    #

    # GG: the function below for the SIMULATED METHOD OF MOMENTS is the one from the SLIDES OF LECTURE 9 (mind the SCROLLER; gotta use HTML and not PDF version)
    function mlogit_smm_overid(α, X, y, D) # GG: D is the number of draws
        K = size(X,2)
        J = length(unique(y))
        N = length(y)
        bigY = zeros(N,J)
        bigỸ = zeros(N,J)
        for j=1:J
            bigY[:,j] = y.==j # GG: bigY is our real data that we're trying to MATCH
        end
        bigα = [reshape(α,K,J-1) zeros(K)] # GG: bigα is our parameter vector that we're gonna try to guess value of, so that we get values so that simulated data match actual ones as closely as possible

        Random.seed!(1234) # GG: cardinal rule, set up the seed; don't forget that!
        for d=1:D
            # draw choices
            ε = rand(Gumbel(0,1),N,J) # GG: if we change the distribution of ε to a multivariate normal, we can do MNP without having to do any Monte Carlo integration. We don't need to have a closed form solution for this.
            # GG: and if we want to do GEV models, we need to specify nest-specific and alternative-specific component; just two errors but same principle.
            ỹ = argmax.(eachrow(X*bigα .+ ε)) # GG: we're goona draw ỹ through ε
            for j=1:J
                bigỸ[:,j] .+= (ỹ.==j)*(1/D) # GG: bigỹ is our SIMULATED data
                # GG: we set the jth column of bigỸ to a dummy if ỹ is equal to j
            end
        end

        g = bigY[:] .- bigỸ[:] # GG: exactly like g in Question 2; very same principle

        J = g'*I*g
        return J
    end
    println(size(X,2))
    println(length(unique(y))-1)
    println("size of svals: ",size(α_true .+ .0001*rand(size(α_true))))
    bigα = [reshape(α_true,size(X,2),length(unique(y))-1) zeros(size(X,2))]
    println("size of bigα: ",size(bigα))
    td = TwiceDifferentiable(th -> mlogit_smm_overid(th, X, y, 2_000), α_true .+ .0001*rand(size(α_true)); autodiff = :forward)
    α_hat_smm_overid = optimize(td, α_true .+ .0001*rand(size(α_true)), LBFGS(), Optim.Options(g_tol=1e-8, x_tol=1e-8, f_tol=1e-8, iterations=100_000, show_trace=true, show_every=5))
    println(α_hat_smm_overid.minimizer)
    return nothing
end

#:::::::::::::::::::::::::::::::::::::::::::::::::::
# Question 6
#:::::::::::::::::::::::::::::::::::::::::::::::::::
@time allwrap()


# A variety of means

# Geometric mean
function geomean(a::RealArray)
    s = 0.0
    n = length(a)
    for i = 1 : n
        @inbounds s += log(a[i])
    end
    return exp(s / n)
end

# Harmonic mean
function harmmean(a::RealArray)
    s = 0.0
    n = length(a)
    for i in 1 : n
        @inbounds s += inv(a[i])
    end
    return n / s
end

# Trimmed mean

function trimmean(x::RealArray, p::FloatingPoint)
    n = length(x)
    n > 0 || error("x can not be empty.")
    rn = min(iround(n * p), n-1)
    if rn == n - 1
        return median(x)
    else
        sx = sort(x)
        nl = rn >> 1
        nh = (rn - nl)
        s = 0.0
        @inbounds for i = (1+nl) : (n-nh)
            s += x[i]
        end
        return s / (n - rn)
    end
end

# Weighted mean
function wmean{T<:Number,W<:Real}(v::AbstractArray{T}, w::AbstractArray{W}; wsum::W=NaN)
	# wsum is the pre-computed sum of weights

	n = length(v)
	length(w) == n || raise_dimerror()
    sv = zero(T)

    if isnan(wsum)
    	sw = zero(W)
    	for i = 1 : n
    		@inbounds wi = w[i]
        	@inbounds sv += v[i] * wi
        	sw += wi
    	end
    else
    	for i = 1 : n
    		@inbounds sv += v[i] * w[i]
    	end
    	sw = wsum
    end

    return sv / sw
end

# Weighted mean (fast version for contiguous arrays using BLAS)
function wmean{T<:BlasReal}(v::Array{T}, w::Array{T}; wsum::T=NaN)
	sw = (isnan(wsum) ? sum(w) : wsum)
	dot(v, w) / sw
end

mean{T<:Number}(v::AbstractArray{T}, w::WeightVec) = wmean(v, values(w); wsum=sum(w))


@doc raw"""
    ExponentiatedKernel([a=1])

The exponentiated kernel is a Mercer kernel given by:

```math
\kappa(\mathbf{x},\mathbf{y}) = \exp\left(a \mathbf{x}^\intercal \mathbf{y} \right)
\qquad a > 0
```
"""
struct ExponentiatedKernel{T<:AbstractFloat} <: MercerKernel{T}
    alpha::HyperParameter{T}
    ExponentiatedKernel{T}(α::Real) where {T<:AbstractFloat} = new{T}(
        HyperParameter(convert(T,α), interval(OpenBound(zero(T)), nothing))
    )
end
ExponentiatedKernel(α::T1 = 1.0) where {T1<:Real} = ExponentiatedKernel{floattype(T1)}(α)

@inline exponentiatedkernel(z::T, α::T) where {T<:AbstractFloat} = exp(α*z)

@inline basefunction(::ExponentiatedKernel) = ScalarProduct()
@inline kappa(κ::ExponentiatedKernel{T}, z::T) where {T} = exponentiatedkernel(z, getvalue(κ.alpha))


"""
    M = time_to_mean(t, T, Mo=0., epoch=0.)

Computes the mean anomaly of an orbit with period `T` from a time `t`, the orbit's mean anomaly
at epoch `Mo`, and the epoch.
"""
function time_to_mean(t, T, Mo=0., epoch=0.)
    return Mo + 2π*(t - epoch)/T
end

function time_to_mean(t, orb::Orbit)
    M = time_to_mean(t, orb.period, orb.Mo, orb.epoch)
    if orb.e < 1
        return wrap_angle(M)
    else
        return M
    end
end

"""
    t = mean_to_time(M, T, Mo=0, epoch=0, tmin=nothing)

Computes the time at which an orbit with period `T`, mean anomaly
at epoch `Mo`, and epoch `epoch` has the specified mean anomaly `M`.

Because this time is not uniquely defined for elliptical orbits, `tmin` can be specified
to ensure that the returned time falls between `tmin` and `tmin + T`.
"""
function mean_to_time(M, T, Mo=0, epoch=0, tmin=nothing)
    t = epoch + (M-Mo)*T/(2π)
    if tmin === nothing
        return t
    else
        return t + T*ceil((tmin-t)/T)
    end
end

mean_to_time(M, orb::Orbit, tmin=nothing) = mean_to_time(M, orb.period, orb.Mo, orb.epoch, orb.e<1 ? tmin : nothing)

"""
    θ = mean_to_true(M, e)

Computes the true anomaly of an orbit with eccentricity `e` at the specified mean anomaly `M`.
"""
function mean_to_true(M, e)
    if e == 1       # Parabolic case
        throw(ArgumentError("Parabolic case (e=1) not implemented"))
    elseif e < 1    # Elliptical case
        E = kepler(M, e)
        return 2*atan(√(1+e)*sin(E/2),  √(1-e)*cos(E/2))
    else            # Hyperbolic case
        H = kepler(M, e)
        return 2*atan(√(e+1)*sinh(H/2), √(e-1)*cosh(H/2))
    end
end

"""
    θ = time_to_true(t, e, T, Mo=0., epoch=0.)

Computes the mean anomaly of an orbit with eccentricity `e` and period `T` from a time `t`, the orbit's mean anomaly
at epoch `Mo`, and the epoch.
"""
function time_to_true(t, e, T, Mo=0, epoch=0)
    M = time_to_mean(t, T, Mo, epoch)
    return mean_to_true(M, e)
end

time_to_true(t, orb::Orbit) = time_to_true(t, orb.e, orb.period, orb.Mo, orb.epoch)

"""
    M = true_to_mean(θ, e)

Computes the mean anomaly of an orbit with eccentricity `e` at the specified true anomaly `θ`.
"""
function true_to_mean(θ, e)
    if e == 1       # Parabolic case
        throw(ArgumentError("Parabolic case (e=1) not implemented"))
    elseif e < 1    # Elliptical case
        E = 2*atan(sin(θ/2)*√(1-e), cos(θ/2)*√(1+e))
        return wrap_angle(E - e*sin(E))
    else            # Hyperbolic case
        H = 2*atanh(tan(θ/2)*√((e-1)/(e+1)))
        # H = copysign(acosh((e+cos(θ))/(1+e*cos(θ))), θ)
        return e*sinh(H) - H
    end
end

true_to_mean(θ, orb::Orbit) = true_to_mean(θ, orb.e)

"""
    t = true_to_time(θ, e, T, Mo=0, epoch=0, tmin=nothing)

Computes the time at which an orbit with eccentricity `e`, period `T`, mean anomaly
at epoch `Mo`, and epoch `epoch` has the specified true anomaly `θ`.

Because this time is not uniquely defined for elliptical orbits, `tmin` can be specified
to ensure that the returned time falls between `tmin` and `tmin + T`.
"""
function true_to_time(θ, e, T, Mo=0, epoch=0, tmin=nothing)
    M = true_to_mean(θ, e)
    return mean_to_time(M, T, Mo, epoch, e<1 ? tmin : nothing)
end

# true_to_time(θ, e, a, μ, Mo=0, epoch=0, tmin=nothing) = true_to_time(θ, e, period(a, μ), Mo, epoch, tmin)
true_to_time(θ, orb::Orbit, tmin=nothing) = true_to_time(θ, orb.e, orb.period, orb.Mo, orb.epoch, tmin)


using PyPlot
using OffsetArrays

function bary_grid(n, a)
    x = OffsetArray{Float64}(undef, 0:n, 0:n)
    y = OffsetArray{Float64}(undef, 0:n, 0:n)
    pts = OffsetArray(range(0, 1, length=n+1), 0:n)
    for i = 0:n
        ξ₁ = pts[i]
        for j = 0:n-i
            ξ₂ = pts[j]
            ξ₃ = 1 - ξ₁ - ξ₂
            x[i,j] = ξ₁ * a[1,1] + ξ₂ * a[1,2] + ξ₃ * a[1,3]
            y[i,j] = ξ₁ * a[2,1] + ξ₂ * a[2,2] + ξ₃ * a[2,3]
        end
        for j = n-i+1:n
            x[i,j] = NaN
            y[i,j] = NaN
        end
    end
    return x, y
end

function bary_func(f, n, a)
    z = OffsetArray{Float64}(undef, 0:n, 0:n)
    pts = OffsetArray(range(0, 1, length=n+1), 0:n)
    ξ = Vector{Float64}(undef, 3)
    for i = 0:n
        ξ[1] = pts[i]
        for j = 0:n-i
            ξ[2] = pts[j]
            ξ[3] = 1 - ξ[1] - ξ[2]
            z[i,j] = f(ξ)
        end
    end
    return z
end

function bary_plot(x, y, z)
    n = size(x, 1) - 1
    xd = [ x[i,n-i] for i = 0:n ]
    yd = [ y[i,n-i] for i = 0:n ]
    zd = [ z[i,n-i] for i = 0:n ]
    idx = [1, 2, 3, 1]
    plot3D(a[1,idx], a[2,idx], zeros(4), "k")
    for i = 0:n
        plot3D(x[i,0:n-i], y[i,0:n-i], z[i,0:n-i], "c")
    end
    for j = 0:n
        plot3D(x[0:n-j,j], y[0:n-j,j], z[0:n-j,j], "c")
    end
    for i = 0:n
        plot3D(xd, yd, zd, "c")
    end
    axis("off")
#    xlabel(L"$x$")
#    ylabel(L"$y$")
end

function midpoints(a)
    prev = [ 3, 1, 2 ]
    succ = [ 2, 3, 1 ]
    m = Matrix{Float64}(undef, 2, 3)
    for i = 1:3
        m[:,i] = ( a[:,prev[i]] + a[:,succ[i]] ) / 2
    end
    return m
end

ψ₁(ξ) = 2 * ξ[1] * ( ξ[1] - 0.5 )
ψ₂(ξ) = 2 * ξ[2] * ( ξ[2] - 0.5 )
ψ₃(ξ) = 2 * ξ[3] * ( ξ[3] - 0.5 )
ψ₄(ξ) = 4 * ξ[2] * ξ[3]
ψ₅(ξ) = 4 * ξ[3] * ξ[1]
ψ₆(ξ) = 4 * ξ[1] * ξ[2]

a = Float64[ 0  1  2
             2  0  3 ]
m = midpoints(a)
n = 20
x, y = bary_grid(n, a)

figure(1)
z = bary_func(ψ₁, n, a)
bary_plot(x, y, z)
plot3D(a[1,1], a[2,1], "ro")
savefig("psi1.pdf")

figure(2)
z = bary_func(ψ₂, n, a)
bary_plot(x, y, z)
plot3D(a[1,2], a[2,2], "ro")
savefig("psi2.pdf")

figure(3)
z = bary_func(ψ₃, n, a)
bary_plot(x, y, z)
plot3D(a[1,3], a[2,3], "ro")
savefig("psi3.pdf")

figure(4)
z = bary_func(ψ₄, n, a)
bary_plot(x, y, z)
plot3D(m[1,1], m[2,1], "ro")
savefig("psi4.pdf")

figure(5)
z = bary_func(ψ₅, n, a)
bary_plot(x, y, z)
plot3D(m[1,2], m[2,2], "ro")
savefig("psi5.pdf")

figure(6)
z = bary_func(ψ₆, n, a)
bary_plot(x, y, z)
plot3D(m[1,3], m[2,3], "ro")
savefig("psi6.pdf")