text:
using StatsPlots
using CSV
using DataFrames

jt = CSV.read(joinpath(@__DIR__, "julia_num_apertures.csv")) |> DataFrame
pt = CSV.read(joinpath(@__DIR__, "python_num_apertures.csv")) |> DataFrame

jt_ell = CSV.read(joinpath(@__DIR__, "julia_num_apertures-ellipse.csv")) |> DataFrame
pt_ell = CSV.read(joinpath(@__DIR__, "python_num_apertures-ellipse.csv")) |> DataFrame

plot(markerstrokealpha=0, xscale=:log10, yscale=:log10, link=:y, legend=:topleft, layout=2, size=(800, 300), dpi=100, bottom_margin=5StatsPlots.mm)

@df pt plot!(:N, :time, c=2, label="photutils", sp=1)
jt_group = groupby(jt, :nt)
@df jt_group[(nt=1,)] plot!(:N, :time, markerstrokealpha=0, c=1, label="Photometry.jl - 1 thread", sp=1)
@df jt_group[(nt=4,)] plot!(:N, :time, c=3,  label="Photometry.jl - 4 threads", sp=1)
@df jt_group[(nt=8,)] plot!(:N, :time, c=3, ls=:dash, label="Photometry.jl - 8 threads", sp=1)
title!("CircularAperture(r=10)", sp=1)
ylabel!("time [s]", sp=1)
xlabel!("number of apertures", sp=1)

@df pt_ell plot!(:N, :time, c=2, label="", sp=2)
jt_ell_group = groupby(jt, :nt)
@df jt_ell_group[(nt=1,)] plot!(:N, :time, markerstrokealpha=0, c=1, label="", sp=2)
@df jt_ell_group[(nt=4,)] plot!(:N, :time, c=3,  label="", sp=2)
@df jt_ell_group[(nt=8,)] plot!(:N, :time, c=3, ls=:dash, label="", sp=2)
title!("EllipticalAperture(a=10, b=10)", sp=2)
xlabel!("number of apertures", sp=2)

savefig(joinpath(@__DIR__, "..", "..", "docs", "src", "assets", "num_apertures_benchmark.png"))

nothing


text:
# Text provided under a Creative Commons Attribution license, CC-BY.  All code is made available under the FSF-approved BSD-3 license.  (c) Lorena A. Barba, Gilbert F. Forsyth 2017. Thanks to NSF for support via CAREER award #1149784.
# [@LorenaABarba](https://twitter.com/LorenaABarba)
# Translation to Julia by [@miguelraz_](https://twitter.com/miguelraz_)

# 12 steps to Navier–Stokes
# =====
# ***

# Did you make it this far? This is the last step! How long did it take you to write your own Navier–Stokes solver in Python following this interactive module? Let us know!

# Step 12: Channel Flow with Navier–Stokes
# ----
# ***

# The only difference between this final step and Step 11 is that we are going to add a source term to the $u$-momentum equation, to mimic the effect of a pressure-driven channel flow. Here are our modified Navier–Stokes equations:

# $$\frac{\partial u}{\partial t}+u\frac{\partial u}{\partial x}+v\frac{\partial u}{\partial y}=-\frac{1}{\rho}\frac{\partial p}{\partial x}+\nu\left(\frac{\partial^2 u}{\partial x^2}+\frac{\partial^2 u}{\partial y^2}\right)+F$$

# $$\frac{\partial v}{\partial t}+u\frac{\partial v}{\partial x}+v\frac{\partial v}{\partial y}=-\frac{1}{\rho}\frac{\partial p}{\partial y}+\nu\left(\frac{\partial^2 v}{\partial x^2}+\frac{\partial^2 v}{\partial y^2}\right)$$

# $$\frac{\partial^2 p}{\partial x^2}+\frac{\partial^2 p}{\partial y^2}=-\rho\left(\frac{\partial u}{\partial x}\frac{\partial u}{\partial x}+2\frac{\partial u}{\partial y}\frac{\partial v}{\partial x}+\frac{\partial v}{\partial y}\frac{\partial v}{\partial y}\right)
# $$

# ### Discretized equations

# With patience and care, we write the discretized form of the equations. It is highly recommended that you write these in your own hand, mentally following each term as you write it.

# The $u$-momentum equation:

# $$
# \begin{split}
# & \frac{u_{i,j}^{n+1}-u_{i,j}^{n}}{\Delta t}+u_{i,j}^{n}\frac{u_{i,j}^{n}-u_{i-1,j}^{n}}{\Delta x}+v_{i,j}^{n}\frac{u_{i,j}^{n}-u_{i,j-1}^{n}}{\Delta y} = \\
# & \qquad -\frac{1}{\rho}\frac{p_{i+1,j}^{n}-p_{i-1,j}^{n}}{2\Delta x} \\
# & \qquad +\nu\left(\frac{u_{i+1,j}^{n}-2u_{i,j}^{n}+u_{i-1,j}^{n}}{\Delta x^2}+\frac{u_{i,j+1}^{n}-2u_{i,j}^{n}+u_{i,j-1}^{n}}{\Delta y^2}\right)+F_{i,j}
# \end{split}
# $$

# The $v$-momentum equation:

# $$
# \begin{split}
# & \frac{v_{i,j}^{n+1}-v_{i,j}^{n}}{\Delta t}+u_{i,j}^{n}\frac{v_{i,j}^{n}-v_{i-1,j}^{n}}{\Delta x}+v_{i,j}^{n}\frac{v_{i,j}^{n}-v_{i,j-1}^{n}}{\Delta y} = \\
# & \qquad -\frac{1}{\rho}\frac{p_{i,j+1}^{n}-p_{i,j-1}^{n}}{2\Delta y} \\
# & \qquad +\nu\left(\frac{v_{i+1,j}^{n}-2v_{i,j}^{n}+v_{i-1,j}^{n}}{\Delta x^2}+\frac{v_{i,j+1}^{n}-2v_{i,j}^{n}+v_{i,j-1}^{n}}{\Delta y^2}\right)
# \end{split}
# $$

# And the pressure equation:

# $$
# \begin{split}
# & \frac{p_{i+1,j}^{n}-2p_{i,j}^{n}+p_{i-1,j}^{n}}{\Delta x^2} + \frac{p_{i,j+1}^{n}-2p_{i,j}^{n}+p_{i,j-1}^{n}}{\Delta y^2} = \\
# & \qquad \rho\left[\frac{1}{\Delta t}\left(\frac{u_{i+1,j}-u_{i-1,j}}{2\Delta x}+\frac{v_{i,j+1}-v_{i,j-1}}{2\Delta y}\right) - \frac{u_{i+1,j}-u_{i-1,j}}{2\Delta x}\frac{u_{i+1,j}-u_{i-1,j}}{2\Delta x} - 2\frac{u_{i,j+1}-u_{i,j-1}}{2\Delta y}\frac{v_{i+1,j}-v_{i-1,j}}{2\Delta x} - \frac{v_{i,j+1}-v_{i,j-1}}{2\Delta y}\frac{v_{i,j+1}-v_{i,j-1}}{2\Delta y}\right]
# \end{split}
# $$

# As always, we need to re-arrange these equations to the form we need in the code to make the iterations proceed.

# For the $u$- and $v$ momentum equations, we isolate the velocity at time step `n+1`:

# $$
# \begin{split}
# u_{i,j}^{n+1} = u_{i,j}^{n} & - u_{i,j}^{n} \frac{\Delta t}{\Delta x} \left(u_{i,j}^{n}-u_{i-1,j}^{n}\right) - v_{i,j}^{n} \frac{\Delta t}{\Delta y} \left(u_{i,j}^{n}-u_{i,j-1}^{n}\right) \\
# & - \frac{\Delta t}{\rho 2\Delta x} \left(p_{i+1,j}^{n}-p_{i-1,j}^{n}\right) \\
# & + \nu\left[\frac{\Delta t}{\Delta x^2} \left(u_{i+1,j}^{n}-2u_{i,j}^{n}+u_{i-1,j}^{n}\right) + \frac{\Delta t}{\Delta y^2} \left(u_{i,j+1}^{n}-2u_{i,j}^{n}+u_{i,j-1}^{n}\right)\right] \\
# & + \Delta t F
# \end{split}
# $$

# $$
# \begin{split}
# v_{i,j}^{n+1} = v_{i,j}^{n} & - u_{i,j}^{n} \frac{\Delta t}{\Delta x} \left(v_{i,j}^{n}-v_{i-1,j}^{n}\right) - v_{i,j}^{n} \frac{\Delta t}{\Delta y} \left(v_{i,j}^{n}-v_{i,j-1}^{n}\right) \\
# & - \frac{\Delta t}{\rho 2\Delta y} \left(p_{i,j+1}^{n}-p_{i,j-1}^{n}\right) \\
# & + \nu\left[\frac{\Delta t}{\Delta x^2} \left(v_{i+1,j}^{n}-2v_{i,j}^{n}+v_{i-1,j}^{n}\right) + \frac{\Delta t}{\Delta y^2} \left(v_{i,j+1}^{n}-2v_{i,j}^{n}+v_{i,j-1}^{n}\right)\right]
# \end{split}
# $$

# And for the pressure equation, we isolate the term $p_{i,j}^n$ to iterate in pseudo-time:

# $$
# \begin{split}
# p_{i,j}^{n} = & \frac{\left(p_{i+1,j}^{n}+p_{i-1,j}^{n}\right) \Delta y^2 + \left(p_{i,j+1}^{n}+p_{i,j-1}^{n}\right) \Delta x^2}{2(\Delta x^2+\Delta y^2)} \\
# & -\frac{\rho\Delta x^2\Delta y^2}{2\left(\Delta x^2+\Delta y^2\right)} \\
# & \times \left[\frac{1}{\Delta t} \left(\frac{u_{i+1,j}-u_{i-1,j}}{2\Delta x} + \frac{v_{i,j+1}-v_{i,j-1}}{2\Delta y}\right) - \frac{u_{i+1,j}-u_{i-1,j}}{2\Delta x}\frac{u_{i+1,j}-u_{i-1,j}}{2\Delta x} - 2\frac{u_{i,j+1}-u_{i,j-1}}{2\Delta y}\frac{v_{i+1,j}-v_{i-1,j}}{2\Delta x} - \frac{v_{i,j+1}-v_{i,j-1}}{2\Delta y}\frac{v_{i,j+1}-v_{i,j-1}}{2\Delta y}\right]
# \end{split}
# $$

# The initial condition is $u, v, p=0$ everywhere, and at the boundary conditions are:

# $u, v, p$ are periodic on $x=0,2$

# $u, v =0$ at $y =0,2$

# $\frac{\partial p}{\partial y}=0$ at $y =0,2$

# $F=1$ everywhere.

# Let's begin by importing our usual run of libraries:

using Plots

# In step 11, we isolated a portion of our transposed equation to make it easier to parse and we're going to do the same thing here.
# One thing to note is that we have periodic boundary conditions throughout this grid,
# so we need to explicitly calculate the values at the leading and trailing edge of our `u` vector.

function build_up_b(rho, dt, dx, dy, u, v)
    b = zeros(u)
    b[2:end, 1:-1] = (rho * (1 / dt * ((u[2:end, 2:] - u[2:end, 0:-2]) / (2 * dx) +
                                      (v[2:, 1:-1] - v[0:-2, 1:-1]) / (2 * dy)) -
                            ((u[2:end, 2:] - u[2:end, 0:-2]) / (2 * dx))^2 -
                            2 * ((u[2:, 1:-1] - u[0:-2, 1:-1]) / (2 * dy) *
                                 (v[2:end, 2:] - v[2:end, 0:-2]) / (2 * dx))-
                            ((v[2:, 1:-1] - v[0:-2, 1:-1]) / (2 * dy))^2))

    # Periodic BC Pressure @ x = 2
    b[2:end, -1] = (rho * (1 / dt * ((u[2:end, 0] - u[2:end,-2]) / (2 * dx) +
                                    (v[2:, -1] - v[0:-2, -1]) / (2 * dy)) -
                          ((u[2:end, 0] - u[2:end, -2]) / (2 * dx))^2 -
                          2 * ((u[2:, -1] - u[0:-2, -1]) / (2 * dy) *
                               (v[2:end, 0] - v[2:end, -2]) / (2 * dx)) -
                          ((v[2:, -1] - v[0:-2, -1]) / (2 * dy))^2))

    # Periodic BC Pressure @ x = 0
    b[2:end, 0] = (rho * (1 / dt * ((u[2:end, 1] - u[2:end, -1]) / (2 * dx) +
                                   (v[2:, 0] - v[0:-2, 0]) / (2 * dy)) -
                         ((u[2:end, 1] - u[2:end, -1]) / (2 * dx))^2 -
                         2 * ((u[2:, 0] - u[0:-2, 0]) / (2 * dy) *
                              (v[2:end, 1] - v[2:end, -1]) / (2 * dx))-
                         ((v[2:, 0] - v[0:-2, 0]) / (2 * dy))^2))

    return b
end

# We'll also define a Pressure Poisson iterative function, again like we did in Step 11.
# Once more, note that we have to include the periodic boundary conditions at the leading and trailing edge.
# We also have to specify the boundary conditions at the top and bottom of our grid.

function pressure_poisson_periodic(p, dx, dy)
    pn = empty(p)

    for q in range(nit):
        pn = p.copy()
        p[2:end, 1:-1] = (((pn[2:end, 2:] + pn[2:end, 0:-2]) * dy^2 +
                          (pn[2:, 1:-1] + pn[0:-2, 1:-1]) * dx^2) /
                         (2 * (dx^2 + dy^2)) -
                         dx^2 * dy^2 / (2 * (dx^2 + dy^2)) * b[2:end, 1:-1])

        # Periodic BC Pressure @ x = 2
        p[2:end, -1] = (((pn[2:end, 0] + pn[2:end, -2])* dy^2 +
                        (pn[2:, -1] + pn[0:-2, -1]) * dx^2) /
                       (2 * (dx^2 + dy^2)) -
                       dx^2 * dy^2 / (2 * (dx^2 + dy^2)) * b[2:end, -1])

        # Periodic BC Pressure @ x = 0
        p[2:end, 0] = (((pn[2:end, 1] + pn[2:end, -1])* dy^2 +
                       (pn[2:, 0] + pn[0:-2, 0]) * dx^2) /
                      (2 * (dx^2 + dy^2)) -
                      dx^2 * dy^2 / (2 * (dx^2 + dy^2)) * b[2:end, 0])

        # Wall boundary conditions, pressure
        p[-1, :] =p[-2, :]  # dp/dy = 0 at y = 2
        p[0, :] = p[1, :]  # dp/dy = 0 at y = 0
        end

    return p
end

Now we have our familiar list of variables and initial conditions to declare before we start.

# Variable declarations
nx = 41
ny = 41
nt = 10
nit = 50
c = 1
dx = 2 / (nx - 1)
dy = 2 / (ny - 1)
x = LinRange(0, 2, nx)
y = LinRange(0, 2, ny)
X, Y = numpy.meshgrid(x, y)

##physical variables
rho = 1
nu = .1
F = 1
dt = .01

#initial conditions
u = zeros((ny, nx))
un = zeros((ny, nx))

v = zeros((ny, nx))
vn = zeros((ny, nx))

p = ones((ny, nx))
pn = ones((ny, nx))

b = zeros((ny, nx))

# For the meat of our computation, we're going to reach back to a trick we used in Step 9 for Laplace's Equation.
# We're interested in what our grid will look like once we've reached a near-steady state.
# We can either specify a number of timesteps `nt` and increment it until we're satisfied with the results,
# or we can tell our code to run until the difference between two consecutive iterations is very small.

# We also have to manage **8** separate boundary conditions for each iteration.
# The code below writes each of them out explicitly.  If you're interested in a challenge,
# you can try to write a function which can handle some or all of these boundary conditions.
# TODO Fix data structure recommendation
# If you're interested in tackling that, you should probably read up on Julia [dictionaries](http://docs.python.org/2/tutorial/datastructures.html#dictionaries).

udiff = 1
stepcount = 0

while udiff > .001
    un = copy(u)
    vn = copy(v)

    b = build_up_b(rho, dt, dx, dy, u, v)
    p = pressure_poisson_periodic(p, dx, dy)

    u[2:end, 2:end] = (un[2:end, 2:end] -
                     un[2:end, 2:end] * dt / dx *
                    (un[2:end, 2:end] - un[2:end, begin:end-1]) -
                     vn[2:end, 2:end] * dt / dy *
                    (un[2:end, 2:end] - un[begin:end-1, 2:end]) -
                     dt / (2 * rho * dx) *
                    (p[2:end, 3:end] - p[2:end, begin:end-1]) +
                     nu * (dt / dx^2 *
                    (un[2:end, 3:end] - 2 * un[2:end, 2:end] + un[2:end, begin:end-1]) +
                     dt / dy^2 *
                    (un[3:end, 2:end] - 2 * un[2:end, 2:end] + un[begin:end-1, 2:end])) +
                     F * dt)

    v[2:end, 2:end] = (vn[2:end, 2:end] -
                     un[2:end, 2:end] * dt / dx *
                    (vn[2:end, 2:end] - vn[2:end, begin:end-1]) -
                     vn[2:end, 2:end] * dt / dy *
                    (vn[2:end, 2:end] - vn[begin:end-1, 2:end]) -
                     dt / (2 * rho * dy) *
                    (p[3:end, 2:end] - p[begin:end-1, 2:end]) +
                     nu * (dt / dx^2 *
                    (vn[2:end, 3:end] - 2 * vn[2:end, 2:end] + vn[2:end, begin:end-1]) +
                     dt / dy^2 *
                    (vn[3:end, 2:end] - 2 * vn[2:end, 2:end] + vn[begin:end-1, 2:end])))

    # Periodic BC u @ x = 2
    u[2:end, end] = (un[2:end, end] - un[2:end, end] * dt / dx *
                  (un[2:end, end] - un[2:end, end-1]) -
                   vn[2:end, end] * dt / dy *
                  (un[2:end, end] - un[begin:end-1, end]) -
                   dt / (2 * rho * dx) *
                  (p[2:end, begin] - p[2:end, end-1]) +
                   nu * (dt / dx^2 *
                  (un[2:end, begin] - 2 * un[2:end,-1] + un[2:end, end-1]) +
                   dt / dy^2 *
                  (un[3:end, end] - 2 * un[2:end, end] + un[begin:end-1, end])) + F * dt)

    # Periodic BC u @ x = 0
    u[2:end, begin] = (un[2:end, begin] - un[2:end, begin] * dt / dx *
                 (un[2:end, begin] - un[2:end, end]) -
                  vn[2:end, begin] * dt / dy *
                 (un[2:end, begin] - un[begin:end-1, begin]) -
                  dt / (2 * rho * dx) *
                 (p[2:end, 2] - p[2:end, end]) +
                  nu * (dt / dx^2 *
                 (un[2:end, 2] - 2 * un[2:end, begin] + un[2:end, end]) +
                  dt / dy^2 *
                 (un[3:end, begin] - 2 * un[2:end, begin] + un[begin:end-1, begin])) + F * dt)

    # Periodic BC v @ x = 2
    v[2:end, end] = (vn[2:end, end] - un[2:end, end] * dt / dx *
                  (vn[2:end, end] - vn[2:end, end-1]) -
                   vn[2:end, end] * dt / dy *
                  (vn[2:end, end] - vn[begin:end-1, end]) -
                   dt / (2 * rho * dy) *
                  (p[3:end, end] - p[begin:end-1, end]) +
                   nu * (dt / dx^2 *
                  (vn[2:end, begin] - 2 * vn[2:end, end] + vn[2:end, end-1]) +
                   dt / dy^2 *
                  (vn[3:end, end] - 2 * vn[2:end, end] + vn[begin:end-1, end])))

    # Periodic BC v @ x = 0
    v[2:end, begin] = (vn[2:end, begin] - un[2:end, begin] * dt / dx *
                 (vn[2:end, begin] - vn[2:end, end]) -
                  vn[2:end, begin] * dt / dy *
                 (vn[2:end, begin] - vn[begin:end-1, begin]) -
                  dt / (2 * rho * dy) *
                 (p[3:end, begin] - p[begin:end-1, begin]) +
                  nu * (dt / dx^2 *
                 (vn[2:end, 2] - 2 * vn[2:end, begin] + vn[2:end, end]) +
                  dt / dy^2 *
                 (vn[3:end, begin] - 2 * vn[2:end, begin] + vn[begin:end-1, begin])))

    # Wall BC: u,v = 0 @ y = 0,2
    u[begin, :] .= 0
    u[end,   :] .= 0
    v[begin, :] .= 0
    v[end,   :] .= 0

    udiff = (sum(u) - sum(un)) / sum(u)
    stepcount += 1
end

# You can see that we've also included a variable `stepcount` to see how many iterations our loop went through before our stop condition was met.

print(stepcount)

# If you want to see how the number of iterations increases as our `udiff` condition gets smaller and smaller,
# try defining a function to perform the `while` loop written above that takes an input
# `udiff` and outputs the number of iterations that the function runs.

# For now, let's look at our results.  We've used the quiver function
# # to look at the cavity flow results and it works well for channel flow, too.

#TODO PLOT
fig = pyplot.figure(figsize = (11,7), dpi=100)
pyplot.quiver(X[::3, ::3], Y[::3, ::3], u[::3, ::3], v[::3, ::3]);

# The structures in the `quiver` command that look like `[::3, ::3]` are useful when dealing with large amounts of data that you want to visualize.  The one used above tells `matplotlib` to only plot every 3rd data point.  If we leave it out, you can see that the results can appear a little crowded.

# TODO PLOT
fig = pyplot.figure(figsize = (11,7), dpi=100)
pyplot.quiver(X, Y, u, v);

## Learn more
# ***

##### What is the meaning of the $F$ term?

# Step 12 is an exercise demonstrating the problem of flow in a channel or pipe.
# If you recall from your fluid mechanics class, a specified
# pressure gradient is what drives Poisseulle flow.

# Recall the $x$-momentum equation:

# $$\frac{\partial u}{\partial t}+u \cdot \nabla u = -\frac{\partial p}{\partial x}+\nu \nabla^2 u$$

# What we actually do in Step 12 is split the pressure into steady and unsteady components $p=P+p'$. The applied steady pressure gradient is the constant $-\frac{\partial P}{\partial x}=F$ (interpreted as a source term), and the unsteady component is $\frac{\partial p'}{\partial x}$. So the pressure that we solve for in Step 12 is actually $p'$, which for a steady flow is in fact equal to zero everywhere.

# Why did we do this?

# Note that we use periodic boundary conditions for this flow. For a flow with a constant pressure gradient, the value of pressure on the left edge of the domain must be different from the pressure at the right edge. So we cannot apply periodic boundary conditions on the pressure directly. It is easier to fix the gradient and then solve for the perturbations in pressure.

# Shouldn't we always expect a uniform/constant $p'$ then?

# That's true only in the case of steady laminar flows. At high Reynolds numbers, flows in channels can become turbulent, and we will see unsteady fluctuations in the pressure, which will result in non-zero values for $p'$.

# In step 12, note that the pressure field itself is not constant, but it's the pressure perturbation field that is. The pressure field varies linearly along the channel with slope equal to the pressure gradient. Also, for incompressible flows, the absolute value of the pressure is inconsequential.

# ##### And explore more CFD materials online

# The interactive module **12 steps to Navier–Stokes** is one of several components
# of the Computational Fluid Dynamics class taught by Prof. Lorena A. Barba in Boston University between 2009 and 2013.

# For a sample of what the othe components of this class are, you can explore
# the **Resources** section of the Spring 2013 version of [the course's Piazza site](https://piazza.com/bu/spring2013/me702/resources).

# ***

ext: jl

max_stars_repo_name: miguelraz/CFDPython

max_issues_repo_name: miguelraz/CFDJulia

max_forks_repo_name: miguelraz/CFDJulia


text:
# tests sub-generators

# generator and sub-generators

@generator SGBoolGen begin # prefix SG (for Sub-Generators) to avoid type name clashes
    start() = choose(Bool)
end

@generator SGIntGen begin
    start() = choose(Int,5,9)
end

@generator SGMain1SGGen(intGen) begin
    start() = choose(intGen)
end

@generator SGMain2SGGen(boolGen, intGen) begin
    start() = map(i->choose(boolGen), 1:choose(intGen))
end

@generator SGNoParamFormGen(intGen) begin
    start() = choose(intGen)
end

@generator SGRepsGen(intGen) begin
    start() = mult(choose(intGen))
end

@testset "subgenerators" begin

	# important to have test for only one sub-generator in order to ensure correct handling of constructors
	@testset "generator with one sub-generator" begin

		ign = SGIntGen()
		gn = SGMain1SGGen(ign)

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Int
		    @test 5 <= td <= 9
		    @mtest_values_vary td
		end

	end

	@testset "generator with two sub-generators" begin

		bgn = SGBoolGen()
		ign = SGIntGen()
		gn = SGMain2SGGen(bgn, ign)

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Vector{Bool}
		    @test 5 <= length(td) <= 9
		    @mtest_that_sometimes minimum(td) != maximum(td) # different Bool values
		    @mtest_values_vary length(td) # different lengths
		end

	end

	@testset "sub-generator called using no-params short form" begin

		ign = SGIntGen()
		gn = SGNoParamFormGen(ign)

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) == Int
		    @test 5 <= td <= 9
		    @mtest_values_vary td
		end

	end

	@testset "sub-generator as a sequence choice point" begin

		ign = SGIntGen()
		gn = SGRepsGen(ign)

		@mtestset reps=Main.REPS alpha=Main.ALPHA begin
		    td = choose(gn)
		    @test typeof(td) <: Array
		    @test all(5 .<= td .<= 9)
		    @mtest_values_vary length(td)
		end

	end

end


text:
for s in [:(+), :(*), :(-), :(/), :(>), :(<), :(>=), :(<=)]
    expr = quote
        @abstract DiffInterpreter Base.$s(args...) = propagate(args...)
    end
    Base.eval(@__MODULE__, expr)
end

for s in [sin, exp, cos, log]
    expr = quote
        @abstract DiffInterpreter $s(args...) = propagate(args...)
    end
    Base.eval(@__MODULE__, expr)
end


text:
using StochasticDiffEq, DiffEqDevTools, ParameterizedFunctions
using Plots; gr()
const N = 100

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

p = [1.5,1.0,3.0,1.0]

function g(du,u,p,t)
  du .= 0.1u
end
u0 = [1.0;1.0]
tspan = (0.0,10.0)
prob = SDEProblem(f,g,u0,tspan,p);

sol = solve(prob,SRIW1(),abstol=1e-4,reltol=1e-4)
plot(sol)

reltols = 1.0 ./ 4.0 .^ (2:4)
abstols = reltols#[0.0 for i in eachindex(reltols)]
setups = [Dict(:alg=>SRIW1())
          Dict(:alg=>EM(),:dts=>1.0./12.0.^((1:length(reltols)) .+ 1.5))
          Dict(:alg=>RKMil(),:dts=>1.0./12.0.^((1:length(reltols)) .+ 1.5),:adaptive=>false)
          Dict(:alg=>SRIW1(),:dts=>1.0./4.0.^((1:length(reltols)) .+ 5),:adaptive=>false)
          Dict(:alg=>SRIW2())
          Dict(:alg=>SOSRI())
          Dict(:alg=>SOSRI2())
          ]
test_dt = 1/10^2
appxsol_setup = Dict(:alg=>SRIW1(),:abstol=>1e-4,:reltol=>1e-4)
wp = WorkPrecisionSet(prob,abstols,reltols,setups,test_dt;
                                     maxiters = 1e7,
                                     verbose=false,save_everystep=false,
                                     parallel_type = :threads,
                                     appxsol_setup = appxsol_setup,
                                     numruns_error=N,error_estimate=:final)
plot(wp)

reltols = 1.0 ./ 4.0 .^ (2:4)
abstols = reltols#[0.0 for i in eachindex(reltols)]
setups = [Dict(:alg=>SRIW1())
          Dict(:alg=>EM(),:dts=>1.0./12.0.^((1:length(reltols)) .+ 1.5))
          Dict(:alg=>RKMil(),:dts=>1.0./12.0.^((1:length(reltols)) .+ 1.5),:adaptive=>false)
          Dict(:alg=>SRIW1(),:dts=>1.0./4.0.^((1:length(reltols)) .+ 5),:adaptive=>false)
          Dict(:alg=>SRIW2())
          Dict(:alg=>SOSRI())
          Dict(:alg=>SOSRI2())
          ]
test_dt = 1e-2
appxsol_setup = Dict(:alg=>SRIW1(),:abstol=>1e-4,:reltol=>1e-4)
wp = WorkPrecisionSet(prob,abstols,reltols,setups,test_dt;
                                     maxiters = 1e7,
                                     verbose=false,save_everystep=false,
                                     parallel_type = :none,
                                     appxsol_setup = appxsol_setup,
                                     numruns_error=N,error_estimate=:weak_final)
plot(wp;legend=:topleft)

sample_size = Int[10;1e2;1e3;1e4]
se = get_sample_errors(prob,setups[6],test_dt,numruns=sample_size,
                                      sample_error_runs = 100_000,solution_runs=100)

plot(wp;legend=:topleft)
times = [wp[i].times for i in 1:length(wp)]
times = [minimum(minimum(t) for t in times),maximum(maximum(t) for t in times)]
plot!([se[end];se[end]],times,color=:orange,linestyle=:dash,label="Sample Error: 1000",lw=3)

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


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

using Markdown
using InteractiveUtils

# ╔═╡ ba27bfd4-3174-40e2-a7f1-59ee10524890
using SparseArrays, Statistics, ForwardDiff, MAT, Optim, LinearAlgebra

# ╔═╡ 6bbe82e8-bc33-4a55-9a0f-c8b2362447f1
cd("/Users/isaacgeng/OneDrive - The Chinese University of Hong Kong/ECON5480-IO")

# ╔═╡ cfee6ce0-7a6a-4e8b-ba1e-e0f137a84a72
md"""
## 1. Data preparation
1.1 Read the input matrice.\
1.2 Calculate the miscelleous matrices for later use.
"""

# ╔═╡ 56d9b941-ee04-4003-9a19-5e1bd62fd1f2
begin
	# load the files into matrixes, use meta-programming to reduce later
	ps2raw = matopen("data/ps2.mat")
	ivraw = matopen("data/iv.mat")
	x2 = read(ps2raw,  "x2")
	id = read(ps2raw, "id")
	s_jt = read(ps2raw, "s_jt")
	x1 = read(ps2raw, "x1")
	v = read(ps2raw,"v")
	demogr = read(ps2raw, "demogr")
	id_demo = read(ps2raw, "id_demo")
	iv = read(ivraw, "iv")
	close(ps2raw)
	close(ivraw)
end

# ╔═╡ fdce56e0-29ee-475e-a10e-22c452e3c64d
ns, nmkt, nbrn, n_inst = 20, 94, 24, 20

# ╔═╡ 9ce7689a-1afc-4e66-8bad-3ab1dca77da3
cdid = kron([1:nmkt]', ones(nbrn, 1))

# ╔═╡ 0da243dc-0f18-4e96-b8d6-b5d33a9fcc65
cdindex = 24:24:2256

# ╔═╡ 92f1eb9d-53b9-4910-86ff-ff2329e87031
IV = [iv[:,2:n_inst+1] x1[:, 2:nbrn+1]]

# ╔═╡ adbd8332-f16a-44ab-956b-45364ae2ff16
md"""
Starting values: 0 means corresponding coef not maxed over.
"""

# ╔═╡ bbdb3cd8-6023-444c-ab75-5789142aa626
θ_2w = [0.3772    3.0888         0    1.1859         0;
             1.8480   16.5980    -.6590         0   11.6245;
            -0.0035   -0.1925         0    0.0296         0;
             0.0810    1.4684         0   -1.5143         0]

# ╔═╡ 0eb95563-5a29-4ebb-a2fa-5bd70e37539e
ind = findall(!iszero,θ_2w)

# ╔═╡ 8e388be0-5958-4486-a689-b1c1c2b7883b
θ_2_init = filter(!iszero,θ_2w)

# ╔═╡ 6514513b-7bc1-42b0-b5b5-2f89db5a53f1
θ_i = getindex.(ind,1)

# ╔═╡ 3ca1524c-20bd-47c1-9578-3a943ad5d307
θ_j = getindex.(ind,2)

# ╔═╡ d40d1b80-8016-4be0-ab98-257c7b377b2a
md"""
create a weight matrix
"""

# ╔═╡ 38df8f14-9609-46c1-adbf-dc037ed2a081
IV'*IV

# ╔═╡ abea1050-a483-42f3-ba46-f37863d6a056
invA = inv(Matrix(IV'*IV))

# ╔═╡ 7ceb006b-9d2b-4392-8b29-68663ac9c3fc
md"""
## 2. Logit results and save the mean utility as initial values for the search below
"""

# ╔═╡ 1f14383f-1c32-4044-9545-abc473f0c80f
# compute the outside good market share by market
temp = cumsum(s_jt, dims=1)

# ╔═╡ 6bbc13da-2b84-4527-b89c-50d418f8bea0
sum1 = temp[cdindex,:] # total market share in each market besides outside option

# ╔═╡ c0625ee4-f0f4-4f28-8ba5-8d91d3bc04dd
sum1[2:size(sum1,1),:] = diff(sum1,dims=1)

# ╔═╡ 9e569165-a847-46fc-8004-da89d2858bb3
outshr = 1.0 .- repeat(sum1,inner=(24,1))

# ╔═╡ 64acc488-00e6-49cc-be70-55c7308d5627
y = log.(s_jt) - log.(outshr)

# ╔═╡ e97a0bb5-fb93-483b-931f-43e6e34fdd46
mid = x1'*IV*invA*IV'

# ╔═╡ 18f25b8c-fca6-45e2-9ba7-2179c2c7f313
t = (mid*x1) \ (mid*y)
# instead of using inv(mid*x1)*(mid*y), use this to utilize qr factorization.

# ╔═╡ bfa9baf8-a0cd-4c7c-a6b3-69c4d4f01c28
mvalold = exp.(x1*t) # fitted log shares

# ╔═╡ 9c93ef5a-7714-435c-a5ad-4fcf7fec049d
oldt2 = zeros(size(θ_2_init))

# ╔═╡ 817414ea-e30d-441e-8e0c-55dafcde1da1
vfull = repeat(v,inner=(24,1))

# ╔═╡ bf3d1cce-cd9c-465a-b3c2-37edec84da3b
dfull = repeat(demogr,inner=(24,1))

# ╔═╡ 2888cb4c-e29f-4b84-8858-98165d8fe531
md"""
## 3. The Random Coefficient Logit Demand Estimation
### 3.1 Function 1: individual choice probabilities
"""

# ╔═╡ 123da232-f975-40b5-8b27-39f33ae2cb0d
md"""
- original ind_sh matlab code
```
function f = ind_sh(expmval,expmu)
# This function computes the "individual" probabilities of choosing each brand
global ns cdindex cdid
eg = expmu.*kron(ones(1,ns),expmval);
temp = cumsum(eg);
sum1 = temp(cdindex,:);
sum1(2:size(sum1,1),:) = diff(sum1);
denom1 = 1./(1+sum1);
denom = denom1(cdid,:);
f = eg.*denom;
```
"""

# ╔═╡ 417375fb-f62d-4415-a768-d4a3ecf042cf
md"""
- We first show the math of market share,  note that the integrad is individual choice probabilities.
$s_{j} = \int \frac{e^{\delta_j + x_j σ ν}}{1+\sum_{i=1}^J e^{\delta_i + x_i σ ν}} dF_ν(ν)$
- where \
$\delta_j = β[1]*p[j] + x[:,j]' *β[2:end] + ξ[j]$
- here, we choose to replicate the nevo's way, that is, we donot write out the formula $\delta_j$ in ind_sh function.
"""

# ╔═╡ 40d6d064-c96e-446e-8ec6-eff5be1bc4d7
function ind_sh(expmval, expmu)
	eg = expmu.*kron(ones(1,ns), expmval)
	# common techinique in nevo: begin
	temp = cumsum(eg,dims=1);
	sum1 = temp[cdindex,:]
	sum1[2:size(sum1, 1),:] = diff(sum1, dims=1)
	# end
	denom1 = 1 ./ (1 .+ sum1)
	# denom. = denom1[cdid,:]
	denom = repeat(denom1, inner=(24,1))
	f = eg.*denom
end

# ╔═╡ dafdd2b7-5992-420d-91e3-a0e64009db43
md"""
- the ind_sh function's second parameter `expmu` is the result of mufunc, we change the name to $\mu$ _func instead.
"""

# ╔═╡ 2d2d119b-bda2-4dc4-8364-0c93f5871bd9
function μ_func(x2,θ_2w)
	n,k = size(x2)
	j = size(θ_2w,2)-1
	μ = Matrix{Real}(undef, n,ns); # 2256, 20
	for i = 1:ns
    	v_i = vfull[:, i:ns:k*ns] # try v_i = vfull[:,3:20:80] if unclear
      	d_i = dfull[:, i:ns:j*ns]
 		μ[:,i] = x2.*v_i*θ_2w[:,1] + x2.*(d_i*θ_2w[:,2:j+1]')*ones(k,1);
	end
	f = μ;
end

# ╔═╡ 644f082d-51bc-4fda-b46b-ae0cfdffe3bb
md"""
- Note that here I use a different sentence to create a empty container where the type is implicitly specified as being subtypes of `Real`. This is needed since later on we use `Autodiff.jl` and `Autodiff.jl` uses dual to calculate its direvatives and Float64 type is not capable of obtaining its dual.
"""

# ╔═╡ 6cb4a7df-4670-4f60-9474-80f143c3671d
md"""
#### 3.1.1 debug and test: ind_sh
-  we test `ind_sh` work or not. And we found those julia syntax difference to matlab:
1. exp() in matlab need to be exp.() in case the matrix is not square
2. the square brackets near `x2.*v_i*θ_2w[:,1]` to remove
3. the square brackets near `d_i*θ_2w[:,2:j+1]'` should be round brackets.
"""

# ╔═╡ 8f3e129d-a68f-433d-99e9-15d7ac2e4490
md"""
- debug console, for record only.
```
exp.(μ_func(x2,θ_2w))
size(x2.*v_i*θ_2w[:,1])
v_i = vfull[:,3:20:80]
```
"""

# ╔═╡ 9a226183-0b7b-448a-986b-722998900181
md"""
- test for $\mu$_func: passed
"""

# ╔═╡ 31a69149-b3df-4f45-a6e4-0209cf0aaccb
μ_func(x2,θ_2w)

# ╔═╡ dcb6c090-f3da-447b-802b-23c8d4881a33
md"""
- test for ind_sh func: passed
"""

# ╔═╡ 8c5e09da-baaf-4b9c-a90d-b70d5130d47d
a, b = mvalold, exp.(μ_func(x2,θ_2w))

# ╔═╡ 6f1d4752-763a-415d-9b54-d12807ebb574
ind_sh(a, b)

# ╔═╡ 4ee9e2b7-ab19-46f3-9817-0b35728bbe19
md"""
### 3.2 Market share function
- original matlab code
```
function f = mktsh(mval, expmu)
% This function computes the market share for each product
% Written by Aviv Nevo, May 1998.
global ns
f = sum((ind_sh(mval,expmu))')/ns;
f = f';
```
"""

# ╔═╡ 5f690e8f-a375-40c3-8a52-6629f04826e5
function mktsh(mval, expmu)
	f = sum(ind_sh(mval, expmu)', dims=1)./ns;
	f = f'
end

# ╔═╡ 8a29a56b-c217-406e-a8f3-3452913b4deb
md"""
#### 3.2.1 test for mktsh: passed
"""

# ╔═╡ 9507da3c-d833-4bb1-a448-eb4d2da76a1c
mktsh(a, b)

# ╔═╡ 16fb264a-ed75-4d57-9cd2-610499653c9b
md"""
### 3.3 Mean utility level function
"""

# ╔═╡ b37444a4-4296-4d8f-81c9-e428d9b2d9a4
function meanval(θ_2, oldt2=oldt2, mvalold=mvalold)
	if maximum(abs.(θ_2 - oldt2)) < 0.01
		tol = 1e-9
		flag = 0
	else
		tol = 1e-6
		flag = 1
	end

	θ_2w = Matrix(sparse(θ_i,θ_j, θ_2))
	expmu = exp.(μ_func(x2, θ_2w))
	norm = 1
	avgnorm = 1

	i = 0

	# contraction mapping to compute mean utility
	while (norm > tol*10^(flag*floor(i/50))) & (avgnorm > 1e-3*tol*10^(flag*floor(i/50)))
		mval = mvalold .* s_jt ./ mktsh(mvalold, expmu);
		t = abs.(mval - mvalold)
		norm = maximum(t)
		avgnorm = mean(t)
		mvalold = mval
		i = i + 1
	end

	println("# of iterations. for delta convergence: "*string(i))

	if flag == 1 & maximum(isnan.(mval)) < 1
		mval = mvalold .* s_jt ./ mktsh(mvalold, expmu)
		mvalold = mval
		oldt2 = θ_2
	end

	log.(mval)
end

# ╔═╡ a2b63874-1949-4bb5-85e1-30cde5795c61
md"""
#### 3.3.1 Test for meanval: passed
"""

# ╔═╡ 6f69fd00-2db3-4b5e-a261-b45d7b04cc8b
meanval(θ_2_init)

# ╔═╡ f4bf30cd-4829-4d18-8cee-844defcbc26a
md"""
#### 3.3.2 (`Q3.3`) Write a routine that saves `avgnorm, norm`, graph and discuss.
"""

# ╔═╡ 7d3f24b3-aedc-4118-927f-87ae54ceefc0
md"""
### 3.4 Jacobian of Mval and Gmmobjg
- I utilize auto-differentiation to save me from the approx calculation process.
"""

# ╔═╡ 06900cab-44da-4d1d-9fab-11ef3d67b3bc
jacob = θ_2 -> ForwardDiff.jacobian(x -> meanval(x), θ_2)

# ╔═╡ 2487fa3d-ce0e-48e0-bada-2698d0c4b881
function jacob!(g,θ_2)
	g = Matrix{Real}(undef, nbrn*nmkt,size(θ_2)[1])
	g = jacob(θ_2)
	return g
end

# ╔═╡ 9019b1a3-8235-4d46-b9ea-6e9fe0a37dd3
md"""
#### 3.4.1 Test for jacob and jacob! function: passed
"""

# ╔═╡ 620ff1f8-b67d-4e8f-93ee-c7a8938e9769
g = Matrix{Real}(undef, 2256,13);

# ╔═╡ 0f1f6bea-091a-460b-8cec-86d577e1c463
jacob(θ_2_init)

# ╔═╡ fe5b09ff-1eea-46c0-8779-b4b4cdfebe7e
jacob!(g, θ_2_init)

# ╔═╡ 03fa5b43-a49f-439c-9f6b-830ddbd750a5
md"""
#### 3.4.2 Comparing the autodiff with matlab result.
"""

# ╔═╡ d3f2e237-2264-46ae-a2d8-59b587c396d2
md"""
- I compared the results of hand written approximated jacobian (**matlab result below**) and found that there is some degree of difference between the audodiff result and the jacobian approximation result.
"""

# ╔═╡ 7508e049-acb9-4761-8c7f-4f2d8383adb2
begin
	jcb_approx_raw = matopen("data/jcb.mat")
	jcb_approx = read(jcb_approx_raw,"jcb")
	close(jcb_approx_raw)
end

# ╔═╡ d68eb579-b818-42f0-b9b4-ad970bfb55e2
jcb_approx

# ╔═╡ b90497b6-0020-487d-9a6e-5047eaf7e8cd
md"""
### 3.5 GMM objective function: gmmobjg
"""

# ╔═╡ ad667085-b7a9-4ec3-8a7e-e2670a3072f5
md"""
- I manually splitted the function into three functions and the code duplcation is hard to avoid due to that `gmmobjg` need to be single valued function to be optimized by using `Optim.jl` package.
"""

# ╔═╡ c3d3c28d-2458-4c91-b346-086fd4f77453
function gmmobjg(θ_2,nargout=0)
	δ = meanval(θ_2)
	# deals with cases were the min algorithm drifts into region where the objective is not defined
	if maximum(isnan.(δ)) == 1
		f = 1e+10
	else
		temp1 = x1' * IV
		temp2 = δ' * IV
		θ_1= (temp1 * invA * temp1') \ (temp1 * invA * temp2')
		gmmresid = δ - x1*θ_1
		temp1 = gmmresid' * IV;
		f1 = temp1 * invA * temp1';
		f = f1
		# donot know what is for
		if nargout > 1
			temp = jacob(mvalold,θ_2)'
			df = 2 * temp * IV * invA * IV' * gmmresid
		end
	end
	# if nargout > 1
	# 	f, df, gmmresid
	# else
	# 	f, gmmresid
	# end
	real(f[1,1])
end

# ╔═╡ 497fc8fb-6c26-4b44-88f3-62df000c05ae
function gmmresid(θ_2, nargout=0)
	δ = meanval(θ_2)
	# deals with cases were the min algorithm drifts into region where the objective is not defined
	if maximum(isnan.(δ)) == 1
		f = 1e+10
	else
		temp1 = x1' * IV
		temp2 = δ' * IV
		θ_1= (temp1 * invA * temp1') \ (temp1 * invA * temp2')
		gmmresid = δ - x1*θ_1
		temp1 = gmmresid' * IV;
		f1 = temp1 * invA * temp1';
		f = f1
		# donot know what is for
		if nargout > 1
			temp = jacob(mvalold,θ_2)'
			df = 2 * temp * IV * invA * IV' * gmmresid
		end
	end
	gmmresid
end

# ╔═╡ 8395e98e-197f-423d-83fc-66a6bcc80afc
function θ_1(θ_2, nargout=0)
	δ = meanval(θ_2)
	# deals with cases were the min algorithm drifts into region where the objective is not defined
	if maximum(isnan.(δ)) == 1
		f = 1e+10
	else
		temp1 = x1' * IV
		temp2 = δ' * IV
		θ_1= (temp1 * invA * temp1') \ (temp1 * invA * temp2')
		gmmresid = δ - x1*θ_1
		temp1 = gmmresid' * IV;
		f1 = temp1 * invA * temp1';
		f = f1
		# donot know what is for
		if nargout > 1
			temp = jacob(mvalold,θ_2)'
			df = 2 * temp * IV * invA * IV' * gmmresid
		end
	end
	θ_1
end

# ╔═╡ d71340f0-6bfa-42d9-b3f8-819ef536eccb
md"""
#### 3.5.1 Test for gmmobjg, gmmresid, $\theta$_1: passed.
"""

# ╔═╡ 1b56809a-656a-41c2-8164-85d532561565
gmmobjg(θ_2_init)

# ╔═╡ b1fb516d-2368-4c85-b774-d12bd12992f2
gmmresid(θ_2_init)

# ╔═╡ 7757a089-7824-4924-b31f-8788d412cc00
ForwardDiff.gradient(gmmobjg,θ_2_init)

# ╔═╡ c3ca972d-77bc-4b44-94cb-f32100d15c8d
md"""
#### 3.5.2 Jacobians wrt gmmobjg
- later I showed that using Jacobian wrt mval or gmmobjg gives you exact same result.
"""

# ╔═╡ 6a7ec604-b739-48dc-a9a9-287cddf3f6c9
function jcb!(g,θ_2)
	g = Matrix{Real}(undef, 1,size(θ_2)[1])
	g = ForwardDiff.gradient(gmmobjg,θ_2)
	g
end

# ╔═╡ fbed38dc-7535-4c32-92d9-f52ef6cb4073
md"""
### 3.6 Variance-covariance matrix of gmmobjg func wrt $\theta_2$
"""

# ╔═╡ 26c81b54-7760-4eed-96c4-f9a9f2f4b532
function var_cov(θ_2,	gmmresid=gmmresid(θ_2))
	N = size(x1,1)
	Z = size(IV,2)
	temp = jacob(θ_2)
	a = [x1 temp]' * IV
	IVres = IV.*(gmmresid*ones(1, Z))
	b = IVres' * IVres
	f = inv(a*invA*a')*a*invA*b*invA*a'*inv(a*invA*a')
end

# ╔═╡ b96160bb-ffb8-43bf-8e7e-2e4dd221b0ee
md"""
#### 3.6.1 Test for var_cov: passed.
"""

# ╔═╡ 9b11c994-4892-4648-b9e1-6fb069a88304
var_cov(θ_2_init)

# ╔═╡ d10bf777-a869-4e99-aa6b-8d9327dfad91
md"""
## 4. Final routine solving for optimal $\theta_2$ to gmmobjg
"""

# ╔═╡ 2040ca28-08ea-4e29-b9a5-f0276122bf78
md"""
- I found that using derivatives on mealval function or on gmmobjg itself has no difference as below shows.
"""

# ╔═╡ d22d94df-1498-439d-8bcf-a5e2379ab70a
Optim.minimizer(optimize(gmmobjg, jacob!,θ_2_init, BFGS()))

# ╔═╡ ca002c77-5a15-4bb7-a726-d02625317ab9
θ2 = Optim.minimizer(optimize(gmmobjg, jcb!,θ_2_init, BFGS()))

# ╔═╡ c8cd2a43-8ece-4e26-a5d4-14630d6572d6
md"""
- extract execution time and final objective function value.
"""

# ╔═╡ 5385600c-5ad4-48eb-beed-a430e5ee788c
ttime = @elapsed res = optimize(gmmobjg, jcb!,θ_2_init, BFGS())

# ╔═╡ 0841ae53-f235-46a0-8b7d-83843898d7c0
fval = Optim.minimum(res)

# ╔═╡ ee1e5433-8dbc-4e62-b667-36ce3a1cb763
begin
	vcov = var_cov(θ2)
	se = real(sqrt(var_cov(θ2)))
end

# ╔═╡ b3328c91-e6f2-4c38-b8a0-75ebe306ca82
begin
	θ2w = Matrix(sparse(θ_i,θ_j,θ2))
	t_new = size(se,1) - size(θ2,1)
	se2w = Matrix(sparse(θ_i, θ_j, se[t_new+1:size(se,1)]))
end

# ╔═╡ 108fa10b-e075-430a-bbfe-b15d1f8d31e3
θ1 = θ_1(θ2)

# ╔═╡ 4e441ea2-1ae9-4052-a238-cf227674614a
begin
	Ω = inv(vcov[2:25,2:25])
	xmd = [x2[1:24,1] x2[1:24,3:4]];
	ymd = θ1[2:25];
	β = (xmd' * Ω * xmd) \ (xmd' * Ω * ymd)
	resmd = ymd - xmd * β
	semd = sqrt.(diag(inv(xmd' * Ω * xmd)))
	mcoef = [β[1]; θ1[1]; β[2:3]];
	semcoef = [semd[1]; se[1]; semd]
end

# ╔═╡ 0b52131c-9488-4175-8572-646719a9a8cb
Rsq = 1 - ((resmd .- mean(resmd))'*(resmd .- mean(resmd)))/ ((ymd .- mean(ymd))'*(ymd .-mean(ymd)))

# ╔═╡ b17c094c-d231-45d6-8bf3-c30d178ff732
Rsq_G = 1-(resmd'*Ω*resmd)/((ymd .- mean(ymd))'*Ω*(ymd .- mean(ymd)))

# ╔═╡ 388ffa99-ec8e-465b-8355-31508de951de
Chisq = size(id,1)*resmd'*Ω*resmd

# ╔═╡ 1a042b8b-2c67-492b-9b86-3b1276a6616f
mcoef1, θ2w1, semcoef1, se2w1 = round.(mcoef,digits=4), round.(θ2w,digits=4), round.(semcoef,digits=4), round.(se2w, digits=4);

# ╔═╡ 473f2f56-4176-4e20-b36a-0d03d8c4c7b5
md"""
#### 4.1 Show the results and comparison to nevo's
"""

# ╔═╡ 3dd0dfeb-928e-45fe-aff3-8c42dae39cd0
s =
"
	       Mean   Sigma  Income Income^2 Age Child \n
Constant: $(vcat(mcoef1[1], θ2w1[1,:]))\n
Constant se: $(vcat(semcoef1[1], se2w1[1,:]))\n
Price: $(vcat(mcoef1[2], θ2w1[2,:]))\n
Price se: $(vcat(semcoef1[1], se2w1[1,:]))\n
Sugar: $(vcat(mcoef1[3], θ2w1[3,:]))\n
Sugar se: $(vcat(semcoef1[3], se2w1[3,:]))\n
Mushy: $(vcat(mcoef1[4], θ2w1[4,:]))\n
Mushy se: $(vcat(semcoef1[4], se2w1[4,:]))\n
GMM Objective: $fval \n
MR R-squared: $Rsq \n
MR Weighted R-squared: $Rsq_G \n
run time: $ttime seconds
";

# ╔═╡ eb8da8fa-1984-4eda-a947-77e4da924c39
Text(s) # show the table of results

# ╔═╡ 4001b7e2-fde2-4d79-ad20-35b6b433b736
md"""
The result is really close to matlab's results wrt estimates.
However, the ses especially the one wrt price is a lot difference.
My estimation of se is generally **smaller** than Nevo's, indicating the benefits of using a forward difference algorithm by `autodiff.jl`.
"""

# ╔═╡ 587f4cbb-d0cb-41ef-b95a-0049bcaa09d4
"
constant
   -1.8866    0.3772    3.0888         0    1.1859         0

   (1,1)       0.2476
   (1,2)       0.1358
   (1,3)       1.1939
   (1,5)       1.0194

price
  -32.4374    1.8480   16.5980   -0.6590         0   11.6245

   (1,1)       7.5809
   (1,2)       1.0913
   (1,3)     173.9858
   (1,4)       9.0578
   (1,6)       5.3098

sugar
    0.1576   -0.0035   -0.1925         0    0.0296         0

   (1,1)       0.2476
   (1,2)       0.0129
   (1,3)       0.0466
   (1,5)       0.0374

mushy
    0.9842    0.0810    1.4684         0   -1.5143         0

   (1,1)       0.0132
   (1,2)       0.2078
   (1,3)       0.6975
   (1,5)       1.1095

";

# ╔═╡ 347c9e65-ba7f-44ce-959e-838466abb44d
md"""
### 5. Conclusion
- This marks the end of the file.
"""

# ╔═╡ Cell order:
# ╠═ba27bfd4-3174-40e2-a7f1-59ee10524890
# ╠═6bbe82e8-bc33-4a55-9a0f-c8b2362447f1
# ╟─cfee6ce0-7a6a-4e8b-ba1e-e0f137a84a72
# ╠═56d9b941-ee04-4003-9a19-5e1bd62fd1f2
# ╠═fdce56e0-29ee-475e-a10e-22c452e3c64d
# ╠═9ce7689a-1afc-4e66-8bad-3ab1dca77da3
# ╠═0da243dc-0f18-4e96-b8d6-b5d33a9fcc65
# ╠═92f1eb9d-53b9-4910-86ff-ff2329e87031
# ╟─adbd8332-f16a-44ab-956b-45364ae2ff16
# ╠═bbdb3cd8-6023-444c-ab75-5789142aa626
# ╠═0eb95563-5a29-4ebb-a2fa-5bd70e37539e
# ╠═8e388be0-5958-4486-a689-b1c1c2b7883b
# ╠═6514513b-7bc1-42b0-b5b5-2f89db5a53f1
# ╠═3ca1524c-20bd-47c1-9578-3a943ad5d307
# ╟─d40d1b80-8016-4be0-ab98-257c7b377b2a
# ╠═38df8f14-9609-46c1-adbf-dc037ed2a081
# ╠═abea1050-a483-42f3-ba46-f37863d6a056
# ╠═7ceb006b-9d2b-4392-8b29-68663ac9c3fc
# ╠═1f14383f-1c32-4044-9545-abc473f0c80f
# ╠═6bbc13da-2b84-4527-b89c-50d418f8bea0
# ╠═c0625ee4-f0f4-4f28-8ba5-8d91d3bc04dd
# ╠═9e569165-a847-46fc-8004-da89d2858bb3
# ╠═64acc488-00e6-49cc-be70-55c7308d5627
# ╠═e97a0bb5-fb93-483b-931f-43e6e34fdd46
# ╠═18f25b8c-fca6-45e2-9ba7-2179c2c7f313
# ╠═bfa9baf8-a0cd-4c7c-a6b3-69c4d4f01c28
# ╠═9c93ef5a-7714-435c-a5ad-4fcf7fec049d
# ╠═817414ea-e30d-441e-8e0c-55dafcde1da1
# ╠═bf3d1cce-cd9c-465a-b3c2-37edec84da3b
# ╟─2888cb4c-e29f-4b84-8858-98165d8fe531
# ╟─123da232-f975-40b5-8b27-39f33ae2cb0d
# ╟─417375fb-f62d-4415-a768-d4a3ecf042cf
# ╠═40d6d064-c96e-446e-8ec6-eff5be1bc4d7
# ╟─dafdd2b7-5992-420d-91e3-a0e64009db43
# ╠═2d2d119b-bda2-4dc4-8364-0c93f5871bd9
# ╟─644f082d-51bc-4fda-b46b-ae0cfdffe3bb
# ╟─6cb4a7df-4670-4f60-9474-80f143c3671d
# ╟─8f3e129d-a68f-433d-99e9-15d7ac2e4490
# ╟─9a226183-0b7b-448a-986b-722998900181
# ╠═31a69149-b3df-4f45-a6e4-0209cf0aaccb
# ╟─dcb6c090-f3da-447b-802b-23c8d4881a33
# ╠═8c5e09da-baaf-4b9c-a90d-b70d5130d47d
# ╠═6f1d4752-763a-415d-9b54-d12807ebb574
# ╟─4ee9e2b7-ab19-46f3-9817-0b35728bbe19
# ╠═5f690e8f-a375-40c3-8a52-6629f04826e5
# ╟─8a29a56b-c217-406e-a8f3-3452913b4deb
# ╠═9507da3c-d833-4bb1-a448-eb4d2da76a1c
# ╟─16fb264a-ed75-4d57-9cd2-610499653c9b
# ╠═b37444a4-4296-4d8f-81c9-e428d9b2d9a4
# ╟─a2b63874-1949-4bb5-85e1-30cde5795c61
# ╠═6f69fd00-2db3-4b5e-a261-b45d7b04cc8b
# ╟─f4bf30cd-4829-4d18-8cee-844defcbc26a
# ╟─7d3f24b3-aedc-4118-927f-87ae54ceefc0
# ╠═06900cab-44da-4d1d-9fab-11ef3d67b3bc
# ╠═2487fa3d-ce0e-48e0-bada-2698d0c4b881
# ╠═9019b1a3-8235-4d46-b9ea-6e9fe0a37dd3
# ╠═620ff1f8-b67d-4e8f-93ee-c7a8938e9769
# ╠═0f1f6bea-091a-460b-8cec-86d577e1c463
# ╠═fe5b09ff-1eea-46c0-8779-b4b4cdfebe7e
# ╟─03fa5b43-a49f-439c-9f6b-830ddbd750a5
# ╟─d3f2e237-2264-46ae-a2d8-59b587c396d2
# ╠═7508e049-acb9-4761-8c7f-4f2d8383adb2
# ╠═d68eb579-b818-42f0-b9b4-ad970bfb55e2
# ╠═b90497b6-0020-487d-9a6e-5047eaf7e8cd
# ╟─ad667085-b7a9-4ec3-8a7e-e2670a3072f5
# ╠═c3d3c28d-2458-4c91-b346-086fd4f77453
# ╠═497fc8fb-6c26-4b44-88f3-62df000c05ae
# ╠═8395e98e-197f-423d-83fc-66a6bcc80afc
# ╟─d71340f0-6bfa-42d9-b3f8-819ef536eccb
# ╠═1b56809a-656a-41c2-8164-85d532561565
# ╠═b1fb516d-2368-4c85-b774-d12bd12992f2
# ╠═7757a089-7824-4924-b31f-8788d412cc00
# ╟─c3ca972d-77bc-4b44-94cb-f32100d15c8d
# ╠═6a7ec604-b739-48dc-a9a9-287cddf3f6c9
# ╠═fbed38dc-7535-4c32-92d9-f52ef6cb4073
# ╠═26c81b54-7760-4eed-96c4-f9a9f2f4b532
# ╟─b96160bb-ffb8-43bf-8e7e-2e4dd221b0ee
# ╠═9b11c994-4892-4648-b9e1-6fb069a88304
# ╟─d10bf777-a869-4e99-aa6b-8d9327dfad91
# ╟─2040ca28-08ea-4e29-b9a5-f0276122bf78
# ╠═d22d94df-1498-439d-8bcf-a5e2379ab70a
# ╠═ca002c77-5a15-4bb7-a726-d02625317ab9
# ╟─c8cd2a43-8ece-4e26-a5d4-14630d6572d6
# ╠═5385600c-5ad4-48eb-beed-a430e5ee788c
# ╠═0841ae53-f235-46a0-8b7d-83843898d7c0
# ╠═ee1e5433-8dbc-4e62-b667-36ce3a1cb763
# ╠═b3328c91-e6f2-4c38-b8a0-75ebe306ca82
# ╠═108fa10b-e075-430a-bbfe-b15d1f8d31e3
# ╠═4e441ea2-1ae9-4052-a238-cf227674614a
# ╠═0b52131c-9488-4175-8572-646719a9a8cb
# ╠═b17c094c-d231-45d6-8bf3-c30d178ff732
# ╠═388ffa99-ec8e-465b-8355-31508de951de
# ╠═1a042b8b-2c67-492b-9b86-3b1276a6616f
# ╟─473f2f56-4176-4e20-b36a-0d03d8c4c7b5
# ╟─3dd0dfeb-928e-45fe-aff3-8c42dae39cd0
# ╠═eb8da8fa-1984-4eda-a947-77e4da924c39
# ╟─4001b7e2-fde2-4d79-ad20-35b6b433b736
# ╠═587f4cbb-d0cb-41ef-b95a-0049bcaa09d4
# ╟─347c9e65-ba7f-44ce-959e-838466abb44d


text:
"""
    baselevel(d, lags)

Computes baselevel activation according exact equation.

# Arguments

- `d`: decay parameter
- `lags`: a vector of time lags between current time and each usage time
"""
function baselevel(d, lags)
    act = 0.0
    for t in lags
        act += t^-d
    end
    return log(act)
end

"""
    baselevel(d, chunk)

Computes baselevel activation according exact equation.

# Arguments

- `actr`: an `ACTR` model object
- `chunk`: a chunk
"""
function baselevel!(actr, chunk)
    @unpack N,L,k,lags = chunk
    d = actr.parms.d
    exact = baselevel(d, lags)
    approx = 0.0
    if N > k
        tk = lags[k]
        x1 = (N - k) * (L^(1 - d) - tk^(1 - d))
        x2 = (1 - d) * (L - tk)
        approx = x1 / x2
    end
    chunk.act_bll = log(exp(exact) + approx)
    return nothing
end

"""
    baselevel!(actr)

Computes baselevel activation for all chunks according exact equation.

# Arguments

- `actr`: an `ACTR` model object
"""
baselevel!(actr) = activation!.(actr, memory.memory)

"""
    compute_activation!(actr::AbstractACTR, chunks::Vector{<:Chunk}; request...)

Computes the activation of a vector of chunks. By default, current time is computed with `get_time`.

# Arguments

- `actr::AbstractACTR`: an `ACTR` object
- `chunks::Vector{<:Chunk}`: a vector of chunks.

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function compute_activation!(actr::AbstractACTR, chunks::Vector{<:Chunk}; request...)
    return compute_activation!(actr, chunks, get_time(actr); request...)
end

"""
    compute_activation!(actr::AbstractACTR, chunks::Vector{<:Chunk}, cur_time::Float64=0.0; request...)

Computes the activation of a vector of chunks

# Arguments

- `actr::AbstractACTR`: an `ACTR` object
- `chunks::Vector{<:Chunk}`: a vector of chunks.
- `cur_time::Float64`: current simulated time in seconds

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function compute_activation!(actr::AbstractACTR, chunks::Vector{<:Chunk}, cur_time::Float64; request...)
    @unpack sa = actr.parms
    sa ? cache_denomitors(actr) : nothing
    # compute activation for each chunk
    for chunk in chunks
        activation!(actr, chunk, cur_time; request...)
    end
    return nothing
end

"""
    compute_activation!(actr, chunk::Chunk; request...)

Computes the activation of a chunk. By default, current time is computed
with `get_time`.

# Arguments

- `actr`: actr object
- `chunk::Chunk`: a chunk.

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function compute_activation!(actr, chunk::Chunk; request...)
    return compute_activation!(actr, chunk, get_time(actr); request...)
end

"""
    compute_activation!(actr, chunk::Chunk, cur_time; request...)

Computes the activation of a chunk

# Arguments

- `actr`: actr object
- `chunk::Chunk`: a chunk.
- `cur_time`: current simulated time in seconds

# Keywords

- `request...`: optional keywords for the retrieval request
"""
compute_activation!(actr, chunk::Chunk, cur_time; request...) = compute_activation!(actr, [chunk], cur_time; request...)

"""
    compute_activation!(actr::AbstractACTR; request...)

Computes the activation of all chunks in declarative memory. By default, current time is computed
with `get_time`.

# Arguments

- `actr::AbstractACTR`: an `ACTR` object

# Keywords

- `request...`: optional keywords for the retrieval request
"""
compute_activation!(actr::AbstractACTR; request...) = compute_activation!(actr, actr.declarative.memory, get_time(actr); request...)

"""
    compute_activation!(actr::AbstractACTR, cur_time::Float64; request...)

Computes the activation of all chunks in declarative memory

# Arguments

- `actr::AbstractACTR`: an `ACTR` object
- `cur_time`: current simulated time in seconds

# Keywords

- `request...`: optional keywords for the retrieval request
"""
compute_activation!(actr::AbstractACTR, cur_time::Float64; request...) = compute_activation!(actr, actr.declarative.memory, cur_time; request...)

"""
    activation!(actr, chunk::Chunk, cur_time; request...)

Computes the activation of a chunk

# Arguments

- `actr`: an `ACTR` object
- `chunk::Chunk`: a chunk.
- `cur_time`: current simulated time in seconds

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function activation!(actr, chunk::Chunk, cur_time=0.0; request...)
    memory = actr.declarative
    @unpack sa_fun,bll,mmp,sa,noise,blc,τ = actr.parms
    reset_activation!(chunk)
    chunk.act_blc = blc + chunk.bl
    if bll
        update_lags!(chunk, cur_time)
        baselevel!(actr, chunk)
    end
    if mmp
        partial_matching!(actr, chunk; request...)
    end
    if sa
        sa_fun(actr, chunk)
    end
    if noise
        add_noise!(actr, chunk)
    end
    total_activation!(chunk)
    return nothing
end

function reset_activation!(chunk)
    a = chunk.act
    chunk.act_mean = zero(a)
    chunk.act_blc = zero(a)
    chunk.act_bll = zero(a)
    chunk.act_pm = zero(a)
    chunk.act_sa = zero(a)
    chunk.act_noise = zero(a)
    chunk.act = zero(a)
end

"""
    total_activation!(chunk)

Assigns sum of all components to of activation to `chunk.act`.

# Arguments

- `chunk`: a chunk object
"""
function total_activation!(chunk)
    chunk.act_mean = chunk.act_blc + chunk.act_bll - chunk.act_pm +
        chunk.act_sa
    chunk.act = chunk.act_mean + chunk.act_noise
    return nothing
end

function add_noise!(actr, chunk)
    @unpack τ,s = actr.parms
    σ = s * pi / sqrt(3)
    chunk.act_noise = rand(Normal(0, σ))
    return nothing
end

function add_noise!(actr)
    @unpack τ,s = actr.parms
    σ = s * pi / sqrt(3)
    actr.parms.τ′ = rand(Normal(τ, σ))
    nothing
end

"""
    partial_matching!(actr, chunk; request...)

Computes activation for partial matching component

# Arguments

- `actr`: an ACTR object
- `chunk`: a chunk

# Keywords

- `request...`: optional keyword arguments for retrieval request
"""
function partial_matching!(actr, chunk; request...)
    p = actr.parms.mmp_fun(actr, chunk; request...)
    chunk.act_pm = p
    return nothing
end

"""
    set_noise!(actr::AbstractACTR, b)

Sets noise true or false.

# Arguments

- `actr`: ACTR object
- `b`: boolean value
"""
function set_noise!(actr, b)
    actr.parms.noise = b
end

"""
    spreading_activation!(actr, chunk)

Computes the spreading activation for a given chunk

# Arguments

- `actr`: an `ACTR` oject
- `chunk`: the chunk for which spreading activation is computed
"""
function spreading_activation!(actr, chunk)
    imaginal = actr.imaginal
    isempty(imaginal.buffer) ? (return nothing) : nothing
    w = compute_weights(imaginal)
    γ = actr.parms.γ; r = zero(γ); sa = zero(γ)
    slots = imaginal.buffer[1].slots
    denoms = imaginal.denoms
    for (v,d) in zip(slots, denoms)
        num = count_values(chunk, v)
        fan = num / (d + 1)
        r = fan == 0 ? zero(γ) : γ + log(fan)
        sa += w * r
    end
    chunk.act_sa = sa
    return nothing
end

"""
    cache_denomitors(actr)

Caches denominator of spreading activation calculations

# Arguments

- `actr`: an `ACTR` oject
"""
function cache_denomitors(actr)
    @unpack imaginal,declarative = actr
    isempty(imaginal.buffer) ? (return nothing) : nothing
    slots = imaginal.buffer[1].slots
    denoms = fill(0, length(slots))
    for (i,v) in enumerate(slots)
        denoms[i] = compute_denom(declarative, v)
    end
    imaginal.denoms = denoms
    return nothing
end

function compute_weights(mod)
    return mod.ω / length(mod.buffer[1].slots)
end

function compute_denom(memory, value)
    denom = 0
    for c in memory.memory
        denom += count_values(c, value)
    end
    return denom
end

function count_values(chunk, value)
    return count(x -> x == value, values(chunk.slots))
end

"""
    update_recent!(actr, chunk)

Adds a new timestamp to chunk and removes oldest timestamp if
length equals k. By default, current time is computed with `get_time`.

# Arguments
- `actr`: an ACT-R model object
- `chunk`: memory chunk object
"""
function update_recent!(actr::AbstractACTR, chunk)
    return update_recent!(chunk, get_time(actr))
end

"""
    update_recent!(chunk, cur_time)

Adds a new timestamp to chunk and removes oldest timestamp if
length equals k.

# Arguments

* `chunk`: memory chunk object
* `cur_time`: current simulated time in seconds
"""
function update_recent!(chunk, cur_time)
    k = chunk.k; recent = chunk.recent
    if length(recent) == k
        pop!(recent)
    end
    pushfirst!(recent, cur_time)
    return nothing
end

"""
    retrieval_prob(actr::AbstractACTR, target::Array{<:Chunk,1}; request...)

Uses the softmax approximation to compute the retrieval probability of retrieving a chunk.
By default, current time is computed from `get_time`.

# Arguments

- `actr::AbstractACTR`: an ACT-R object
- `target::Array{<:Chunk,1}`: a vector chunks in the numerator of the softmax function

# Keywords

- `request...`: optional keyword pairs representing a retrieval request
"""
function retrieval_prob(actr::AbstractACTR, target::Array{<:Chunk,1}; request...)
    return retrieval_prob(actr, target, get_time(actr); request...)
end

"""
    retrieval_prob(actr::AbstractACTR, target::Array{<:Chunk,1}, cur_time; request...)

Computes the retrieval probability of one chunk from a set of chunks defined in `target`.
Retrieval probability is computed with the softmax approximation.

# Arguments

- `actr::AbstractACTR`: an actr object
- `target::Array{<:Chunk,1}`: a vector chunks in the numerator of the softmax function
- `cur_time`: current time in seconds

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function retrieval_prob(actr::AbstractACTR, target::Array{<:Chunk,1}, cur_time; request...)
    @unpack τ,s,noise = actr.parms
    σ = s * sqrt(2)
    chunks = retrieval_request(actr; request...)
    filter!(x -> (x ∈ chunks), target)
    isempty(target) ? (return (0.0,1.0)) : nothing
    set_noise!(actr, false)
    compute_activation!(actr, chunks, cur_time; request...)
    set_noise!(actr, noise)
    denom = fill(target[1].act_mean, length(chunks) + 1)
    map!(x -> exp(x.act_mean / σ), denom, chunks)
    denom[end] = exp(τ / σ)
    num = map(x -> exp(x.act_mean / σ), target)
    prob = sum(num) / sum(denom)
    fail = denom[end] / sum(denom)
    return prob,fail
end

"""
    retrieval_prob(actr::AbstractACTR, chunk::Chunk; request...)

Uses the softmax approximation to compute the retrieval probability of retrieving a chunk.
By default, current time is computed from `get_time`.

# Arguments

- `actr::AbstractACTR`: an ACT-R object
- `chunk::Chunk`: a chunk

# Keywords

- `request...`: optional keyword pairs representing a retrieval request
"""
function retrieval_prob(actr::AbstractACTR, chunk::Chunk; request...)
    return retrieval_prob(actr, chunk, get_time(actr); request...)
end

"""
    retrieval_prob(actr::AbstractACTR, chunk::Chunk, cur_time=0.0; request...)

Uses the softmax approximation to compute the retrieval probability of retrieving a chunk.

# Arguments

- `actr::AbstractACTR`: an ACT-R object
- `chunk::Chunk`: a chunk
- `cur_time`: current simulated time in seconds

# Keywords

- `request...`: optional keyword pairs representing a retrieval request
"""
function retrieval_prob(actr::AbstractACTR, chunk::Chunk, cur_time; request...)
    @unpack τ,s,noise = actr.parms
    σ = s * sqrt(2)
    chunks = retrieval_request(actr; request...)
    !(chunk ∈ chunks) ? (return (0.0,1.0)) : nothing
    set_noise!(actr, false)
    compute_activation!(actr, chunks, cur_time; request...)
    set_noise!(actr, noise)
    v = fill(chunk.act_mean, length(chunks) + 1)
    map!(x -> exp(x.act_mean / σ), v, chunks)
    v[end] = exp(τ / σ)
    prob = exp(chunk.act_mean / σ) / sum(v)
    fail = v[end] / sum(v)
    return prob,fail
end

"""
    retrieval_probs(actr::AbstractACTR; request...)

Computes the retrieval probability for each chunk matching the retrieval request. By default,
current time is computed from `get_time`.

# Arguments

- `actr::AbstractACTR`: an actr object

# Keywords

- `request...`: optional keyword pairs representing a retrieval request
"""

function retrieval_probs(actr::AbstractACTR; request...)
    return retrieval_probs(actr, get_time(actr); request...)
end

"""
    retrieval_probs(actr::AbstractACTR, cur_time; request...)

Computes the retrieval probability for each chunk matching the retrieval request.

# Arguments

- `actr::AbstractACTR`: an actr object
- `cur_time`: current simulated time in seconds

# Keywords

- `request...`: optional keyword pairs representing a retrieval request
"""
function retrieval_probs(actr::AbstractACTR, cur_time; request...)
    @unpack τ,s,noise = actr.parms
    σ = s * sqrt(2)
    chunks = retrieval_request(actr; request...)
    isempty(chunks) ? (return ([0.0],chunks)) : nothing
    set_noise!(actr, false)
    compute_activation!(actr, chunks, cur_time; request...)
    set_noise!(actr, noise)
    v = Array{typeof(chunks[1].act),1}(undef, length(chunks) + 1)
    map!(x -> exp(x.act_mean / σ), v, chunks)
    v[end] = exp(τ / σ)
    p = v ./ sum(v)
    return p,chunks
end

"""
    update_lags!(chunk::Chunk, cur_time)

Compute lags for each use of a chunk.

# Arguments

- `chunk::Chunks`: a chunk
- `cur_time`: current simulated time in seconds.

"""
function update_lags!(chunk::Chunk, cur_time)
    chunk.L = cur_time - chunk.time_created
    chunk.lags = cur_time .- chunk.recent
    return nothing
end

"""
    update_lags!(actr::AbstractACTR, cur_time)

Compute lags for each use of a chunk. Applies to all chunks in declarative memory.

# Arguments

- `actr::AbstractACTR`: an ACTR object
- `cur_time`: current simulated time in seconds
"""
update_lags!(actr::AbstractACTR, cur_time) = update_lags!(actr.declarative, cur_time)

"""
    update_lags!(memory::Declarative, cur_time)

Compute lags for each use of a chunk. Applies to all chunks in declarative memory.

# Arguments

- `memory::Declarative`: a declarative memory object object
- `cur_time`: current simulated time in seconds
"""
update_lags!(memory::Declarative, cur_time) = update_lags!.(memory.memory, cur_time)

"""
    update_chunk!(chunk, cur_time)

Increments number of uses and adds `cur_time` as the most recent time of use.

# Arguments

- `chunk`: a chunk object
- `cur_time': current simulated time in seconds
"""
function update_chunk!(chunk, cur_time)
    update_recent!(chunk, cur_time)
    chunk.N += 1
    return nothing
end

"""
    add_chunk!(actr::AbstractACTR; slots...)

Adds new chunk to declarative memory or updates existing chunk with new use. The default time is
computed from `get_time`.

# Arguments

- `actr::AbstractACTR`: an ACT-R model object

# Keywords

- `slots...`: optional keyword arguments corresponding to slot-value pairs, e.g. name=:Bob
"""
add_chunk!(actr::AbstractACTR; slots...) = add_chunk!(actr.declarative, get_time(actr); slots...)

"""
    add_chunk!(memory::Declarative, cur_time=0.0; act=0.0, slots...)

Adds new chunk to declarative memory or updates existing chunk with new use

# Arguments

- `memory::Declarative`: declarative memory object
- `cur_time=0.0`: current simulated time in seconds
- `bl=0.0`: baselevel constant for new/updated chunk

# Keywords

- `slots...`: optional keyword arguments corresponding to slot-value pairs, e.g. name=:Bob
"""
function add_chunk!(memory::Declarative, cur_time=0.0; bl::T=0.0, slots...) where {T}
    chunk = get_chunks_exact(memory; slots...)
    if isempty(chunk)
        c = Chunk(;act=zero(T), bl, time_created=cur_time, recent=[cur_time], slots...)
        push!(memory.memory, c)
    else
        update_chunk!(chunk[1], cur_time)
    end
    return nothing
end

"""
    add_chunk!(actr::AbstractACTR, cur_time=0.0; slots...)

Adds new chunk to declarative memory or updates existing chunk with new use

# Arguments

- `actr::AbstractACTR`: an ACT-R model object
- `cur_time=0.0`: current simulated time in seconds

# Keywords

- `slots...`: optional keyword arguments corresponding to slot-value pairs, e.g. name=:Bob
"""
add_chunk!(actr::AbstractACTR, cur_time; slots...) = add_chunk!(actr.declarative, cur_time; slots...)

"""
    get_chunks_exact(memory::Vector{<:Chunk}; criteria...)

Returns all chunks that matches a set criteria and has the same number of slots.

# Arguments

- `memory::Vector{<:Chunk}`: vector of chunk objects

# Keywords

-`criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
function get_chunks_exact(memory::Vector{<:Chunk}; criteria...)
    c = filter(x -> match_exact(x, criteria), memory)
    return c
end

get_chunks_exact(d::Declarative; criteria...) = get_chunks_exact(d.memory; criteria...)

function match_exact(chunk::Chunk, request)
    slots = chunk.slots
    length(slots) ≠ length(request) ? (return false) : nothing
    for (k,v) in request
        if !(k ∈ keys(slots)) || (slots[k] != v)
            return false
        end
    end
    return true
end

"""
    get_chunks(memory::Vector{<:Chunk}; criteria...)

Returns all chunks that matches a set criteria.

# Arguments

- `memory::Vector{<:Chunk}`: vector of chunk objects

# Keywords

 - `check_value=true`: check slot value

-`criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
function get_chunks(memory::Vector{<:Chunk}; check_value=true, criteria...)
    c = filter(x -> _match(x, criteria; check_value), memory)
    return c
end

"""
    get_chunks(memory::Vector{<:Chunk}, funs...; criteria...)

Returns all chunks that matches a set `criteria` which are evaluted according to the list of functions in `funs`.

# Arguments

- `memory::Vector{<:Chunk}`: vector of chunk objects
- `funs...`: a list of functions

# Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
function get_chunks(memory::Vector{<:Chunk}, funs...; check_value=true, criteria...)
    c = filter(x -> _match(x, funs, criteria; check_value), memory)
    return c
end

"""
    get_chunks(d::Declarative; criteria...)

Returns all chunks that matches a set criteria.

# Arguments

- `d::Declarative`: declarative memory object

# Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
get_chunks(d::Declarative; check_value=true, criteria...) = get_chunks(d.memory; check_value, criteria...)

"""
    get_chunks(actr::AbstractACTR; )

Returns all chunks that matches a set criteria

# Arguments

- `actr::AbstractACTR`: an ACTR Object

#Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
get_chunks(actr::AbstractACTR; check_value=true, criteria...) = get_chunks(actr.declarative.memory; check_value, criteria...)

"""
    get_chunks(d::Declarative, funs...; criteria...)

Returns all chunks that matches a set criteria using `funs...` as matching functions

# Arguments

- `d::Declarative`: declarative memory object
- `funs...`: a list of functions

# Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
get_chunks(d::Declarative, funs...; check_value=true, criteria...) = get_chunks(d.memory, funs...; check_value, criteria...)

"""
    get_chunks(actr::AbstractACTR, funs...; criteria...)

Returns all chunks that matches a set criteria using `funs...` as matching functions.

# Arguments

- `actr::AbstractACTR`: an ACT-R model object
- `funs...`: a list of functions

# Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
get_chunks(actr::AbstractACTR, funs...; check_value=true, criteria...) = get_chunks(actr.declarative.memory, funs...; check_value, criteria...)

"""
    first_chunk(memory::Vector{<:Chunk}; criteria...)

Returns the first chunk in memory that matches a set of criteria

# Arguments

- `memory::Vector{<:Chunk}`: a vector of chunks

# Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
function first_chunk(memory::Vector{<:Chunk}; check_value=true, criteria...)
    chunk = Array{eltype(memory),1}()
    for m in memory
        if _match(m, criteria; check_value)
            push!(chunk, m)
            return chunk
        end
    end
    return chunk
end

"""
    first_chunk(d::Declarative; criteria...)

Returns the first chunk in memory that matches a set of criteria

# Arguments

- `d::Declarative`: a declarative memory object

# Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
first_chunk(d::Declarative; check_value=true, criteria...) = first_chunk(d.memory; check_value, criteria...)

"""
    first_chunk(a::AbstractACTR; criteria...)

Returns the first chunk in memory that matches a set of criteria

# Arguments

- `actr::AbstractACTR`: an ACT-R model object

# Keywords

 - `check_value=true`: check slot value

- `criteria...`: optional keyword arguments corresponding to critiria for matching chunk
"""
first_chunk(actr::AbstractACTR; check_value=true, criteria...) = first_chunk(actr.declarative.memory; check_value, criteria...)

"""
    match(chunk::Chunk, request)

Returns a boolean indicating whether a request matches a chunk.
False is returned if the slot does not exist in the chunk or the value
of the slot does not match the request value.

# Arguments

- `chunk::Chunk`: chunk object
- `request`: a NamedTuple of slot value pairs

# Keywords

 - `check_value=true`: check slot value
"""
function _match(chunk::Chunk, request; check_value=true)
    slots = chunk.slots
    for (k,v) in request
        if !(k ∈ keys(slots))
            return false
        elseif check_value
            if (slots[k] != v)
                return false
            end
        end
    end
    return true
end

"""
    match(chunk::Chunk, f, request)

Returns a boolean indicating whether a request matches a chunk.
False is returned if the slot does not exist in the chunk or the value
of the slot does not match the request value.

# Arguments

- `chunk`: a chunk object
- `f`: a list of functions such as `!=, ==`
- `request`: a NamedTuple of slot value pairs

# Keywords

 - `check_value=true`: check slot value
"""
function _match(chunk::Chunk, f, request; check_value=true)
    slots = chunk.slots
    i = 0
    for (k,v) in request
        i += 1
        if !(k ∈ keys(slots))
            return false
        elseif check_value
            if !(f[i](slots[k], v))
                return false
            end
        end
    end
    return true
end

"""
    match(chunk::Chunk; request...)

Returns a boolean indicating whether a request matches a chunk.
False is returned if the slot does not exist in the chunk or the value
of the slot does not match the request value.

# Arguments

- `chunk::Chunk`: a chunk object

# Keywords

- `request...`: optional keyword arguments corresponding to critiria for matching chunk
"""
function match(chunk::Chunk; check_value=true, request...)
    return _match(chunk, request; check_value)
end
"""
    match(chunk::Chunk, funs...; request...)

Returns a boolean indicating whether a request matches a chunk.
False is returned if the slot does not exist in the chunk or the value
of the slot does not match the request value.

- `chunk`: chunk object
- `funs...`: a list of functions such as `!=, ==`
- `request...`: a NamedTuple of slot value pairs
"""
function match(chunk::Chunk, funs...; check_value=true, request...)
    return _match(chunk, funs, request; check_value)
end

"""
    get_subset(actr; request...)

Returns a filtered subset of the retrieval request when partial matching is on.
By default, slot values for isa and retrieved must match exactly.

# Arguments

- `actr`: an ACTR object

# Keywords

- `request...`: an option set of keyword arguments corresponding to a retrieval request.
"""
function get_subset(actr; request...)
    return Iterators.filter(x -> any(s -> s == x[1], actr.declarative.filtered),
    request)
end

"""
    retrieval_request(actr::AbstractACTR; request...)

Returns chunks matching a retrieval request.

# Arguments

- `memory`: declarative memory object

# Keywords

- `request...`: optional keyword arguments corresponding to retrieval request e.g. dog = :fiddo
"""
function retrieval_request(actr::AbstractACTR; request...)
    @unpack mmp = actr.parms
    !mmp ? (return get_chunks(actr; request...)) : nothing
    chunks = get_chunks(actr; check_value=false, request...)
    c = get_subset(actr; request...)
    return get_chunks(chunks; check_value=true, c...)
end

"""
    modify!(c; args...)

Modifies fields of an object

# Arguments

- `c`: an object

# Keywords

- `args...`: optional keywords for field and values
"""
function modify!(c; args...)
    for (k,v) in args
        setfield!(c, k, v)
    end
    return nothing
end

"""
    modify!(c::NamedTuple; args...)

Modifies fields of NamedTuple

# Arguments

- `c`: a NamedTuple

# Keywords

- `args...`: optional keywords for field and values
"""
function modify!(c::NamedTuple; args...)
    for (k,v) in args
        c[k][1] = v
    end
    return nothing
end

"""
    retrieve(actr::AbstractACTR; request...)

Retrieves a chunk given a retrieval request. By default, current time is
computed with `get_time`.

# Arguments

- `actr`: an ACT-R object

# Keywords

- `request...`: optional keyword arguments representing a retrieval request, e.g. person=:bob
"""
function retrieve(actr::AbstractACTR; request...)
    return retrieve(actr, get_time(actr); request...)
end

"""
    retrieve(actr::AbstractACTR, cur_time; request...)

Retrieves a chunk given a retrieval request

# Arguments

- `actr`: an ACT-R object
- `cur_time`: current simulated time in seconds

# Keywords

- `request...`: optional keyword arguments representing a retrieval request, e.g. person=:bob
"""
function retrieve(actr::AbstractACTR, cur_time; request...)
    @unpack declarative,parms = actr
    arr = Array{eltype(declarative.memory),1}()
    chunks = retrieval_request(actr; request...)
    # add noise to threshold even if result of request is empty
    actr.parms.noise ? add_noise!(actr) : (parms.τ′ = parms.τ)
    isempty(chunks) ? (return arr) : nothing
    compute_activation!(actr, chunks, cur_time; request...)
    best = get_max_active(chunks)
    if best[1].act >= parms.τ′
        return best
    end
    return arr
end

"""
    get_max_active(chunks)

Returns the chunk with maximum activation

# Arguments

- `chunks`: a vector of chunk objects
"""
function get_max_active(chunks)
    a = -Inf
    mx = chunks[1]
    for c in chunks
        if c.act > a
            a = c.act
            mx = c
        end
    end
    return [mx]
end

"""
    compute_RT(actr, chunk)

Generates a reaction time for retrieving a chunk based
on the current activation levels of a chunk. If the vector is empty, time
for a retrieval failure is returned.

# Arguments

- `actr`: ACTR object
- `chunk`: a vector that is empty or contains one chunk
"""
function compute_RT(actr, chunk)
    @unpack τ′,lf = actr.parms
    if isempty(chunk)
        return lf * exp(-τ′)
    end
    return lf * exp(-chunk[1].act)
end

"""
    compute_RT(actr, chunk)

Generates a reaction time for retrieving a chunk based
on the current activation levels of a chunk.

# Arguments

- `actr`: ACTR object
- `chunk`: a chunk
"""
function compute_RT(actr, chunk::Chunk)
    @unpack lf = actr.parms
    return lf * exp(-chunk.act)
end

"""
    get_parm(actr, p)

Returns the value of a parameter

# Arguments

- `actr`: ACTR object
- ` p`: parameter name
"""
function get_parm(actr, p)
    misc = actr.parms.misc
    if p in keys(misc)
        return misc[p]
    end
    return getfield(actr.parms, p)
end

"""
    blend_chunks(actr; request...)

Computes blended value over chunks given a retrieval request. By default,
values are blended over the set of slots formed by the set difference between all
slots of a chunk and the slots specified in the retrieval request. The default time used
in activation calculations is taken from `get_time`. Currently, blended
is only supported for numeric slot-values.

# Arguments

- `actr::AbstractACTR`: an `ACTR` model object

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function blend_chunks(actr; request...)
    return blend_chunks(actr, get_time(actr); request...)
end

"""
    blend_chunks(actr::AbstractACTR, cur_time::Float64; request...)

Computes blended value over chunks given a retrieval request. By default,
values are blended over the set of slots formed by the set difference between all
slots of a chunk and the slots specified in the retrieval request. Currently, blended
is only supported for numeric slot-values.

# Arguments

- `actr::AbstractACTR`: an `ACTR` model object
- `cur_time::Float64`: current simulated time

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function blend_chunks(actr::AbstractACTR, cur_time::Float64; request...)
    blended_slots = setdiff(keys(actr.declarative.memory[1].slots), keys(request))
    return blend_chunks(actr, blended_slots, cur_time; request...)
end

"""
    blend_chunks(actr::AbstractACTR; request...)

Computes blended value over chunks given a retrieval request. By default,
values are blended over the set of slots formed by the set difference between all
slots of a chunk and the slots specified in the retrieval request. The default time used
in activation calculations is taken from `get_time`. Currently, blended
is only supported for numeric slot-values.

# Arguments

- `actr::AbstractACTR`: an `ACTR` model object

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function blend_chunks(actr::AbstractACTR; request...)
    return blend_chunks(actr, get_time(actr); request...)
end

"""
    blend_chunks(actr::AbstractACTR, blended_slots; request...)

Computes blended value over chunks given a retrieval request. By default,
values are blended over the set of slots formed by the set difference between all
slots of a chunk and the slots specified in the retrieval request. The default time used
in activation calculations is taken from `get_time(actr`). Currently, blended
is only supported for numeric slot-values.

# Arguments

- `actr::AbstractACTR`: an `ACTR` model object
- `blended_slots`: a set of slots over which slot-values are blended

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function blend_chunks(actr::AbstractACTR, blended_slots; request...)
    return blend_chunks(actr, blended_slots, get_time(actr); request...)
end

"""
    blend_chunks(actr, blended_slots, cur_time; request...)

Computes blended value over chunks given a retrieval request. Values are blended
over the slots specified in `blended_slots`. Currently, blended is only supported
for numeric slot-values.

# Arguments

- `actr`: an `ACTR` model object
- `blended_slots`: a set of slots over which slot-values are blended
- `cur_time`: current simulated time

# Keywords

- `request...`: optional keywords for the retrieval request
"""
function blend_chunks(actr, blended_slots, cur_time; request...)
    chunks = retrieval_request(actr; request...)
    compute_activation!(actr, chunks, cur_time; request...)
    probs = soft_max(actr, chunks)
    v = blend_slots(chunks, probs, blended_slots)
    return v
end

function blend_slots(chunks, probs, slots)
    T = typeof(probs)
    v::T = map(s -> blend_slots(chunks, probs, s), slots)
    return v
end

"""
    blend_slots(chunks, probs, slot::Symbol)

Computes an expected value over chunks for a specified slot.

# Arguments

- `chunks`: a set of chunks over which slot-values are blended
- `probs`: a vector of retrieval probabilities
- `slot::Symbol`: a slot over which slot-values are blended
"""
function blend_slots(chunks, probs, slot::Symbol)
    T = eltype(probs)
    v::T = 0.0
    for (c,p) in zip(chunks, probs)
        v += p * c.slots[slot]
    end
    return v
end

# non numeric
# blend_chunk(actr, chunks, probs, chunk) = blend_chunk(actr, chunks, probs; chunk.slots...)

# function blend_chunk(actr, chunks, probs; slots...)
#     penalties = actr.parms.mmp_fun.(actr, chunks; slots...)
#     return probs' * penalties
# end

function soft_max(actr, chunks)
    σ = actr.parms.s * sqrt(2)
    v = map(x -> exp(x.act / σ), chunks)
    return v ./ sum(v)
end

"""
    blended_activation(chunks)

Computes a blended activation value by exponentiating, summing and taking the
log of activations across a set of chunks.

# Arguments

- `chunks`: a set of chunks over which slot-values are blended
"""
function blended_activation(chunks)
    exp_act = map(x -> exp(x.act_mean), chunks)
    return log(sum(exp_act))
end

"""
    compute_RT(blended_act)

Computes retrieval time for a given blended activation value.

# Arguments

- `blended_act`: a blended activation value
"""
compute_RT(blended_act) = exp(-blended_act)


text:
# * Media

using Unitful
import Unitful: Length, Area, Frequency
import Base: show
using Printf

@doc raw"""
    Medium(B, C)

The [Sellmeier
equations](https://en.wikipedia.org/wiki/Sellmeier_equation) are used
to describe dispersion in glass using a series of resonances:

```math
n^2(\lambda) =
1 + \sum_i \frac {B_i\lambda^2}{\lambda^2-C_i}.
```
"""
struct Medium{T<:Real,A<:Area{T}}
    B::Vector{T}
    C::Vector{A}
end

function show(io::IO, m::Medium)
    write(io, "Medium(")
    show(io, m.B)
    write(io, @sprintf(", [%s])",
                       join(string.(m.C), ", ")))
end

# ** BK7
"""
    BK7

[Borosilicate glass](https://en.wikipedia.org/wiki/Borosilicate_glass)
is commonly used in optical lenses
"""
BK7 = Medium([1.03961212, 0.231792344, 1.01046945],
             [6.00069867e-3u"μm"^2, 2.00179144e-2u"μm"^2, 1.03560653e2u"μm"^2])

# ** SiO₂

"""
    SiO₂

[Fused silica](https://en.wikipedia.org/wiki/Fused_quartz#Optical_properties)
"""
SiO₂ = Medium([0.6961663, 0.4079426, 0.8974794],
              [0.0684043u"μm", 0.1162414u"μm", 9.896161u"μm"].^2)

# * Refractive index

function n²(m::Medium, λ::Length)
    if !isfinite(λ)
        1.0
    else
        1.0 + sum([Bᵢ*λ^2/(λ^2-Cᵢ)
                   for (Bᵢ,Cᵢ) ∈ zip(m.B, m.C)]) |> NoUnits
    end
end

maybe_complex(x) = x<0 ? complex(x) : x
refractive_index(m::Medium, λ::Length) = √(maybe_complex(n²(m, λ)))
refractive_index(m::Medium, f::Frequency) = refractive_index(m, u"c"/f)

(m::Medium)(x) = refractive_index(m, x)

# * Dispersion

"""
    dispersion(m, d, f[, f₀=0u"Hz])

Calculate dispersion through a medium `m` of length `d`. Optionally,
remove central frequency k-vector, to keep pulse temporally centred.
"""
function dispersion(m::Medium, d::Length, f::AbstractVector{F}, f₀::Frequency = 0u"Hz") where {F<:Frequency}
    n = real(m.(f))

    # Wavevector in vacuum
    k₀ = 2π*f./u"c"
    # Wavevector in the medium
    k = n.*k₀

    if !iszero(f₀)
        # Slope of dispersion relation at central frequency
        ∂k∂f = (2π/u"c")*ForwardDiff.derivative(f -> real(m(f*u"Hz"))*f, f₀/u"Hz" .|> NoUnits)
        k -= ∂k∂f*f
    end

    exp.(-im*k*d)
end
# * Exports

export Medium, BK7, SiO₂, refractive_index, dispersion


text:
# For the conversion to go through ForwardDiff, I need
#Base.convert(FDT::Type{ForwardDiff.Dual{T,V,N}}, p::AIBECS.Parameters) where {T,V,N} = AIBECS.Parameters(convert(Vector{FDT}, vec(p))...)

@testset "Derivatives" begin
    nt = length(T_all)
    n = nt * nb
    x₀ = p₀.xgeo * ones(n)
    @testset "∇ₓF" begin
        @test ForwardDiff.jacobian(x -> F(x, p₀), x₀) ≈ ∇ₓF(x₀, p₀) rtol = 1e-14
        @test ForwardDiff.jacobian(x -> F(x, 2p₀), x₀) ≈ ∇ₓF(x₀, 2p₀) rtol = 1e-14
        @test ForwardDiff.jacobian(x -> F(x, p₀), 2x₀) ≈ ∇ₓF(2x₀, p₀) rtol = 1e-14
        @test ForwardDiff.jacobian(x -> F(x, p₀), -x₀) ≈ ∇ₓF(-x₀, p₀) rtol = 1e-14
    end
    @testset "∇ₚf" begin
        @test ForwardDiff.jacobian(p -> [f(x₀, p)], p₀) ≈ ∇ₚf(x₀, p₀) rtol = 1e-14
        @test ForwardDiff.jacobian(p -> [f(x₀, p)], 2p₀) ≈ ∇ₚf(x₀, 2p₀) rtol = 1e-14
        @test ForwardDiff.jacobian(p -> [f(2x₀, p)], p₀) ≈ ∇ₚf(2x₀, p₀) rtol = 1e-14
    end
    @testset "∇ₓf" begin
        @test ForwardDiff.jacobian(x -> [f(x, p₀)], x₀) ≈ ∇ₓf(x₀, p₀) rtol = 1e-14
        @test ForwardDiff.jacobian(x -> [f(x, 2p₀)], x₀) ≈ ∇ₓf(x₀, 2p₀) rtol = 1e-14
        @test ForwardDiff.jacobian(x -> [f(x, p₀)], 2x₀) ≈ ∇ₓf(2x₀, p₀) rtol = 1e-14
    end
end


text:
using Documenter
push!(LOAD_PATH, joinpath(@__DIR__, "../src/"))
using Q

makedocs(
    modules   = [Q, Q._k],
    sitename  = "Q.jl - Julia for kdb+",
    linkcheck = true,
    pages = [
        "Home" => "index.md",
        "Manual" => [
            "manual/common-api.md",
            "manual/client.md",
            "manual/server.md",
            "manual/low-level.md",
        ]
    ],
)

deploydocs(
    deps   = Deps.pip("mkdocs", "python-markdown-math"),
    repo = "github.com/enlnt/Q.jl.git",
    julia  = "0.6",
    osname = "linux",
)


text:
using Benchmark

################################################################################
# We could decide which pooling to use inside a loop, or we could
# write a different loop for different pooling. The latter leads
# to code duplication, but better performance. This benchmark is
# to show how much performance different this two ways have. It
# turns out the difference is not so significant for this case.
#
# | Row | Function       | Average | Relative | Replications |
# |-----|----------------|---------|----------|--------------|
# | 1   | "loop_in_cond" | 5.44311 | 1.0      | 10           |
# | 2   | "cond_in_loop" | 5.88683 | 1.08152  | 10           |
################################################################################

function cond_in_loop(input, kernel, pooling)
  width, height, channels, num = size(input)
  pooled_width = width - kernel[1] + 1
  pooled_height = height - kernel[2] + 1

  output = zeros(pooled_width, pooled_height, channels, num)
  kernel_size = prod(kernel)

  for n = 1:num
    for c = 1:channels
      for pw = 1:pooled_width
        for ph = 1:pooled_height
          hstart = ph
          wstart = pw
          hend = min(hstart + kernel[2] - 1, height)
          wend = min(wstart + kernel[1] - 1, width)
          if pooling == :max
            maxval = -Inf
            for u = wstart:wend
              for v = hstart:hend
                maxval = max(maxval, input[u,v,c,n])
              end
            end
            output[pw,ph,c,n] = maxval
          elseif pooling == :mean
            the_sum = 0
            for u = wstart:wend
              for v = hstart:hend
                the_sum += input[u,v,c,n]
              end
            end
            output[pw,ph,c,n] = the_sum / kernel_size
          end
        end
      end
    end
  end

  return output
end

function loop_in_cond(input, kernel, pooling)
  width, height, channels, num = size(input)
  pooled_width = width - kernel[1] + 1
  pooled_height = height - kernel[2] + 1

  output = zeros(pooled_width, pooled_height, channels, num)
  if pooling == :max
    for n = 1:num
      for c = 1:channels
        for pw = 1:pooled_width
          for ph = 1:pooled_height
            hstart = ph
            wstart = pw
            hend = min(hstart + kernel[2] - 1, height)
            wend = min(wstart + kernel[1] - 1, width)
            maxval = -Inf
            for u = wstart:wend
              for v = hstart:hend
                maxval = max(maxval, input[u,v,c,n])
              end
            end
            output[pw,ph,c,n] = maxval
          end
        end
      end
    end
  elseif pooling == :mean
    kernel_size = prod(kernel)

    for n = 1:num
      for c = 1:channels
        for pw = 1:pooled_width
          for ph = 1:pooled_height
            hstart = ph
            wstart = pw
            hend = min(hstart + kernel[2] - 1, height)
            wend = min(wstart + kernel[1] - 1, width)
            the_sum = 0
            for u = wstart:wend
              for v = hstart:hend
                the_sum += input[u,v,c,n]
              end
            end
            output[pw,ph,c,n] = the_sum / kernel_size
          end
        end
      end
    end
  end

  return output
end

input = rand(28, 28, 50, 128)
kernel = (5,5)
cond_in_loop() = (cond_in_loop(input,kernel,:max),
                  cond_in_loop(input,kernel,:mean))
loop_in_cond() = (loop_in_cond(input,kernel,:max),
                  loop_in_cond(input,kernel,:mean))
maxpool1, meanpool1 = cond_in_loop()
maxpool2, meanpool2 = cond_in_loop()
@assert all(abs.(maxpool1-maxpool2) .< 1e-10)
@assert all(abs.(meanpool1-meanpool2) .< 1e-10)

println("Running benchmark")
df = compare([loop_in_cond, cond_in_loop], 10)
println("$df")


text:
module WeirdMatrices

using LinearAlgebra
import LinearAlgebra: Matrix, AbstractMatrix, Array,
    Adjoint, Transpose,
    transpose, adjoint, diag, mul!, _modify!, MulAddMul
import Base: size, similar,
    +, -, *, /, ==,
    conj, copy, real, imag,
    require_one_based_indexing, getindex

using BandedMatrices
using SparseArrays

include("skew_tridiagonal.jl")
include("one_block_matrix.jl")
include("block_sparse_diagonal.jl")
include("column_sparse_matrix.jl")

end


text:
using Rayuela

# === Main
Xt = fvecs_read(Int(1e4), "../data/sift/sift_learn.fvecs")
m  = 8   # Number of codebooks
h  = 256 # Number of entries in each codebook
lr = 0.5f0 # learning rate
H  = 32 # depth for search during encoding

n_its = 250

C, B = train_rvq(Xt, m, h, 25, true)
qerr = qerror(Xt, B, C)
@show B[:, 1]
print( "Error after initialization is $qerr\n" )

train_competitiveq(Xt, C, n_its, H, B, lr )


text:
using NASAEarthObservations
using Test

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