text:
#
# Differential operators for regular grids
#

∂x_caa(i, j, k, grid, u::AbstractArray) = δx_f2c(grid, u, i, j, k) / grid.Δx
∂y_aca(i, j, k, grid, v::AbstractArray) = δy_f2c(grid, v, i, j, k) / grid.Δy
∂z_aac(i, j, k, grid, w::AbstractArray) = δz_f2c(grid, w, i, j, k) / grid.Δz
∂x_faa(i, j, k, grid, ϕ::AbstractArray) = δx_c2f(grid, ϕ, i, j, k) / grid.Δx
∂y_afa(i, j, k, grid, ϕ::AbstractArray) = δy_c2f(grid, ϕ, i, j, k) / grid.Δy
∂z_aaf(i, j, k, grid, ϕ::AbstractArray) = δz_c2f(grid, ϕ, i, j, k) / grid.Δz

#
# Differentiation and interpolation operators for functions
#
# The geometry of the Oceananigans grid (in one dimension) is
#
# face   cell   face   cell   face
#
#         i-1            i
#          ↓             ↓
#   |      ×      |      ×      |
#   ↑             ↑             ↑
#  i-1            i            i+1
#
# As a result the interpolation of a quantity ϕᵢ from a cell i to face i
# (this operation is denoted "▶x_faa" in the code below) is
#
# (▶ϕ)ᵢ = (ϕᵢ + ϕᵢ₋₁) / 2 .
#
# Conversely, the interpolation of a quantity uᵢ from face i to cell i is given by
#
# (▶u)ᵢ = (uᵢ₊₁ + uᵢ) / 2.
#
# Derivative operators are defined similarly. Using the symbol "∂" to denote
# differentiation, we have
#
# (∂ϕ)ᵢ = (ϕᵢ - ϕᵢ₋₁) / Δ
#
# for the differentation of a CellField ending on face i, and
#
# (∂u)ᵢ = (uᵢ₊₁ - uᵢ) / Δ
#
# for the differentation of a FaceField ending on cell i, where Δ is the spacing
# between cell points or face points, appropriately.
#

#
# Differential operators
#

"""
    ∂x_faa(i, j, k, grid, F, args...)

Differentiate the function or callable object

    `F(i, j, k, grid, args...)`

located at `caa` in `x`, across `faa`.
"""
function ∂x_faa(i, j, k, grid, F::Function, args...)
    i⁻ = decmod1(i, grid.Nx)
    return (F(i, j, k, grid, args...) - F(i⁻, j, k, grid, args...)) / grid.Δx
end

"""
    ∂x_caa(i, j, k, grid, F, args...)

Differentiate the function or callable object

    `F(i, j, k, grid, args...)`

located at `faa` in `x`, across `caa`.
"""
function ∂x_caa(i, j, k, grid, F::Function, args...)
    i⁺ = incmod1(i, grid.Nx)
    return (F(i⁺, j, k, grid, args...) - F(i, j, k, grid, args...)) / grid.Δx
end

"""
    ∂y_afa(i, j, k, grid, F, args...)

Differentiate the function or callable object

    `F(i, j, k, grid, args...)`

located at `aca` in `y`, across `afa`.
"""
function ∂y_afa(i, j, k, grid, F::Function, args...)
    j⁻ = decmod1(j, grid.Ny)
    return (F(i, j, k, grid, args...) - F(i, j⁻, k, grid, args...)) / grid.Δy
end

"""
    ∂y_aca(i, j, k, grid, F, args...)

Differentiate the function or callable object

    `F(i, j, k, grid, args...)`

located at `afa` in `y`, across `aca`.
"""
function ∂y_aca(i, j, k, grid, F::Function, args...)
    j⁺ = incmod1(j, grid.Ny)
    return (F(i, j⁺, k, grid, args...) - F(i, j, k, grid, args...)) / grid.Δy
end

"""
    ∂z_aaf(i, j, k, grid, F, args...)

Differentiate the function or callable object

    `F(i, j, k, grid, args...)`

located at `aac` in `z`, across `aaf`.
"""
function ∂z_aaf(i, j, k, grid::Grid{T}, F::Function, args...) where T
    if k < 2
        return -zero(T)
    else
        return (F(i, j, k-1, grid, args...) - F(i, j, k, grid, args...)) / grid.Δz
    end
end

"""
    ∂z_aac(i, j, k, grid, F, args...)

Differentiate the function or callable object

    `F(i, j, k, grid, args...)`

located at `aaf` in `z`, across `aac`.
"""
function ∂z_aac(i, j, k, grid::Grid{T}, F::Function, args...) where T
    if k == grid.Nz
        return -zero(T)
    else
        return (F(i, j, k, grid, args...) - F(i, j, k+1, grid, args...)) / grid.Δz
    end
end

#
# Double differentiation
#

∂x²_caa(i, j, k, grid, u) = ∂x_faa(i, j, k, grid, ∂x_caa, u)
∂x²_faa(i, j, k, grid, ϕ) = ∂x_caa(i, j, k, grid, ∂x_faa, ϕ)

∂y²_aca(i, j, k, grid, v) = ∂y_afa(i, j, k, grid, ∂y_aca, v)
∂y²_afa(i, j, k, grid, ϕ) = ∂y_aca(i, j, k, grid, ∂y_afa, ϕ)

∂z²_aac(i, j, k, grid, w) = ∂z_aaf(i, j, k, grid, ∂z_aac, w)
∂z²_aaf(i, j, k, grid, ϕ) = ∂z_aac(i, j, k, grid, ∂z_aaf, ϕ)

#
# Interpolation operations for functions
#

"""
    ▶x_faa(i, j, k, grid, F, args...)

Interpolate the function or callable object

    `F(i, j, k, grid, args...)`

from `caa` to `faa`."
"""
function ▶x_faa(i, j, k, grid::Grid{T}, F::Function, args...) where T
    i⁻¹ = decmod1(i, grid.Nx)
    return T(0.5) * (F(i, j, k, grid, args...) + F(i⁻¹, j, k, grid, args...))
end

"""
    ▶x_caa(i, j, k, grid, F, args...)

Interpolate the function or callable object

    `F(i, j, k, grid, args...)`

from `faa` to `caa`."
"""
function ▶x_caa(i, j, k, grid::Grid{T}, F::Function, args...) where T
    i⁺¹ = incmod1(i, grid.Nx)
    return T(0.5) * (F(i⁺¹, j, k, grid, args...) + F(i, j, k, grid, args...))
end

"""
    ▶y_afa(i, j, k, grid, F, args...)

Interpolate the function or callable object

    `F(i, j, k, grid, args...)`

from `aca` to `afa`.
"""
function ▶y_afa(i, j, k, grid::Grid{T}, F::Function, args...) where T
    j⁻¹ = decmod1(j, grid.Ny)
    return T(0.5) * (F(i, j, k, grid, args...) + F(i, j⁻¹, k, grid, args...))
end

"""
    ▶y_aca(i, j, k, grid, F, args...)

Interpolate the function or callable object

    `F(i, j, k, grid, args...)`

from `afa` to `aca`."
"""
function ▶y_aca(i, j, k, grid::Grid{T}, F::Function, args...) where T
    j⁺¹ = incmod1(j, grid.Ny)
    return T(0.5) * (F(i, j⁺¹, k, grid, args...) + F(i, j, k, grid, args...))
end

"""
    ▶z_aaf(i, j, k, grid, F, args...)

Interpolate the function or callable object

    `F(i, j, k, grid, args...)`

from `aac` to `aaf`.
"""
function ▶z_aaf(i, j, k, grid::Grid{T}, F::Function, args...) where T
    if k == 1
        return T(0.5) * F(i, j, k, grid, args...)
    else
        return T(0.5) * (F(i, j, k, grid, args...) + F(i, j, k-1, grid, args...))
    end
end

"""
    ▶z_aac(i, j, k, grid::Grid{T}, F, args...) where T

Interpolate the function or callable object

    `F(i, j, k, grid, args...)`

from `aaf` to `aac`.
"""
function ▶z_aac(i, j, k, grid::Grid{T}, F::Function, args...) where T
    if k == grid.Nz
        return T(0.5) * F(i, j, k, grid, args...)
    else
        return T(0.5) * (F(i, j, k+1, grid, args...) + F(i, j, k, grid, args...))
    end
end

# Convenience operators for "interpolating constants"
▶x_faa(i, j, k, grid, F::Number, args...) = F
▶x_caa(i, j, k, grid, F::Number, args...) = F
▶y_afa(i, j, k, grid, F::Number, args...) = F
▶y_aca(i, j, k, grid, F::Number, args...) = F
▶z_aaf(i, j, k, grid, F::Number, args...) = F
▶z_aac(i, j, k, grid, F::Number, args...) = F

#
# Double interpolation: 12 operators
#

"""
    ▶xy_cca(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `y`, from `ffa` to `cca`.
"""
▶xy_cca(i, j, k, grid, F, args...) = ▶y_aca(i, j, k, grid, ▶x_caa, F, args...)

"""
    ▶xy_fca(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `y`, from `cfa` to `fca`.
"""
▶xy_fca(i, j, k, grid, F, args...) = ▶y_aca(i, j, k, grid, ▶x_faa, F, args...)

"""
    ▶xy_ffa(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `y`, from `cca` to `ffa`.
"""
▶xy_ffa(i, j, k, grid, F, args...) = ▶y_afa(i, j, k, grid, ▶x_faa, F, args...)

"""
    ▶xy_cfa(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `y`, from `fca` to `cfa`.
"""
▶xy_cfa(i, j, k, grid, F, args...) = ▶y_afa(i, j, k, grid, ▶x_caa, F, args...)

"""
    ▶xz_cac(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `z`, from `faf` to `cac`.
"""
▶xz_cac(i, j, k, grid, F, args...) = ▶z_aac(i, j, k, grid, ▶x_caa, F, args...)

"""
    ▶xz_fac(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `z`, from `caf` to `fac`.
"""
▶xz_fac(i, j, k, grid, F, args...) = ▶z_aac(i, j, k, grid, ▶x_faa, F, args...)

"""
    ▶xz_faf(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `z`, from `cac` to `faf`.
"""
▶xz_faf(i, j, k, grid, F, args...) = ▶z_aaf(i, j, k, grid, ▶x_faa, F, args...)

"""
    ▶xz_caf(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `x` and `z`, from `fac` to `caf`.
"""
▶xz_caf(i, j, k, grid, F, args...) = ▶z_aaf(i, j, k, grid, ▶x_caa, F, args...)

"""
    ▶yz_acc(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `y` and `z`, from `aff` to `acc`.
"""
▶yz_acc(i, j, k, grid, F, args...) = ▶z_aac(i, j, k, grid, ▶y_aca, F, args...)

"""
    ▶yz_afc(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `y` and `z`, from `acf` to `afc`.
"""
▶yz_afc(i, j, k, grid, F, args...) = ▶z_aac(i, j, k, grid, ▶y_afa, F, args...)

"""
    ▶yz_aff(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `y` and `z`, from `ffa` to `cca`.
"""
▶yz_aff(i, j, k, grid, F, args...) = ▶z_aaf(i, j, k, grid, ▶y_afa, F, args...)

"""
    ▶yz_acf(i, j, k, grid, F, args...)

Interpolate the function

    `F(i, j, k, grid, args...)`

in `y` and `z`, from `afc` to `acf`.
"""
▶yz_acf(i, j, k, grid, F, args...) = ▶z_aaf(i, j, k, grid, ▶y_aca, F, args...)

"""
    ∂x_caa(κ, Σ, i, j, k, grid, closure, u, v, w, T, S)

Compute ``\\partial_x ( \\kappa \\Sigma )``, where `κ` and `Σ`
are diffusivity and strain functions that return quantities
at the location `faa`, where `a` is either `f` or `c`.
The quantity returned has the location `caa`.
"""
function ∂x_caa(κ::Function, Σ::Function, i, j, k, grid, closure, eos, g, u, v, w, T, S)
    i⁺¹ = incmod1(i, grid.Nx)

    Σ⁺¹ = Σ(i⁺¹, j, k, grid, u, v, w)
    Σ⁻¹ = Σ(i,   j, k, grid, u, v, w)

    κ⁺¹ = κ(i⁺¹, j, k, grid, closure, eos, g, u, v, w, T, S)
    κ⁻¹ = κ(i,   j, k, grid, closure, eos, g, u, v, w, T, S)

    return (κ⁺¹ * Σ⁺¹ - κ⁻¹ * Σ⁻¹) / grid.Δx
end

"""
    ∂x_faa(κ, Σ, i, j, k, grid, closure, u, v, w, T, S)

Compute ``\\partial_x ( \\kappa \\Sigma )``, where `κ` and `Σ`
are diffusivity and strain functions that return quantities
at the location `caa`, where `a` is either `f` or `c`.
The quantity returned has the location `faa`.
"""
function ∂x_faa(κ::Function, Σ::Function, i, j, k, grid, closure, eos, g, u, v, w, T, S)
    i⁻¹ = decmod1(i, grid.Nx)

    Σ⁺¹ = Σ(i,   j, k, grid, u, v, w)
    Σ⁻¹ = Σ(i⁻¹, j, k, grid, u, v, w)

    κ⁺¹ = κ(i,   j, k, grid, closure, eos, g, u, v, w, T, S)
    κ⁻¹ = κ(i⁻¹, j, k, grid, closure, eos, g, u, v, w, T, S)

    return (κ⁺¹ * Σ⁺¹ - κ⁻¹ * Σ⁻¹) / grid.Δx
end

"""
    ∂y_aca(κ, Σ, i, j, k, grid, closure, u, v, w, T, S)

Compute ``\\partial_y ( \\kappa \\Sigma )``, where `κ` and `Σ`
are diffusivity and strain functions that return quantities
at the location `afa`, where `a` is either `f` or `c`.
The quantity returned has the location `aca`.
"""
function ∂y_aca(κ::Function, Σ::Function, i, j, k, grid, closure, eos, g, u, v, w, T, S)
    j⁺¹ = incmod1(j, grid.Ny)

    Σ⁺¹ = Σ(i, j⁺¹, k, grid, u, v, w)
    Σ⁻¹ = Σ(i, j,   k, grid, u, v, w)

    κ⁺¹ = κ(i, j⁺¹, k, grid, closure, eos, g, u, v, w, T, S)
    κ⁻¹ = κ(i, j,   k, grid, closure, eos, g, u, v, w, T, S)

    return (κ⁺¹ * Σ⁺¹ - κ⁻¹ * Σ⁻¹) / grid.Δy
end

"""
    ∂y_afa(κ, Σ, i, j, k, grid, closure, u, v, w, T, S)

Compute ``\\partial_y ( \\kappa \\Sigma )``, where `κ` and `Σ`
are diffusivity and strain functions that return quantities
at the location `aca`, where `a` is either `f` or `c`.
The quantity returned has the location `afa`.
"""
function ∂y_afa(κ::Function, Σ::Function, i, j, k, grid, closure, eos, g, u, v, w, T, S)
    j⁻¹ = decmod1(j, grid.Ny)

    Σ⁺¹ = Σ(i, j,   k, grid, u, v, w)
    Σ⁻¹ = Σ(i, j⁻¹, k, grid, u, v, w)

    κ⁺¹ = κ(i, j,   k, grid, closure, eos, g, u, v, w, T, S)
    κ⁻¹ = κ(i, j⁻¹, k, grid, closure, eos, g, u, v, w, T, S)

    return (κ⁺¹ * Σ⁺¹ - κ⁻¹ * Σ⁻¹) / grid.Δy
end

"""
    ∂z_aac(κ, Σ, i, j, k, grid, closure, u, v, w, T, S)

Compute ``\\partial_z ( \\kappa \\Sigma )``, where `κ` and `Σ`
are diffusivity and strain functions that return quantities
at the location `aaf`, where `a` is either `f` or `c`.
The quantity returned has the location `aac`.
"""
function ∂z_aac(κ::Function, Σ::Function, i, j, k, grid, closure, eos, g, u, v, w, T, S)
    Σ⁻¹ = Σ(i, j, k, grid, u, v, w)
    κ⁻¹ = κ(i, j, k, grid, closure, eos, g, u, v, w, T, S)

    if k == grid.Nz
        return κ⁻¹ * Σ⁻¹ / grid.Δz
    else
        Σ⁺¹ = Σ(i, j, k+1, grid, u, v, w)
        κ⁺¹ = κ(i, j, k+1, grid, closure, eos, g, u, v, w, T, S)

        # Reversed due to inverse convention for z-indexing
        return (κ⁻¹ * Σ⁻¹ - κ⁺¹ * Σ⁺¹) / grid.Δz
    end
end

"""
    ∂z_aaf(κ, Σ, i, j, k, grid, closure, u, v, w, T, S)

Compute ``\\partial_z ( \\kappa \\Sigma )``, where `κ` and `Σ`
are diffusivity and strain functions that return quantities
at the location `aac`, where `a` is either `f` or `c`.
The quantity returned has the location `aaf`.
"""
function ∂z_aaf(κ::Function, Σ::Function, i, j, k, grid, closure, eos, g, u, v, w, T, S)
    Σ⁺¹ = Σ(i, j, k, grid, u, v, w)
    κ⁺¹ = κ(i, j, k, grid, closure, eos, g, u, v, w, T, S)

    if k == 1
        return - κ⁺¹ * Σ⁺¹ / grid.Δz
    else
        Σ⁻¹ = Σ(i, j, k-1, grid, u, v, w)
        κ⁻¹ = κ(i, j, k-1, grid, closure, eos, g, u, v, w, T, S)

        # Reversed due to inverse convention for z-indexing
        return (κ⁻¹ * Σ⁻¹ - κ⁺¹ * Σ⁺¹) / grid.Δz
    end
end

# At fcc
∂x_2ν_Σ₁₁(ijk...) = 2 * ∂x_faa(ν_ccc, Σ₁₁, ijk...)
∂y_2ν_Σ₁₂(ijk...) = 2 * ∂y_aca(ν_ffc, Σ₁₂, ijk...)
∂z_2ν_Σ₁₃(ijk...) = 2 * ∂z_aac(ν_fcf, Σ₁₃, ijk...)

# At cfc
∂x_2ν_Σ₂₁(ijk...) = 2 * ∂x_caa(ν_ffc, Σ₂₁, ijk...)
∂y_2ν_Σ₂₂(ijk...) = 2 * ∂y_afa(ν_ccc, Σ₂₂, ijk...)
∂z_2ν_Σ₂₃(ijk...) = 2 * ∂z_aac(ν_cff, Σ₂₃, ijk...)

# At ccf
∂x_2ν_Σ₃₁(ijk...) = 2 * ∂x_caa(ν_fcf, Σ₃₁, ijk...)
∂y_2ν_Σ₃₂(ijk...) = 2 * ∂y_aca(ν_cff, Σ₃₂, ijk...)
∂z_2ν_Σ₃₃(ijk...) = 2 * ∂z_aaf(ν_ccc, Σ₃₃, ijk...)

"""
    ∂ⱼ_2ν_Σ₁ⱼ(i, j, k, grid, closure, eos, g, u, v, w, T, S)

Return the ``x``-component of the turbulent diffusive flux divergence:

`∂x(2 ν Σ₁₁) + ∂y(2 ν Σ₁₁) + ∂z(2 ν Σ₁₁)`

at the location `fcc`.
"""
∂ⱼ_2ν_Σ₁ⱼ(i, j, k, grid, closure::IsotropicDiffusivity, args...) = (
    ∂x_2ν_Σ₁₁(i, j, k, grid, closure, args...)
  + ∂y_2ν_Σ₁₂(i, j, k, grid, closure, args...)
  + ∂z_2ν_Σ₁₃(i, j, k, grid, closure, args...)
)

"""
    ∂ⱼ_2ν_Σ₂ⱼ(i, j, k, grid, closure, eos, g, u, v, w, T, S)

Return the ``y``-component of the turbulent diffusive flux divergence:

`∂x(2 ν Σ₂₁) + ∂y(2 ν Σ₂₂) + ∂z(2 ν Σ₂₂)`

at the location `ccf`.
"""
∂ⱼ_2ν_Σ₂ⱼ(i, j, k, grid, closure::IsotropicDiffusivity, args...) = (
    ∂x_2ν_Σ₂₁(i, j, k, grid, closure, args...)
  + ∂y_2ν_Σ₂₂(i, j, k, grid, closure, args...)
  + ∂z_2ν_Σ₂₃(i, j, k, grid, closure, args...)
)

"""
    ∂ⱼ_2ν_Σ₃ⱼ(i, j, k, grid, closure, eos, g, u, v, w, T, S)

Return the ``z``-component of the turbulent diffusive flux divergence:

`∂x(2 ν Σ₃₁) + ∂y(2 ν Σ₃₂) + ∂z(2 ν Σ₃₃)`

at the location `ccf`.
"""
∂ⱼ_2ν_Σ₃ⱼ(i, j, k, grid, closure::IsotropicDiffusivity, args...) = (
    ∂x_2ν_Σ₃₁(i, j, k, grid, closure, args...)
  + ∂y_2ν_Σ₃₂(i, j, k, grid, closure, args...)
  + ∂z_2ν_Σ₃₃(i, j, k, grid, closure, args...)
)

"""
    κ_∂x_ϕ(i, j, k, grid, ϕ, κ, closure, eos, g, u, v, w, T, S)

Return `κ ∂x ϕ`, where `κ` is a function that computes
diffusivity at cell centers (location `ccc`), and `ϕ` is an array of scalar
data located at cell centers.
"""
function κ_∂x_ϕ(i, j, k, grid, ϕ, κ, closure, args...)
    κ = ▶x_faa(i, j, k, grid, κ, closure, args...)
    ∂x_ϕ = ∂x_faa(i, j, k, grid, ϕ)
    return κ * ∂x_ϕ
end

"""
    κ_∂y_ϕ(i, j, k, grid, ϕ, κ, closure, eos, g, u, v, w, T, S)

Return `κ ∂y ϕ`, where `κ` is a function that computes
diffusivity at cell centers (location `ccc`), and `ϕ` is an array of scalar
data located at cell centers.
"""
function κ_∂y_ϕ(i, j, k, grid, ϕ, κ, closure, args...)
    κ = ▶y_afa(i, j, k, grid, κ, closure, args...)
    ∂y_ϕ = ∂y_afa(i, j, k, grid, ϕ)
    return κ * ∂y_ϕ
end

"""
    κ_∂z_ϕ(i, j, k, grid, ϕ, κ, closure, eos, g, u, v, w, T, S)

Return `κ ∂z ϕ`, where `κ` is a function that computes
diffusivity at cell centers (location `ccc`), and `ϕ` is an array of scalar
data located at cell centers.
"""
function κ_∂z_ϕ(i, j, k, grid, ϕ, κ, closure, args...)
    κ = ▶z_aaf(i, j, k, grid, κ, closure, args...)
    ∂z_ϕ = ∂z_aaf(i, j, k, grid, ϕ)
    return κ * ∂z_ϕ
end

"""
    ∂x_κ_∂x_ϕ(i, j, k, grid, ϕ, κ, closure, eos, g, u, v, w, T, S)

Return `(∂x κ ∂x) ϕ`, where `κ` is a function that computes
diffusivity at cell centers (location `ccc`), and `ϕ` is an array of scalar
data located at cell centers.
"""
∂x_κ_∂x_ϕ(i, j, k, grid, ϕκ...) = ∂x_caa(i, j, k, grid, κ_∂x_ϕ, ϕκ...)

"""
    ∂y_κ_∂y_ϕ(i, j, k, grid, ϕ, κ, closure, eos, g, u, v, w, T, S)

Return `(∂y κ ∂y) ϕ`, where `κ` is a function that computes
diffusivity at cell centers (location `ccc`), and `ϕ` is an array of scalar
data located at cell centers.
"""
∂y_κ_∂y_ϕ(i, j, k, grid, ϕκ...) = ∂y_aca(i, j, k, grid, κ_∂y_ϕ, ϕκ...)

"""
    ∂z_κ_∂z_ϕ(i, j, k, grid, ϕ, κ, closure, eos, g, u, v, w, T, S)

Return `(∂z κ ∂z) ϕ`, where `κ` is a function that computes
diffusivity at cell centers (location `ccc`), and `ϕ` is an array of scalar
data located at cell centers.
"""
∂z_κ_∂z_ϕ(i, j, k, grid, ϕκ...) = ∂z_aac(i, j, k, grid, κ_∂z_ϕ, ϕκ...)

"""
    ∇_κ_∇_ϕ(i, j, k, grid, ϕ, closure, eos, g, u, v, w, T, S)

Return the diffusive flux divergence `∇ ⋅ (κ ∇ ϕ)` for the turbulence
`closure`, where `ϕ` is an array of scalar data located at cell centers.
"""
∇_κ_∇ϕ(i, j, k, grid, ϕ, closure::IsotropicDiffusivity, args...) = (
      ∂x_κ_∂x_ϕ(i, j, k, grid, ϕ, κ_ccc, closure, args...)
    + ∂y_κ_∂y_ϕ(i, j, k, grid, ϕ, κ_ccc, closure, args...)
    + ∂z_κ_∂z_ϕ(i, j, k, grid, ϕ, κ_ccc, closure, args...)
    )


text:
__precompile__()

module Ftools

	export @stack, @queue, List ,list
	export rep
	export @set

	mutable struct List
		   first::Any
		   second::List
		   List() = (x = new(nothing); x.second = x )
		   List(first, second::List) = new(first,second)
	       end

	list(value::Any) = List(value, List())
	list() 		 = List()

	function Base.show( io::IO, x::List)
		x.first == nothing  && return
		print(io, x.first, ",")
		show(io, x.second)
	end

	function lextra(l::List, r::Array)
		if l.first == nothing
			return r
		else
			lextra(l.second, push!(r, l.first))
		end
	end

	function Base.values(l::List)
		res = []
		lextra(l, res)
	end

	function llength(l::List, r::Int=0)
		if l.first == nothing
			return r
		else
			llength(l.second, r+1)
		end
	end

	Base.keys( l::List ) = fieldnames( typeof(l) )
	Base.length(l::List) = llength(l)
	Base.copy(l::List) = List(l.first, l.second)

	function Base.pop!(l::List)
		l.first == nothing && return nothing
		tmp = l.second.first == nothing ? List() : copy(l.second)
		res = l.first
		l.first =  tmp.first
		l.second = tmp.second
		res
	end

	function Base.push!(l::List, a::Any)

		if l.first == nothing
		   l.first, l.second =  a, List()
		else
			push!(l.second, a)
		end
	end

	function Base.push!(a::Any, l::List )
		l.second = copy(l)
		l.first = a
	end

	function Base.push!(l::List, a::Any , test::Any)
		if l.second.first == test
		   l.second.second = List(a, l.second.second)
		else
			push!(l.second, a, test)
		end
	end

	include("Stack.jl")
	include("Queue.jl")
	include("tools.jl")

end

ext: jl

max_stars_repo_name: frankill/ftest

max_issues_repo_name: frankill/ftest

max_forks_repo_name: frankill/ftest


text:
# this should be run after running GovDebt.jl
# targeted moments:
varlinc_target  = 0.7  # var(log(annual income))
vardlinc_target = 0.23 # var(one year changes in log(annual income))
mpc_target      = 0.16
pc0_target      = 0.10 # fraction with zero bond holdings
# parameters: pLH, pHL, zlo, sH/sL. We normalize so s and z both have mean 1.

using Distributions
#using PyPlot
using Roots
using JLD
using Optim

function mollifier(z::AbstractFloat,ehi::AbstractFloat,elo::AbstractFloat)
    # mollifier function
    In = 0.443993816237631
    if z<ehi && z>elo
        temp = -1.0 + 2.0*(z-elo)/(ehi-elo)
        out = (2.0/(ehi-elo))*exp(-1.0/(1.0-temp*temp))/In
    else
        out = 0.0
    end

    return out
end

ni = 10000 # number of individuals in the sample we use for simulating skill process
uz = rand(ni,8) # uniform random variable used to get transitory shock z; need 8 quarters to get 1 year change
us = rand(ni,8) # uniform random variable used to get persistent markov chain shock s

# we will create a single z sample which has zlo=0, zhi = 2. We will then rescale it as we change zlo.
nz = 1000
zgrid = collect(linspace(0.,2.,nz))
zprob = zeros(nz)
for i=1:nz
	zprob[i] = 2*mollifier(zgrid[i],2.,0.)/nz
end
zprob = zprob/sum(zprob)
zcdf = cumsum(zprob)

zave = 0.5*zgrid[1:nz-1]+0.5*zgrid[2:nz]

function zsample(uz::Array{Float64,2}, zgrid::AbstractArray, zcdf::AbstractArray, ni::Int, nz::Int)
	zs = zeros(ni,8)
	for i=1:ni
		for t=1:8
			for iz=1:nz-1
				if uz[i,t]>zcdf[iz] && uz[i,t] <= zcdf[iz+1]
					zs[i,t] = zave[iz]
				end
			end
		end
	end
	return zs
end
tic()
zs = zsample(uz,zgrid,zcdf,ni,nz)
toc()

function ssample(us::Array{Float64,2}, P::Array{Float64,2}, πss::AbstractArray, ni::Int)
	shist = ones(Int,ni,8)
	for i=1:ni
		if us[i,1] > πss[1]
			shist[i,1] = 2
		end
		for t=2:8
			if us[i,t] > P[shist[i,t-1],1]
				shist[i,t] = 2
			end
		end
	end
	return shist
end

function ln_annual_inc(zhist::Array{Float64,2}, us::Array{Float64,2}, zlo::AbstractFloat, P::Array{Float64,2}, πss::AbstractArray, sgrid::AbstractArray, ni::Int)
  	s_inds = ssample(us,P,πss,ni)
	linc1 = zeros(ni)
	linc2 = zeros(ni)
	for i=1:ni
		inc1 = 0.
		for t=1:4
#            @show zhist[i, t]-1
#            @show zlo
#            @show sgrid[1]
#            @show sgrid[2]
#            @show s_inds[i,t]-1
			zshock = 1. + (1. - zlo)*(zhist[i,t]-1.)
			sshock = (sgrid[1] + (sgrid[2] - sgrid[1])*(s_inds[i,t] - 1.))
#            @show zshock
#            @show sshock
			inc1 += zshock*sshock
		end
		inc2 = 0.
		for t=5:8
			zshock = 1. + (1. - zlo)*(zhist[i,t]-1.)
			sshock = (sgrid[1] + (sgrid[2] - sgrid[1])*(s_inds[i,t] - 1.))
			inc2 += zshock*sshock
		end
		linc1[i] += log(inc1)
     	linc2[i] += log(inc2)
	end
	return (linc1, linc2)
end

function skill_moments(sH_by_sL::Real, zlo::Real, pLH::AbstractFloat,
                       pHL::AbstractFloat, zhist::Array{Float64,2},
                       us::Array{Float64,2}, ni::Int = 10000)
	πL  = pHL / (pLH+pHL)
	πss = [πL ; 1-πL]
	P   = [[1.0 -pLH pLH] ; [pHL 1.0-pHL]]
	slo = 1.0 / (πL + (1-πL) * sH_by_sL)
	shi = sH_by_sL * slo
	sgrid = [slo; shi]
	(linc1, linc2) = ln_annual_inc(zs, us, zlo, P, πss, sgrid, ni)
	return (var(linc1), var(linc2 - linc1))
end

tic()
(varlinc, vardlinc) = skill_moments(1.2, 0.5, 0.1, 0.1, zs, us)
toc()

function best_fit(pLH::AbstractFloat, pHL::AbstractFloat, varlinc_target::AbstractFloat,
                  vardlinc_target::AbstractFloat, sH_by_sL_grid::AbstractArray,
                  zlo_grid::AbstractArray, n1::Int, n2::Int, zhist::Array{Float64,2},
                  us::Array{Float64,2}, ni::Int = 10000)
	n = n1*n2
	dist = zeros(n)
	varlinc = zeros(n)
	vardlinc = zeros(n)
	for i1=1:n1
		for i2=1:n2
			i = n2*(i1-1)+i2
			(varlinc[i],vardlinc[i]) = skill_moments(sH_by_sL_grid[i1], zlo_grid[i2], pLH, pHL, zs, us)
			dist[i] = abs(varlinc[i] - varlinc_target) + abs(vardlinc[i] - vardlinc_target)
		end
	end
	(mindist,imin) = findmin(dist)
	i2 = mod(imin-1,n2)+1
	i1 = Int((imin-i2)/n2 + 1)
	return (sH_by_sL_grid[i1], zlo_grid[i2], varlinc[imin], vardlinc[imin])
end

np = 100
sH_by_sL_grid = linspace(3.,9.,np)
zlo_grid      = linspace(0.+eps(), 0.8-eps(), np)
# tic()
# varlinc_target  = 0.7
# vardlinc_target = 0.23
# (sH_by_sL, zlo, varlinc, vardlinc) = best_fit(0.02,0.03,varlinc_target,vardlinc_target,sH_by_sL_grid, zlo_grid,np,np,zs,us)
# toc()

# out = [varlinc varlinc_target vardlinc vardlinc_target sH_by_sL zlo]

"""
THIS ONE GETS CALLED
"""
function best_fit(pLH::AbstractFloat, pHL::AbstractFloat, varlinc_target::AbstractFloat, vardlinc_target::AbstractFloat, sH_by_sL_grid::AbstractArray, zlo_grid::AbstractArray,
				  n1::Int, n2::Int, zhist::Array{Float64,2}, us::Array{Float64,2}, T::Int = 8, ni::Int = 10000)
	n = n1*n2
	dist = zeros(n)
	varlinc = zeros(n)
	vardlinc = zeros(n)
	for i1=1:n1
		for i2=1:n2
			i = n2*(i1-1)+i2
			(varlinc[i], vardlinc[i]) = skill_moments(sH_by_sL_grid[i1], zlo_grid[i2], pLH, pHL, zs, us)
			dist[i] = abs(varlinc[i] - varlinc_target) + abs(vardlinc[i] - vardlinc_target)
		end
	end
	(mindist,imin) = findmin(dist)
	i2 = mod(imin-1,n2)+1
	i1 = Int((imin-i2)/n2 + 1)
	return (sH_by_sL_grid[i1], zlo_grid[i2], varlinc[imin], vardlinc[imin])
end

"""

"""
loss(x::Vector{Float64}, target::Vector{Float64}) = sum(abs.(x-target))

function best_fit(pLH::T, pHL::T, target::Vector{T}, lower::Vector{T},
                  upper::Vector{T}, us::Array{Float64,2}, max_iter::Int = 300,
                  initial_guess::Vector{Float64} = [6.3, 0.03]) where T<:AbstractFloat

    skill_moments_f(x) = loss(collect(skill_moments(x[1], x[2], pLH, pHL, zs, us)), target)

    res = optimize(skill_moments_f, lower, upper, initial_guess, Fminbox(NelderMead()),
                   Optim.Options(f_calls_limit = max_iter))

    sH_by_sL_argmin, zlo_argmin = Optim.minimizer(res)
    min_varlinc, min_vardlinc   = skill_moments(sH_by_sL_argmin, zlo_argmin, pLH, pHL, zs, us)

    return (sH_by_sL_argmin, zlo_argmin, min_varlinc, min_vardlinc)
end

function ave_mpc(m::AbstractArray, c::AbstractArray, agrid::AbstractArray, aswts::AbstractArray, na::Int, ns::Int)
	mpc = 0.
	for iss=1:ns
		i = na*(iss-1)+1
		mpc += aswts[i]*m[i]*(c[i+1] - c[i])/(agrid[2] - agrid[1])
		for ia=2:na
			i = na*(iss-1)+ia
			mpc += aswts[i]*m[i]*(c[i] - c[i-1])/(agrid[ia] - agrid[ia-1])
		end
	end
	return mpc
end

function frac_zero(m::AbstractArray, c::AbstractArray, agrid::AbstractArray, aswts::AbstractArray, ns::Int)
	return sum(aswts.*m.*(c.==repeat(agrid,ns)))
end

nLH = 2
nHL = 2

pLH_grid = linspace(0.01,0.1,nLH)
pHL_grid = linspace(0.01,0.1,nHL)
nn = nLH*nHL
sH_by_sL_grid2 = zeros(nn)
zlo_grid2 = zeros(nn)
varlinc_grid = zeros(nn)
vardlinc_grid = zeros(nn)
mpc_grid  = zeros(nn)
pc0_grid  = zeros(nn)

# RECA
sH_by_sL_lo = 3.0
sH_by_sL_hi = 9.0
zlo_lo      = 1e-18
zlo_hi      = 0.8-eps()

lower = [sH_by_sL_lo, zlo_lo]
upper = [sH_by_sL_hi, zlo_hi]
target = [varlinc_target, vardlinc_target]

for iLH = 1:nLH
	for iHL = 1:nHL
		i = nHL*(iLH-1)+iHL

        @time out = best_fit(pLH_grid[iLH], pHL_grid[iHL], target, lower, upper, us)
        @show "Old Method"
		@time (sH_by_sL_grid2[i], zlo_grid2[i], varlinc_grid[i], vardlinc_grid[i]) = best_fit(pLH_grid[iLH],pHL_grid[iHL],varlinc_target,vardlinc_target,sH_by_sL_grid, zlo_grid,np,np,zs,us)
		(f, sgrid, swts) = persistent_skill_process(sH_by_sL_grid2[i], pLH_grid[iLH], pHL_grid[iHL], ns)
		(agrid, awts) = cash_grid(sgrid, ω, H, r, η, γ, T, zlo, na)
		aswts = kron(swts, awts)
		qfunction(x) = mollifier(x, 2. - zlo_grid[i], zlo_grid[i])
		println("parameters:")
		println([sH_by_sL_grid2[i] zlo_grid2[i] pLH_grid[iLH] pHL_grid[iHL]])
		(ell, c, m, β, report) = findss(na, ns, βlo, βhi, R, ω, H, η, T, γ, bg, qfunction, agrid, sgrid,aswts, Win, f)
		mpc_grid[i] = ave_mpc(m,c,agrid,aswts,na,ns)
		pc0_grid[i] = frac_zero(m,c,agrid,aswts,ns)
	end
end


text:
"""
function to find steady state for MPO L given input guess A in left
canonical form (normalisation not necessary). Matrix as linear map.
"""

function steady_state_search!(At,L,Lconj,maxsweeps,converge_rat,eigs_tol)

    na = length(At)
    FL = Array{Array{Complex{Float64},4},1}(na)
    FR = Array{Array{Complex{Float64},4},1}(na)
    ae = 0.0
    aestart = 0.0

    FL[1] = ones(1,1,1,1)
    for jj = 1:na-1
        FL[jj+1] = update_lenv_two_op(At[jj],L[jj],Lconj[jj],At[jj],FL[jj])
    end
    FR[na] = ones(1,1,1,1)

    for sweep = 1:maxsweeps

        println("sweeping left")
        #sweep left
        for jj = na:-1:2
            HM = LinearMap(x->A_prod_L_OO_R_eigs(FL[jj],x,L[jj],Lconj[jj],FR[jj]),
                size(At[jj][:],1),Complex128; ishermitian=true)
            ae,av = eigs(HM; nev = 1, which = :SR, v0 = At[jj][:],
                maxiter = 1000, tol = eigs_tol)
            println("site ", jj, " evalue is ", ae[1])
            if jj == na
                aestart = ae[1]
            end
            #Q,R = qr(reshape(av,size(FL[jj],1),size(H[jj],2)*size(FR[jj],1))')
            #av = qrfact!(reshape(av,size(FL[jj],1),size(H[jj],2)*size(FR[jj],1))')
            #av = full(av[:Q])
            av,rv = qr(reshape(av,size(FL[jj],1),size(L[jj],2)*size(FR[jj],1))')
            At[jj] = reshape(av',size(av,2),size(At[jj],2),size(At[jj],3))
            FR[jj-1] = update_renv_two_op(At[jj],L[jj],Lconj[jj],At[jj],FR[jj])
            mo = At[jj-1]
            @tensor mo[-1,-2,-3] := mo[-1,-2,1]*rv[-3,1]
            At[jj-1] = mo
        end

        #HM = LinearMap(x->A_prod_LOR_eigs(FL[1],x,H[1],FR[1]),
        #       size(At[1][:],1),Complex128; ishermitian=true)
        #ae,av = eigs(HM; nev = 1, which = :SR, v0 = At[1][:],
        #        maxiter = 1000, tol = eigs_tol)
        #println("site ", 1, " evalue is ", ae[1])
        #At[1] = reshape(av,size(FL[1],1),size(H[1],2),size(FR[1],1))

        println("sweeping right")
        #sweep right
        for jj = 1:na-1;
            HM = LinearMap(x->A_prod_L_OO_R_eigs(FL[jj],x,L[jj],Lconj[jj],FR[jj]),
                size(At[jj][:],1),Complex128; ishermitian=true)
            ae,av = eigs(HM; nev = 1, which = :SR, v0 = At[jj][:],
                maxiter = 1000, tol = eigs_tol)
            println("site ", jj, " evalue is ", ae[1])
            #Q,R = qr(reshape(av,size(FL[jj],1)*size(H[jj],2),size(FR[jj],1)));
            #av = qrfact!(reshape(av,size(FL[jj],1)*size(H[jj],2),size(FR[jj],1)))
            #av = full(av[:Q])
            av, rv = qr(reshape(av,size(FL[jj],1)*size(L[jj],2),size(FR[jj],1)))
            At[jj] = reshape(av,size(At[jj],1),size(At[jj],2),size(av,2))
            FL[jj+1] = update_lenv_two_op(At[jj],L[jj],Lconj[jj],At[jj],FL[jj]);
            mo = At[jj+1]
            @tensor mo[-1,-2,-3] := rv[-1,1]*mo[1,-2,-3]
            At[jj+1] = mo
        end

        HM = LinearMap(x->A_prod_L_OO_R_eigs(FL[na],x,L[na],Lconj[na],FR[na]),
                size(At[na][:],1),Complex128; ishermitian=true)
        ae,av = eigs(HM; nev = 1, which = :SR, v0 = At[na][:],
                maxiter = 1000, tol = eigs_tol)
        println("site ", na, " evalue is ", ae[1])
        At[na] = reshape(av,size(FL[na],1),size(L[na],2),size(FR[na],1))

        if abs(ae[1]) > converge_rat*abs(aestart) || abs(ae[1])<1e-10
            return
        end

    end

    no = sum(abs2(At[na][:]))
    return ae

end

function steady_state_search_two_site!(At,L,Lconj,maxsweeps,converge_rat,eigs_tol,D_lim,svd_tol,Dims)

    na = length(At)
    FL = Array{Array{Complex{Float64},4},1}(na)
    FR = Array{Array{Complex{Float64},4},1}(na)
    ae = 0.0
    aestart = 0.0

    FL[1] = ones(1,1,1,1)
    for jj = 1:na-2
        FL[jj+1] = update_lenv_two_op(At[jj],L[jj],Lconj[jj],At[jj],FL[jj])
    end
    FR[na] = ones(1,1,1,1)

    for sweep = 1:maxsweeps

        println("sweeping left")
        #sweep left
        for jj = na:-1:2
            vin = neighbour_contract(At[jj-1],At[jj])
            HM = LinearMap(x->A_prod_L_OO_R_two_site(FL[jj-1],x,L[jj-1],Lconj[jj-1],
                L[jj],Lconj[jj],FR[jj]),
                size(vin,1),Complex128; ishermitian=true)
            ae,av = eigs(HM; nev = 1, ncv = 20, which = :SR, v0 = vin,
                maxiter = 2000, tol = eigs_tol)
            if jj == na
                aestart = ae[1]
            end
            U,S,Vd = svdecon(reshape(av,Dims[jj-1]*size(At[jj-1],2),
                size(At[jj],2)*Dims[jj+1]))
            S = S[S.>0]
            stot = sum(S.^2)
            Dims[jj] = min(sum(cumsum(S.^2/stot) .< 1-svd_tol)+1,D_lim)
            println("site ", jj, " evalue is ", ae[1], " bond dim ", Dims[jj])
            At[jj] = reshape(Vd[:,1:Dims[jj]]',Dims[jj],size(At[jj],2),Dims[jj+1])
            FR[jj-1] = update_renv_two_op(At[jj],L[jj],Lconj[jj],At[jj],FR[jj])
            At[jj-1] = reshape(U[:,1:Dims[jj]]*diagm(S[1:Dims[jj]]),Dims[jj-1],size(At[jj-1],2),Dims[jj])
        end

        #HM = LinearMap(x->A_prod_L_OO_R_eigs(FL[1],x,L[1],Lconj[1],FR[1]),
        #       size(At[1][:],1),Complex128; ishermitian=true)
        #ae,av = eigs(HM; nev = 1, which = :SR, v0 = At[1][:],
        #        maxiter = 2000, tol = eigs_tol)
        #println("site ", 1, " evalue is ", ae[1])
        #At[1] = reshape(av,size(FL[1],1),size(L[1],2),size(FR[1],1))

        println("sweeping right")
        #sweep right
        for jj = 1:na-1;
            vin = neighbour_contract(At[jj],At[jj+1])
            HM = LinearMap(x->A_prod_L_OO_R_two_site(FL[jj],x,L[jj],Lconj[jj],
                L[jj+1],Lconj[jj+1],FR[jj+1]),
                size(vin,1),Complex128; ishermitian=true)
            ae,av = eigs(HM; nev = 1, ncv = 20, which = :SR, v0 = vin,
                maxiter = 2000, tol = eigs_tol)
            U,S,Vd = svdecon(reshape(av,Dims[jj]*size(At[jj],2),
                size(At[jj+1],2)*Dims[jj+2]))
            S = S[S.>0]
            stot = sum(S.^2)
            Dims[jj+1] = min(sum(cumsum(S.^2/stot) .< 1-svd_tol)+1,D_lim)
            println("site ", jj, " evalue is ", ae[1], " bond dim ", Dims[jj+1])
            At[jj] = reshape(U[:,1:Dims[jj+1]],Dims[jj],size(At[jj],2),Dims[jj+1])
            At[jj+1] = reshape(diagm(S[1:Dims[jj+1]])*(Vd[:,1:Dims[jj+1]]'),Dims[jj+1],size(At[jj+1],2),Dims[jj+2])
            FL[jj+1] = update_lenv_two_op(At[jj],L[jj],Lconj[jj],At[jj],FL[jj]);
        end

        #HM = LinearMap(x->A_prod_L_OO_R_eigs(FL[na],x,L[na],Lconj[na],FR[na]),
        #        size(At[na][:],1),Complex128; ishermitian=true)
        #ae,av = eigs(HM; nev = 1, which = :SR, v0 = At[na][:],
        #        maxiter = 2000, tol = eigs_tol)
        #println("site ", na, " evalue is ", ae[1])
        #At[na] = reshape(av,size(FL[na],1),size(L[na],2),size(FR[na],1))

        if abs(ae[1]) > converge_rat*abs(aestart) || abs(ae[1])<1e-10
            return
        end

    end

    no = sum(abs2(At[na][:]))
    return ae, Dims

end


text:
function fastmath_variants(pairs...)
    result = []
    for (func_name, func_def) in pairs
        push!(result, eval(parse(func_name)) => func_def)
        try # not all functions have a fast variant
            fast = eval(parse("$(func_name)_fast"))
            push!(result, fast => func_def)
        end
    end
    result
end

function openlibm_variants(pairs...)
    result = []
    for (func_name, func_def) in pairs
        push!(result, (Symbol("$(func_name)"), "libopenlibm") => func_def)
        push!(result, (Symbol("$(func_name)f"), "libopenlibm") => func_def)
    end
    result
end

"""
Helper to get around the extraneous checks that Julia performs
when multiplying bools to floats.
"""
@noinline function case(test::Bool, a::Number)
    Base.select_value(test, a, zero(a))
end

basefuncs = Dict(

    Base.FastMath.box => (typ, value) -> "$value",
    Base.box => (typ, value) -> "$value",
    Base.sitofp => (typ, value) -> "(($typ)$value)",
    Base.fptosi => (typ, value) -> "(($typ)$value)",
    Base.uitofp => (typ, value) -> "(($typ)$value)",
    Base.fptoui => (typ, value) -> "(($typ)$value)",
    Base.checked_trunc_sint => (typ, value) -> "(($typ)$value)",
    Base.fpext => (typ, value) -> "(($typ)$value)",
    Base.fptrunc => (typ, value) -> "(($typ)$value)",
    Base.sext_int => (typ, value) -> "(($typ)$value)",
    Base.convert => (typ, value) -> "(($typ)$value)",
    Base.typeassert => (value, typ) -> "$value",

    Base.round => (value) -> "round($value)",
    Base.floor => (value) -> "floor($value)",
    Base.ceil => (value) -> "ceil($value)",
    Base.clamp => (val,min,max) -> "clamp($val,$min,$max)",

    ISPC.case => (test, a) -> "((float)$test * $a)",
    Core.Intrinsics.select_value => (test, a, b) -> "($test ? $a : $b)",

    (===) => (a, b) -> "($a == $b)",

    (*) => (a, b) -> "($a * $b)",
    (-) => (a, b) -> "($a - $b)",
    (+) => (a, b) -> "($a + $b)",
    (<) => (a, b) -> "($a < $b)",
    (!) => (a) -> "(!$a)",

    Base.FastMath.sub_fast => (a,b) -> "($a - $b)",
    Base.FastMath.add_fast => (a,b) -> "($a + $b)",
    Base.FastMath.mul_fast => (a,b) -> "($a * $b)",
    Base.FastMath.div_fast => (a,b) -> "($a / $b)",

    fastmath_variants(
        "Base.min" => (a) -> "min($a)",
        "Base.max" => (a) -> "max($a)",

        "Base.not_int" => (a) -> "__not($a)",
        "Base.or_int" => (a,b) -> "($a || $b)",
        "Base.xor_int" => (a,b) -> "($a ^ $b)",
        "Base.and_int" => (a,b) -> "($a && $b)",

        "Base.add_int" => (a,b) -> "($a + $b)",
        "Base.sub_int" => (a,b) -> "($a - $b)",
        "Base.mul_int" => (a,b) -> "($a * $b)",
        "Base.sdiv_int" => (a,b) -> "($a / $b)",
        "Base.udiv_int" => (a,b) -> "($a / $b)",

        "Base.eq_int" => (a,b) -> "($a == $b)",
        "Base.slt_int" => (a,b) -> "($a < $b)",
        "Base.sle_int" => (a,b) -> "($a <= $b)",
        "Base.ult_int" => (a,b) -> "($a < $b)",
        "Base.ule_int" => (a,b) -> "($a <= $b)",

        "Base.eq_float" => (a,b) -> "($a == $b)",
        "Base.ne_float" => (a,b) -> "($a != $b)",
        "Base.lt_float" => (a,b) -> "($a < $b)",
        "Base.le_float" => (a,b) -> "($a <= $b)",
        "Base.neg_float" => (a) -> "(-$a)",
        "Base.abs_float" => (a) -> "abs($a)",
        "Base.add_float" => (a,b) -> "($a + $b)",
        "Base.sub_float" => (a,b) -> "($a - $b)",
        "Base.mul_float" => (a,b) -> "($a * $b)",
        "Base.div_float" => (a,b) -> "($a / $b)",

        "Base.powi_llvm" => (a,b) -> "pow($a, $b)",
        "Base.floor_llvm" => (a) -> "floor($a)",
        "Base.sqrt_llvm" => (a) -> "sqrt($a)",
        "Base.sqrt" => (a) -> "sqrt($a)",
        "Base.inv" => (a) -> "rcp($a)",

        "Base.nan_dom_err" => (a, b) -> "$a",

    )...,

    openlibm_variants(
        "sin" => (a) -> "sin($a)",
        "cos" => (a) -> "cos($a)",
        "tan" => (a) -> "tan($a)",
        "asin" => (a) -> "sin($a)",
        "acos" => (a) -> "cos($a)",
        "atan" => (a) -> "tan($a)",
        "exp" => (a) -> "exp($a)",
        "log" => (a) -> "log($a)",
        "pow" => (a,b) -> "pow($a,$b)",
        "sqrt" => (a) -> "sqrt($a)",
    )...
)


text:
"""
**Heteroscedastic Likelihood**

Gaussian with heteroscedastic noise given by another gp: ``p(y|f,g) = \\mathcal{N}(y|f,(\\lambda\\sigma(g))^{-1})``

```julia
HeteroscedasticLikelihood([kernel=RBFKernel(),[priormean=0.0]])
```
Augmentation is described here (#TODO)
"""
struct HeteroscedasticLikelihood{T<:Real} <: RegressionLikelihood{T}
    kernel::LatentArray{Kernel{T}}
    μ₀::LatentArray{PriorMean{T}}
    λ::LatentArray{T}
    c::LatentArray{Vector{T}}
    ϕ::LatentArray{Vector{T}}
    γ::LatentArray{Vector{T}}
    θ::LatentArray{Vector{T}}
    μ::LatentArray{Vector{T}}
    Σ::LatentArray{Symmetric{T,Matrix{T}}}
    σg::LatentArray{Vector{T}}
    K::LatentArray{Symmetric{T,Matrix{T}}}
    invK::LatentArray{Symmetric{T,Matrix{T}}}
    κ::LatentArray{Matrix{T}}
    K̃::LatentArray{Vector{T}}
    function HeteroscedasticLikelihood{T}(kernel::Vector{<:Kernel{T}},priormean::Vector{<:PriorMean}) where {T<:Real}
            new{T}(kernel,priormean)
    end
    function HeteroscedasticLikelihood{T}(kernel::AbstractVector{<:Kernel},μ₀::AbstractVector{<:PriorMean},λ::AbstractVector{T},c::AbstractVector{<:AbstractVector{T}},ϕ::AbstractVector{<:AbstractVector{T}},γ::AbstractVector{<:AbstractVector{T}},θ::AbstractVector{<:AbstractVector{T}},μ::AbstractVector{<:AbstractVector{T}},Σ::AbstractVector{<:Symmetric{T,Matrix{T}}},σg::AbstractVector{<:AbstractVector{T}},K::AbstractVector{<:Symmetric{T,Matrix{T}}},invK::AbstractVector{<:Symmetric{T,Matrix{T}}}) where {T<:Real}
         new{T}(kernel,μ₀,λ,c,ϕ,γ,θ,μ,Σ,σg,K,invK)
     end
     function HeteroscedasticLikelihood{T}(kernel::AbstractVector{<:Kernel},μ₀::AbstractVector{<:PriorMean},λ::AbstractVector{T},c::AbstractVector{<:AbstractVector{T}},ϕ::AbstractVector{<:AbstractVector{T}},γ::AbstractVector{<:AbstractVector{T}},θ::AbstractVector{<:AbstractVector{T}},μ::AbstractVector{<:AbstractVector{T}},Σ::AbstractVector{<:Symmetric{T,Matrix{T}}},σg::AbstractVector{<:AbstractVector{T}},K::AbstractVector{<:Symmetric{T,Matrix{T}}},invK::AbstractVector{<:Symmetric{T,Matrix{T}}},κ::AbstractVector{<:AbstractMatrix{T}},K̃::AbstractVector{<:AbstractVector{T}}) where {T<:Real}
          new{T}(kernel,μ₀,λ,c,ϕ,γ,θ,μ,Σ,σg,K,invK,κ,K̃)
      end
end

function HeteroscedasticLikelihood(kernel::Union{Kernel{T},AbstractVector{<:Kernel{T}}}=RBFKernel(),priormean::PriorMean{T}=ConstantMean(0.0)) where {T<:Real}
    if typeof(kernel) <: AbstractVector
        HeteroscedasticLikelihood{T}(kernel,[priormean])
    else
        HeteroscedasticLikelihood{T}([kernel],[priormean])
    end
end

function pdf(l::HeteroscedasticLikelihood,y::Real,f::Real,g::Real)
    pdf(Normal(y,inv(sqrt(l.λ[1]*logistic(g)))),f) #WARNING multioutput invalid
end

function logpdf(l::HeteroscedasticLikelihood,y::Real,f::Real,g::Real)
    logpdf(Normal(y,inv(sqrt(l.λ[1]*logistic(g)))),f) #WARNING multioutput invalid
end

function Base.show(io::IO,model::HeteroscedasticLikelihood{T}) where T
    print(io,"Gaussian likelihood with heteroscedastic noise")
end

function init_likelihood(likelihood::HeteroscedasticLikelihood{T},inference::Inference{T},nLatent::Integer,nSamplesUsed::Integer,nFeatures::Integer) where {T<:Real}
    kernel = Vector{Kernel}(undef,nLatent)
    if length(likelihood.kernel) != nLatent
        kernel[:] = [deepcopy(likelihood.kernel[1]) for _ in 1:nLatent]
    else
        kernel[:] = likelihood.kernel
    end
    μ₀ = [deepcopy(likelihood.μ₀[1]) for _ in 1:nLatent]
    λ = ones(T,nLatent)
    c = [ones(T,nSamplesUsed) for _ in 1:nLatent]
    ϕ = [ones(T,nSamplesUsed) for _ in 1:nLatent]
    γ = [ones(T,nSamplesUsed) for _ in 1:nLatent]
    θ = [ones(T,nSamplesUsed) for _ in 1:nLatent]
    μ = [zeros(T,nSamplesUsed) for _ in 1:nLatent]
    Σ = [Symmetric(Matrix(Diagonal(one(T)*I,nFeatures))) for _ in 1:nLatent]
    σg = [ones(T,nSamplesUsed) for _ in 1:nLatent]
    K = [Symmetric(Matrix(Diagonal(one(T)*I,nFeatures))) for _ in 1:nLatent]
    invK = [Symmetric(Matrix(Diagonal(one(T)*I,nFeatures))) for _ in 1:nLatent]
    κ = [ones(T,nSamplesUsed,nFeatures) for _ in 1:nLatent]
    K̃ = [ones(T,nSamplesUsed) for _ in 1:nLatent]
    HeteroscedasticLikelihood{T}(kernel,μ₀,λ,c,ϕ,γ,θ,μ,Σ,σg,K,invK,κ,K̃)
end

function local_updates!(model::VGP{T,HeteroscedasticLikelihood{T}}) where {T}
    model.likelihood.ϕ .= broadcast((μ,Σ,y)->0.5*(abs2.(μ-y)+Σ),model.μ,diag.(model.Σ),model.y)
    model.likelihood.c .= broadcast((μ,Σ)->sqrt.(Σ+abs2.(μ)),model.likelihood.μ,diag.(model.likelihood.Σ))
    model.likelihood.γ .= broadcast((λ,ϕ,μ,c)->0.5*λ*ϕ.*safe_expcosh.(-0.5*μ,0.5*c),model.likelihood.λ,model.likelihood.ϕ,model.likelihood.μ,model.likelihood.c)
    model.likelihood.θ .= broadcast((γ,c)->0.5*(0.5.+γ)./c.*tanh.(0.5*c),model.likelihood.γ,model.likelihood.c)
    model.likelihood.K .= broadcast(kernel->Symmetric(kernelmatrix(model.X,kernel)+getvariance(kernel)*T(jitter)*I),model.likelihood.kernel)
    model.likelihood.invK .= inv.(model.likelihood.K)
    model.likelihood.Σ .= broadcast((θ,invK)->Symmetric(inv(Diagonal(θ)+invK)),model.likelihood.θ,model.likelihood.invK)
    model.likelihood.μ .= broadcast((Σ,invK,μ₀,γ)->Σ*(invK*μ₀+0.5*(0.5.-γ)),model.likelihood.Σ,model.likelihood.invK,model.likelihood.μ₀,model.likelihood.γ)
    model.likelihood.σg .=  broadcast((μ,Σ)->expectation.(logistic,Normal.(μ,sqrt.(diag(Σ)))),model.likelihood.μ,model.likelihood.Σ)
    model.likelihood.λ .= 0.5*broadcast((ϕ,σg)->model.inference.nSamples/dot(ϕ,σg),model.likelihood.ϕ,model.likelihood.σg)
end

function local_autotuning!(model::VGP{T,<:HeteroscedasticLikelihood}) where {T}
    Jnn = kernelderivativematrix.([model.X],model.likelihood.kernel)
    f_l,f_v,f_μ₀ = hyperparameter_local_gradient_function(model)
    grads_l = map(compute_hyperparameter_gradient,model.likelihood.kernel,fill(f_l,model.nLatent),Jnn,1:model.nLatent)
    grads_v = map(f_v,model.likelihood.kernel,1:model.nPrior)
    grads_μ₀ = map(f_μ₀,1:model.nLatent)

    apply_gradients_lengthscale!.(model.likelihood.kernel,grads_l) #Send the derivative of the matrix to the specific gradient of the model
    apply_gradients_variance!.(model.likelihood.kernel,grads_v) #Send the derivative of the matrix to the specific gradient of the model
    update!.(model.likelihood.μ₀,grads_μ₀)

    model.inference.HyperParametersUpdated = true
end

function local_updates!(model::SVGP{T,HeteroscedasticLikelihood{T}}) where {T}
    model.likelihood.ϕ .= broadcast((κ,μ,Σ,y,K̃)->0.5*(abs2.(κ*μ-y)+opt_diag(κ*Σ,κ)+K̃),model.κ,model.μ,diag.(model.Σ),model.inference.y,model.K̃)
    model.likelihood.c .= broadcast((κ,μ,Σ,K̃)->sqrt.(opt_diag(κ*Σ,κ)+abs2.(κ*μ)+K̃),model.likelihood.κ,model.likelihood.μ,diag.(model.likelihood.Σ),model.likelihood.K̃)
    model.likelihood.γ .= broadcast((λ,ϕ,κμ,c)->0.5*λ*ϕ.*safe_expcosh.(-0.5*μ,0.5*c),model.likelihood.λ,model.likelihood.ϕ,model.likelihood.κ.*model.likelihood.μ,model.likelihood.c)
    model.likelihood.θ .= broadcast((γ,c)->0.5*(0.5.+γ)./c.*tanh.(0.5*c),model.likelihood.γ,model.likelihood.c)
    model.likelihood.K .= broadcast(kernel->Symmetric(kernelmatrix(model.Z,kernel)+getvariance(kernel)*T(jitter)*I),model.likelihood.kernel)
    model.likelihood.invK .= inv.(model.likelihood.K)
    model.likelihood.Knm .= broadcast(kernel->kernelmatrix(model.inference.x,model.Z,kernel),model.likelihood.kernel)
    model.likelihood.κ .= model.likelihood.Knm.*model.invK
    model.likelihood.K̃ .=  broadcast(kernel->kerneldiagmatrix(model.inference.x,kernel),model.likelihood.kernel)
    model.likelihood.Σ .= broadcast((θ,κ,invK)->Symmetric(inv(model.inference.ρ*transpose(κ)*Diagonal(θ)*κ+invK)),model.likelihood.θ,model.likelihood.κ,model.likelihood.invK)
    model.likelihood.μ .= broadcast((Σ,invK,μ₀,γ)->Σ*(invK*μ₀+0.5*transpose(κ)*(0.5.-γ)),model.likelihood.Σ,model.likelihood.invK,model.likelihood.μ₀,model.likelihood.γ)
    model.likelihood.σg .=  broadcast((μ,Σ)->expectation.(logistic,Normal.(μ,sqrt.(diag(Σ)))),model.likelihood.μ,model.likelihood.Σ)
    model.likelihood.λ .= 0.5*broadcast((ϕ,σg)->model.inference.nSamples/dot(ϕ,σg),model.likelihood.ϕ,model.likelihood.σg)
end

@inline ∇E_μ(model::AbstractGP{T,<:HeteroscedasticLikelihood,<:AnalyticVI}) where {T} = 0.5*hadamard.(model.inference.y,model.likelihood.λ.*model.likelihood.σg)
@inline ∇E_μ(model::AbstractGP{T,<:HeteroscedasticLikelihood,<:AnalyticVI},i::Int) where {T} = 0.5*model.likelihood.λ[i]*model.inference.y[i].*model.likelihood.σg[i]

@inline ∇E_Σ(model::AbstractGP{T,<:HeteroscedasticLikelihood,<:AnalyticVI}) where {T} = 0.5*model.likelihood.λ.*model.likelihood.σg
@inline ∇E_Σ(model::AbstractGP{T,<:HeteroscedasticLikelihood,<:AnalyticVI},i::Int) where {T} = 0.5*model.likelihood.λ[i].*model.likelihood.σg[i]

function proba_y(model::VGP{T,HeteroscedasticLikelihood{T},AnalyticVI{T}},X_test::AbstractMatrix{T}) where {T<:Real}
    μf, σ²f = predict_f(model,X_test,covf=true)
    μg, σ²g = _predict_f.(model.likelihood.μ,model.likelihood.Σ,model.likelihood.invK,model.likelihood.kernel,[X_test],[model.X],covf=true)[1]#WARNING Only valid for 1D output
    return μf,σ²f.+broadcast((λ,μ,σ)->expectation.(x->inv(λ*logistic(x)),Normal.(μ,sqrt.(σ))),model.likelihood.λ,μg,σ²g)
end

function proba_y(model::SVGP{T,HeteroscedasticLikelihood{T},AnalyticVI{T}},X_test::AbstractMatrix{T}) where {T}
    μf, σ²f = predict_f(model,X_test,covf=true)
    μg, σ²g = _predict_f.(model.likelihood.μ,model.likelihood.Σ,model.likelihood.invK,model.likelihood.kernel,[X_test],model.Z,covf=true)

    return μf,σ²f.+broadcast((λ,μ,σ)->dot.([pred_weights],eachcol(inv.(λ*logistic.(pred_nodes.*sqrt2*sqrt.(σ)'.+μ')))),model.likelihood.λ,μg,σ²g)
end

function ELBO(model::AbstractGP{T,HeteroscedasticLikelihood{T},<:AnalyticVI}) where {T}
    return expecLogLikelihood(model) - GaussianKL(model) - GaussianKL_g(model) - PoissonKL(model) - PolyaGammaKL(model)
end

function expecLogLikelihood(model::VGP{T,HeteroscedasticLikelihood{T}}) where {T}
    tot = model.nFeatures*(0.5*sum(log,model.likelihood.λ)-model.nLatent*(log(2*sqrt(twoπ))))
    tot += 0.5*sum(broadcast((μ,γ,Σ,θ)->dot(μ,(0.5 .- γ)) - dot(abs2.(μ),θ)-dot(Σ,θ),model.likelihood.μ,model.likelihood.γ,diag.(model.likelihood.Σ),model.likelihood.θ))
    return tot
end

function GaussianKL_g(model::AbstractGP{T,<:HeteroscedasticLikelihood}) where {T}
    sum(broadcast(GaussianKL,model.likelihood.μ,model.likelihood.μ₀,model.likelihood.Σ,model.likelihood.invK))
end

function PoissonKL(model::VGP{T,<:HeteroscedasticLikelihood}) where {T}
    return sum(broadcast(PoissonKL,model.likelihood.γ,
            broadcast((λ,y,μ,Σ)->0.5*λ*(abs2.(y-μ)+Σ),model.likelihood.λ,model.y,model.μ,diag.(model.Σ)),
            broadcast((λ,y,μ,Σ)->log.(0.5*λ*(abs2.(μ-y)+Σ)),model.likelihood.λ,model.y,model.μ,diag.(model.Σ)))) #TODO
end

function PolyaGammaKL(model::VGP{T,<:HeteroscedasticLikelihood}) where {T}
    sum(broadcast(PolyaGammaKL,broadcast(γ->0.5.+γ,model.likelihood.γ),model.likelihood.c,model.likelihood.θ))
end


text:
using Plots

@test TO.pos(10) == 10
@test TO.pos(-10) == 0
@test TO.pos(-1e-4) == 0

# Conversion between Matrices and Trajectories
n,m,N = 5,3,21
X = rand(n,N)
X_ = [X[:,k] for k = 1:N]
@test to_array(X_) == X
@test TO.to_trajectory(X) == X_
@test to_dvecs(X) == X_
@test vec(X_) == vec(X)

Z = rand(n,m,N)
Z_ = [Z[:,:,k] for k = 1:N]
@test to_array(Z_) == Z
@test to_dvecs(Z) == Z_

D_ = [Diagonal(rand(n)) for k = 1:N]
D = zeros(n,n,N)
for k = 1:N
    D[:,:,k] = D_[k]
end
@test to_array(D_) == D

# Test copy methods
X_ = [rand(n) for k = 1:N]
X2_ = [rand(n) for k = 1:N]
copyto!(X2_, X_)
@test X2_[1] == X_[1]
X_[1][1] = 100
@test X2_[1][1] != 100

X2_
copyto!(X2_, X)
@test X2_[1] == X[:,1]

copyto!(X, X_)
@test X[1,1] == 100

# Vector of different lengths
x,y,z = rand(n), rand(n), rand(m)
W = [x;y;z]
p = [n,n,m]
@test to_dvecs(W,p) == [x,y,z]

# Others
@test isdir(TO.root_dir())

# Plotting
function test_plot()
    plot()
    plt = TO.plot_circle!([1,1,5], 1, color=:black)

    circles = [[-1,1.5,1], [-1,-1,0.5], [1,-1,0.25]]
    TO.plot_obstacles(circles)

    display(plt)

    time = range(0,10,length=N)
    X = [[sin(t); cos(t); rand(3)] for t in time]
    plot_trajectory!(X)
    plot_trajectory!(to_dvecs(X).*2)

    plot_vertical_lines!(plt, [-1,0,1])

    X_ = to_dvecs(X)
    plot(X_)
    plot!(X_, color=:black)
    plot(X_,1:2)
    return nothing
end
test_plot()

# Constraints
@test TO.circle_constraint([0,0,0], 1, 0, 1) == 0
@test TO.circle_constraint([0,0,0], 1, 0, 0.5) < 0
@test TO.circle_constraint([0.75,0,0], 1, 0, 0.5) > 0
@test TO.circle_constraint([0,0,0], [1, 0], 1) == 0
@test TO.circle_constraint([0,0,0], [1, 0], 0.5) < 0
@test TO.circle_constraint([0.75,0,0], [1, 0], 0.5) > 0

@test TO.sphere_constraint([0,0,0], 1, 0, 0, 1) == 0
@test TO.sphere_constraint([0,0,0], 1, 0, 0, 0.5) < 0
@test TO.sphere_constraint([0.75,0,0], 1, 0, 0, 0.5) > 0
@test TO.sphere_constraint([0,0,0], [1, 0, 0], 1) == 0
@test TO.sphere_constraint([0,0,0], [1, 0, 0], 0.5) < 0
@test TO.sphere_constraint([0.75,0,0], [1, 0, 0], 0.5) > 0

max_issues_repo_name: MehrdadZareian/GMARA

max_forks_repo_name: MehrdadZareian/GMARA


text:
function solve()
    n = parse(Int, readline())
    as = [parse(Int, x) for x in split(readline())]

    ret = 0
    for l=1:n
        x = as[l]
        for r=l:n
            x = min(x, as[r])
            ret = max(ret, x*(r-l+1))
        end
    end
    ret
end

println(solve())

max_stars_repo_name: aishikawa/atcoder-julia

max_issues_repo_name: aishikawa/atcoder-julia

max_forks_repo_name: aishikawa/atcoder-julia


text:
#!/home/kfischer/julia/julia
using Gallium
using Arsenic
pid = parse(Int,ARGS[1])
sess, modules = Gallium.Ptrace.attach(pid)
if !isdefined(Base, :active_repl)
    term = Base.Terminals.TTYTerminal(get(ENV, "TERM", @static is_windows() ? "" : "dumb"), STDIN, STDOUT, STDERR)
    active_repl = Base.REPL.LineEditREPL(term, true)
    eval(Base,:(active_repl = $active_repl))
    eval(Base,:(have_color = true))
    eval(Base,:(is_interactive = true))
end
ASTInterpreter.RunDebugREPL(Arsenic.compute_stack(modules, sess))
exit(0)


text:
using Classes
using DataStructures

"""
    @or(args...)

Return `a` if a !== nothing, else return `b`. Evaluates each expression
at most once.
"""
macro or(a, b)
    # @info "or($a, $b)"
    tmp = gensym(:tmp)
    expr = quote
        $tmp = $a
        ($tmp === nothing ? $b : $tmp)
    end
    esc(expr)
end

# Having all our structs/classes subtype these simplifies "show" methods
abstract type MimiStruct end
@class MimiClass <: Class

const AbstractMimiType = Union{MimiStruct, AbstractMimiClass}

# To identify components, @defcomp creates a variable with the name of
# the component whose value is an instance of this type.
struct ComponentId <: MimiStruct
    module_obj::Union{Nothing, Module}
    comp_name::Symbol
end

# Modules cannot be deepcopied, thus the override
Base.deepcopy_internal(x::ComponentId, dict::IdDict) = ComponentId(x.module_obj, x.comp_name)

# ComponentPath identifies the path through multiple composites to a leaf comp.
struct ComponentPath <: MimiStruct
    names::NTuple{N, Symbol} where N
end

ComponentPath(names::Vector{Symbol}) = ComponentPath(Tuple(names))
ComponentPath(names::Vararg{Symbol}) = ComponentPath(Tuple(names))

ComponentPath(path::ComponentPath, name::Symbol) = ComponentPath(path.names..., name)

ComponentPath(path1::ComponentPath, path2::ComponentPath) = ComponentPath(path1.names..., path2.names...)

ComponentPath(::Nothing, name::Symbol) = ComponentPath(name)

const ParamPath = Tuple{ComponentPath, Symbol}

#
# Dimensions
#

abstract type AbstractDimension <: MimiStruct end

const DimensionKeyTypes   = Union{AbstractString, Symbol, Int, Float64}
const DimensionRangeTypes = Union{UnitRange{Int}, StepRange{Int, Int}}

struct Dimension{T <: DimensionKeyTypes} <: AbstractDimension
    dict::OrderedDict{T, Int}

    function Dimension(keys::Vector{T}) where {T <: DimensionKeyTypes}
        dict = OrderedDict(collect(zip(keys, 1:length(keys))))
        return new{T}(dict)
    end

    function Dimension(rng::T) where {T <: DimensionRangeTypes}
        return Dimension(collect(rng))
    end

    Dimension(i::Int) = Dimension(1:i)

    # Support Dimension(:foo, :bar, :baz)
    function Dimension(keys::T...) where {T <: DimensionKeyTypes}
        vector = [key for key in keys]
        return Dimension(vector)
    end
end

#
# Simple optimization for ranges since indices are computable.
# Unclear whether this is really any better than simply using
# a dict for all cases. Might scrap this in the end.
#
mutable struct RangeDimension{T <: DimensionRangeTypes} <: AbstractDimension
    range::T
 end


text:
# An implementation of the Lanczos version of the conjugate gradient method
# for a family of shifted systems of the form (A + αI) x = b.
#
# The implementation follows
# A. Frommer and P. Maass, Fast CG-Based Methods for Tikhonov-Phillips
# Regularization, SIAM Journal on Scientific Computing, 20(5),
# pp. 1831-1850, 1999.
#
# Dominique Orban, <dominique.orban@gerad.ca>
# Princeton, NJ, March 2015.

export cg_lanczos, cg_lanczos_shift_seq

"""The Lanczos version of the conjugate gradient method to solve the
symmetric linear system

  Ax = b

The method does _not_ abort if A is not definite.
"""
function cg_lanczos(A :: AbstractLinearOperator, b :: AbstractVector{T};
                    atol :: Float64=1.0e-8, rtol :: Float64=1.0e-6, itmax :: Int=0,
                    check_curvature :: Bool=false, verbose :: Bool=false) where T <: Number

  n = size(b, 1);
  (size(A, 1) == n & size(A, 2) == n) || error("Inconsistent problem size");
  verbose && @printf("CG Lanczos: system of %d equations in %d variables\n", n, n);

  # Initial state.
  x = zeros(T, n);
  β = @knrm2(n, b)
  β == 0 && return x, LanczosStats(true, [0.0], false, 0.0, 0.0, "x = 0 is a zero-residual solution")
  v = b / β;
  v_prev = copy(v);
  p = copy(b);
  iter = 0;
  itmax == 0 && (itmax = 2 * n);

  # Initialize some constants used in recursions below.
  σ = β;
  ω = 0;
  γ = 1;
  Anorm2 = 0.0;
  β_prev = 0.0;

  # Define stopping tolerance.
  rNorm = σ;
  rNorms = [rNorm;];
  ε = atol + rtol * rNorm;
  verbose && @printf("%5d  %8.1e\n", iter, rNorm);

  indefinite = false;
  solved = rNorm <= ε;
  tired = iter >= itmax;
  status = "unknown";

  # Main loop.
  while ! (solved || tired || (check_curvature & indefinite))
    # Form next Lanczos vector.
    v_next = A * v;
    δ = @kdot(n, v, v_next);  # BLAS.dot(n, v, 1, v_next, 1) doesn't seem to pay off.

    # Check curvature. Exit fast if requested.
    # It is possible to show that σⱼ² (δⱼ - ωⱼ₋₁ / γⱼ₋₁) = pⱼᵀ A pⱼ.
    γ = 1 / (δ - ω / γ);
    indefinite |= (γ <= 0.0);
    (check_curvature & indefinite) && continue;

    @kaxpy!(n, -δ, v, v_next)  # Faster than v_next = Av - δ * v;
    if iter > 0
      @kaxpy!(n, -β, v_prev, v_next)  # Faster than v_next = v_next - β * v_prev;
      @. v_prev .= v
    end
    β = @knrm2(n, v_next)
    @. v .= v_next / β
    Anorm2 += β_prev^2 + β^2 + δ^2;  # Use ‖T‖ as increasing approximation of ‖A‖.
    β_prev = β;

    # Compute next CG iterate.
    @kaxpy!(n, γ, p, x)  # Faster than x = x + γ * p;

    ω = β * γ;
    σ = -ω * σ;
    ω = ω * ω;
    @kscal!(n, ω, p)
    @kaxpy!(n, σ, v, p)  # Faster than p = σ * v + ω * p;
    rNorm = abs(σ);
    push!(rNorms, rNorm);
    iter = iter + 1;
    verbose && @printf("%5d  %8.1e\n", iter, rNorm);
    solved = rNorm <= ε;
    tired = iter >= itmax;
  end

  status = tired ? "maximum number of iterations exceeded" : (check_curvature & indefinite) ? "negative curvature" : "solution good enough given atol and rtol"
  stats = LanczosStats(solved, rNorms, indefinite, sqrt(Anorm2), 0.0, status);  # TODO: Estimate Acond.
  return (x, stats);
end

"""The Lanczos version of the conjugate gradient method to solve a family
of shifted systems

  (A + αI) x = b  (α = α₁, α₂, ...)

The method does _not_ abort if A + αI is not definite.
"""
function cg_lanczos_shift_seq(A :: AbstractLinearOperator, b :: AbstractVector{Tb}, shifts :: AbstractVector{Ts};
                              atol :: Float64=1.0e-8, rtol :: Float64=1.0e-6, itmax :: Int=0,
                              check_curvature :: Bool=false, verbose :: Bool=false) where {Tb, Ts <: Number}

  n = size(b, 1);
  (size(A, 1) == n & size(A, 2) == n) || error("Inconsistent problem size");

  nshifts = size(shifts, 1);
  verbose && @printf("CG Lanczos: system of %d equations in %d variables with %d shifts\n", n, n, nshifts);

  # Initial state.
  ## Distribute x similarly to shifts.
  x = zeros(Tb, n, nshifts);
  β = @knrm2(n, b)
  β == 0 && return x, LanczosStats(true, [0.0], false, 0.0, 0.0, "x = 0 is a zero-residual solution")
  v = b / β;
  v_prev = copy(v);

  # Initialize each p to b.
  p = b * ones(nshifts)';

  # Initialize some constants used in recursions below.
  σ = β * ones(nshifts);
  δhat = zeros(nshifts);
  ω = zeros(Tb, nshifts);
  γ = ones(Tb, nshifts);

  # Define stopping tolerance.
  rNorms = β * ones(Tb, nshifts);
  rNorms_history = [rNorms;];
  ε = atol + rtol * β;

  # Keep track of shifted systems that have converged.
  converged = rNorms .<= ε;
  iter = 0;
  itmax == 0 && (itmax = 2 * n);

  # Keep track of shifted systems with negative curvature if required.
  indefinite = falses(nshifts);

  # Build format strings for printing.
  if verbose
    fmt = "%5d" * repeat("  %8.1e", nshifts) * "\n";
    # precompile printf for our particular format
    local_printf(data...) = Core.eval(Main, :(@printf($fmt, $(data)...)))
    local_printf(iter, rNorms...)
  end

  solved = all(converged);
  tired = iter >= itmax;
  status = "unknown";

  # Main loop.
  while ! (solved || tired)
    # Form next Lanczos vector.
    v_next = A * v;
    δ = @kdot(n, v, v_next)
    @kaxpy!(n, -δ, v, v_next)  # Faster than v_next = Av - δ * v;
    if iter > 0
      @kaxpy!(n, -β, v_prev, v_next)  # Faster than v_next = v_next - β * v_prev;
      v_prev = v;
    end
    β = @knrm2(n, v_next)
    v = v_next / β;

    # Check curvature: v'(A + sᵢI)v = v'Av + sᵢ ‖v‖² = δ + sᵢ because ‖v‖ = 1.
    # It is possible to show that σⱼ² (δⱼ + sᵢ - ωⱼ₋₁ / γⱼ₋₁) = pⱼᵀ (A + sᵢ I) pⱼ.
    for i = 1 : nshifts
      δhat[i] = δ + shifts[i]
      γ[i] = 1 ./ (δhat[i] - ω[i] ./ γ[i])
    end
    indefinite .|= (γ .<= 0.0);

    # Compute next CG iterate for each shifted system that has not yet converged.
    # Stop iterating on indefinite problems if requested.
    not_cv = check_curvature ? findall(.! (converged .| indefinite)) : findall(.! converged);

    # Loop is a bit faster than the vectorized version.
    for i in not_cv
      x[:, i] += γ[i] * p[:, i];  # Strangely, this is faster than a loop.

      ω[i] = β * γ[i];
      σ[i] *= -ω[i];
      ω[i] *= ω[i];
      p[:, i] = v * σ[i] + p[:, i] * ω[i];  # Faster than loop.

      # Update list of systems that have converged.
      rNorms[i] = abs(σ[i]);
      converged[i] = rNorms[i] <= ε;
    end

    length(not_cv) > 0 && append!(rNorms_history, rNorms);

    # Is there a better way than to update this array twice per iteration?
    not_cv = check_curvature ? findall(.! (converged .| indefinite)) : findall(.! converged);
    iter = iter + 1;
    verbose && local_printf(iter, rNorms...)

    solved = length(not_cv) == 0;
    tired = iter >= itmax;
  end

  status = tired ? "maximum number of iterations exceeded" : "solution good enough given atol and rtol"
  stats = LanczosStats(solved, reshape(rNorms_history, nshifts, round(Int, sum(size(rNorms_history))/nshifts))', indefinite, 0.0, 0.0, status);  # TODO: Estimate Anorm and Acond.
  return (x, stats);
end


text:
# Test basic extensions
@testset "Basics" begin
    # initialize the model, references, and other info
    m = InfiniteModel()
    @infinite_parameter(m, t in [0, 10])
    con = ScalarConstraint(zero(AffExpr), MOI.EqualTo(0.0))
    object = ConstraintData(con, Int[], "c1", MeasureIndex[], false)
    idx = InfOptConstraintIndex(1)
    cref = InfOptConstraintRef(m, idx)
    bad_cref = InfOptConstraintRef(m, InfOptConstraintIndex(-1))
    # owner_model
    @testset "JuMP.owner_model" begin
      @test owner_model(cref) === m
    end
    # index
    @testset "JuMP.index" begin
      @test JuMP.index(cref) == InfOptConstraintIndex(1)
    end
    # test Base.:(==) of constraint references
    @testset "Base.:(==) References" begin
        @test !(cref == bad_cref)
        @test cref == InfOptConstraintRef(m, idx)
        @test cref != InfOptConstraintRef(InfiniteModel(), idx)
    end
    # test broadcastable
    @testset "Base.broadcastable Reference" begin
        @test isa(Base.broadcastable(cref), Base.RefValue)
    end
    # _add_data_object
    @testset "_add_data_object" begin
        @test InfiniteOpt._add_data_object(m, object) == idx
    end
    # _data_dictionary
    @testset "_data_dictionary" begin
        @test InfiniteOpt._data_dictionary(cref) === m.constraints
    end
    # is_valid
    @testset "JuMP.is_valid" begin
      @test is_valid(m, cref)
    end
    # _data_object
    @testset "_data_object" begin
        @test InfiniteOpt._data_object(cref) === object
        @test_throws ErrorException InfiniteOpt._data_object(bad_cref)
    end
    # _core_constraint_object
    @testset "_core_constraint_object" begin
        @test InfiniteOpt._core_constraint_object(cref) === con
    end
    # _set_core_constraint_object
    @testset "_set_core_constraint_object" begin
        @test InfiniteOpt._set_core_constraint_object(cref, con) isa Nothing
        con2 = ScalarConstraint(zero(AffExpr), MOI.LessThan(0.0))
        @test InfiniteOpt._set_core_constraint_object(cref, con2) isa Nothing
        @test InfiniteOpt._core_constraint_object(cref) === con2
        @test InfiniteOpt._set_core_constraint_object(cref, con) isa Nothing
        @test InfiniteOpt._core_constraint_object(cref) === con
    end
    # _object_numbers
    @testset "_object_numbers" begin
        @test InfiniteOpt._object_numbers(cref) == []
    end
    # _measure_dependencies
    @testset "_measure_dependencies" begin
        @test InfiniteOpt._measure_dependencies(cref) == MeasureIndex[]
    end
    @testset "_is_info_constraint" begin
        @test !InfiniteOpt._is_info_constraint(cref)
    end
    # constraint_object
    @testset "JuMP.constraint_object" begin
      @test constraint_object(cref) == con
    end
    # JuMP.name
    @testset "JuMP.name" begin
        @test name(cref) == "c1"
        @test name(bad_cref) == ""
    end
    # JuMP.set_name
    @testset "JuMP.set_name" begin
        @test isa(set_name(cref, "new"), Nothing)
        @test_throws ErrorException set_name(bad_cref, "")
        @test name(cref) == "new"
    end
    # test has_domain_restrictions
    @testset "has_domain_restrictions" begin
        @test !has_domain_restrictions(cref)
    end
    # test domain_restrictions
    @testset "domain_restrictions" begin
        @test domain_restrictions(cref) == DomainRestrictions()
    end
    # _make_constraint_ref
    @testset "_make_constraint_ref" begin
        @test InfiniteOpt._make_constraint_ref(m, idx) == cref
    end
    # constraint_by_name
    @testset "JuMP.constraint_by_name" begin
        # test normal
        @test constraint_by_name(m, "new") == cref
        @test isa(constraint_by_name(m, "test2"), Nothing)
        # prepare constraint with duplicate name
        idx2 = InfiniteOpt._add_data_object(m, object)
        cref2 = InfOptConstraintRef(m, idx2)
        @test isa(set_name(cref2, "new"), Nothing)
        # test for duplciate name error
        @test_throws ErrorException constraint_by_name(m, "new")
    end
    # _delete_data_object
    @testset "_delete_data_object" begin
        @test InfiniteOpt._delete_data_object(cref) isa Nothing
        @test length(InfiniteOpt._data_dictionary(cref)) == 1
        @test !is_valid(m, cref)
    end
end

# Test definition methods
@testset "Definition" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1])
    @infinite_parameter(m, pars[1:2] in [0, 1])
    @variable(m, inf, Infinite(par))
    @variable(m, inf2, Infinite(par, pars))
    @variable(m, pt, Point(inf, 0.5))
    @variable(m, x)
    @variable(m, rv, SemiInfinite(inf2, 0.5, pars))
    data = TestData(par, 1, 1)
    meas = @measure(inf + par - x, data, name = "test")
    mindex = MeasureIndex(1)
    dinf = @deriv(inf, par)
    @finite_parameter(m, fin == 42)
    # test _check_restrictions
    @testset "_check_restrictions" begin
        # test normal
        rs = DomainRestrictions(par => [0, 1])
        @test InfiniteOpt._check_restrictions(rs) isa Nothing
        rs = DomainRestrictions(pars => [0, 1])
        @test InfiniteOpt._check_restrictions(rs) isa Nothing
        rs = DomainRestrictions(pars => 0)
        @test InfiniteOpt._check_restrictions(rs) isa Nothing
        rs = DomainRestrictions(pars[1] => [0, 1])
        @test InfiniteOpt._check_restrictions(rs) isa Nothing
        # test errors
        rs = DomainRestrictions(fin => [0, 1])
        @test_throws ErrorException InfiniteOpt._check_restrictions(rs)
        rs = DomainRestrictions(par => [-1, 1])
        @test_throws ErrorException InfiniteOpt._check_restrictions(rs)
        rs = DomainRestrictions(par => [0, 2])
        @test_throws ErrorException InfiniteOpt._check_restrictions(rs)
        rs = DomainRestrictions(pars[1] => 0)
        @test_throws ErrorException InfiniteOpt._check_restrictions(rs)
        rs = DomainRestrictions(pars[1] => 0, pars[2] => [0, 3])
        @test_throws ErrorException InfiniteOpt._check_restrictions(rs)
    end
    # test build_constraint for DomainRestrictedConstraint
    @testset "build_constraint (DomainRestrictedConstraint)" begin
        c = JuMP.ScalarConstraint(inf + x, MOI.EqualTo(0.0))
        rs = DomainRestrictions(par => [0, 1])
        @test build_constraint(error, inf + x, MOI.EqualTo(0.0), rs).constraint isa ScalarConstraint
        @test build_constraint(error, inf + x, MOI.EqualTo(0.0), rs).restrictions == rs
    end
    # test _validate_restrictions
    @testset "_validate_restrictions" begin
        # test normal
        rs = DomainRestrictions(par => [0, 1])
        @test InfiniteOpt._validate_restrictions(m, rs) isa Nothing
        @test num_supports(par) == 1
        # test error
        @infinite_parameter(InfiniteModel(), par2 ~ Normal())
        rs = DomainRestrictions(par2 => [0, 1])
        err = VariableNotOwned{GeneralVariableRef}
        @test_throws err InfiniteOpt._validate_restrictions(m, rs)
        # test support addition
        rs = DomainRestrictions(par => 0)
        @test InfiniteOpt._validate_restrictions(m, rs) isa Nothing
        @test supports(par, label = UserDefined) == [0, 0.5]
        rs = DomainRestrictions(pars => 0)
        @test InfiniteOpt._validate_restrictions(m, rs) isa Nothing
        @test supports(pars) == zeros(2, 1)
    end
    # test _update_var_constr_mapping
    @testset "_update_var_constr_mapping" begin
        # add dumby constraints
        con = ScalarConstraint(zero(AffExpr), MOI.EqualTo(0.0))
        object1 = ConstraintData(con, Int[], "c", MeasureIndex[], false)
        object2 = ConstraintData(con, Int[], "c", MeasureIndex[], false)
        idx1 = InfOptConstraintIndex(1)
        idx2 = InfOptConstraintIndex(2)
        cref1 = InfOptConstraintRef(m, idx1)
        cref2 = InfOptConstraintRef(m, idx2)
        @test InfiniteOpt._add_data_object(m, object1) == idx1
        @test InfiniteOpt._add_data_object(m, object2) == idx2
        # test initial use of variable
        @test isa(InfiniteOpt._update_var_constr_mapping([inf, par, meas, rv], cref1),
                  Nothing)
        @test InfiniteOpt._constraint_dependencies(inf) == [idx1]
        @test InfiniteOpt._constraint_dependencies(par) == [idx1]
        @test InfiniteOpt._constraint_dependencies(meas) == [idx1]
        @test InfiniteOpt._constraint_dependencies(rv) == [idx1]
        @test InfiniteOpt._measure_dependencies(cref1) == [mindex]
        # test secondary use of variable
        @test isa(InfiniteOpt._update_var_constr_mapping([inf, par, meas, rv], cref2),
                  Nothing)
        @test InfiniteOpt._constraint_dependencies(inf) == [idx1, idx2]
        @test InfiniteOpt._constraint_dependencies(par) == [idx1, idx2]
        @test InfiniteOpt._constraint_dependencies(meas) == [idx1, idx2]
        @test InfiniteOpt._constraint_dependencies(rv) == [idx1, idx2]
        @test InfiniteOpt._measure_dependencies(cref2) == [mindex]
        # Undo changes
        empty!(InfiniteOpt._constraint_dependencies(inf))
        empty!(InfiniteOpt._constraint_dependencies(par))
        empty!(InfiniteOpt._constraint_dependencies(meas))
        empty!(InfiniteOpt._constraint_dependencies(rv))
        empty!(InfiniteOpt._measure_dependencies(cref1))
        empty!(InfiniteOpt._measure_dependencies(cref2))
        empty!(m.constraints)
    end
    # test add_constraint
    @testset "JuMP.add_constraint" begin
        # test rejects bad variables
        @variable(InfiniteModel(), z)
        con = ScalarConstraint(z, MOI.EqualTo(42.0))
        err = VariableNotOwned{GeneralVariableRef}
        @test_throws err add_constraint(m, con, "a")
        # test restricted constraint
        rs = DomainRestrictions(par => [0, 1])
        con = ScalarConstraint(inf + pt, MOI.EqualTo(42.0))
        con = DomainRestrictedConstraint(con, rs)
        idx = InfOptConstraintIndex(1)
        cref = InfOptConstraintRef(m, idx)
        @test add_constraint(m, con, "a") == cref
        @test name(cref) == "a"
        @test domain_restrictions(cref) == rs
        @test constraint_object(cref).set == MOI.EqualTo(42.0)
        # test infinite constraint
        con = ScalarConstraint(inf + pt, MOI.EqualTo(42.0))
        idx = InfOptConstraintIndex(2)
        cref = InfOptConstraintRef(m, idx)
        @test add_constraint(m, con, "b") == cref
        @test name(cref) == "b"
        @test domain_restrictions(cref) == DomainRestrictions()
        @test constraint_object(cref).set == MOI.EqualTo(42.0)
        # test finite constraint
        con = ScalarConstraint(x + pt, MOI.EqualTo(42.0))
        idx = InfOptConstraintIndex(3)
        cref = InfOptConstraintRef(m, idx)
        @test add_constraint(m, con, "d") == cref
        @test name(cref) == "d"
        @test !InfiniteOpt._is_info_constraint(cref)
        @test !optimizer_model_ready(m)
        @test used_by_constraint(pt)
        # test vector constraint
        con = VectorConstraint([inf + pt, 2inf], MOI.Zeros(2))
        idx = InfOptConstraintIndex(4)
        cref = InfOptConstraintRef(m, idx)
        @test add_constraint(m, con, "e") == cref
        @test constraint_object(cref) isa VectorConstraint
        @test !has_domain_restrictions(cref)
    end
    # test macro
    @testset "JuMP.@constraint" begin
        # test scalar constraint
        idx = InfOptConstraintIndex(5)
        cref = InfOptConstraintRef(m, idx)
        @test @constraint(m, f, x + pt -2 <= 2) == cref
        @test InfiniteOpt._core_constraint_object(cref) isa ScalarConstraint
        # test restricted scalar constraint
        rs = DomainRestrictions(par => [0, 1])
        idx = InfOptConstraintIndex(6)
        cref = InfOptConstraintRef(m, idx)
        @test @constraint(m, g, inf + meas - dinf <= 2, rs) == cref
        @test InfiniteOpt._object_numbers(cref) == [1]
        @test InfiniteOpt._core_constraint_object(cref) isa ScalarConstraint
        @test used_by_constraint(dinf)
    end
    # test deprecation
    @testset "@BDconstraint" begin
        @test_macro_throws ErrorException @BDconstraint(m, bd(par == 0), x == 0)
    end
end

# Test parameter reference methods
@testset "Parameter References" begin
    m = InfiniteModel()
    @infinite_parameter(m, t in [0, 1])
    @infinite_parameter(m, y in [0, 1])
    @infinite_parameter(m, x[1:3] in [0, 1])
    @variable(m, z)
    @constraint(m, c1, z >= 0)
    @constraint(m, c2, z + t + x[1] >= 0)
    # test parameter_refs
    @testset "parameter_refs" begin
        @test parameter_refs(c1) == ()
        @test parameter_refs(c2) == (t, x)
    end
end

# Test domain restrictions methods
@testset "Domain Restrictions" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 10])
    @infinite_parameter(m, pars[1:2] in [0, 10])
    @variable(m, inf, Infinite(par))
    @variable(m, pt, Point(inf, 0.5))
    @variable(m, x >= 0, Int)
    @constraint(m, c1, inf + x == 0, DomainRestrictions(par => [0, 1]))
    @constraint(m, c2, x * pt + x == 2)
    # test has_domain_restrictions
    @testset "has_domain_restrictions" begin
        @test has_domain_restrictions(c1)
        @test !has_domain_restrictions(c2)
        @test_deprecated has_parameter_bounds(c1)
    end
    # test domain_restrictions
    @testset "domain_restrictions" begin
        @test domain_restrictions(c1) == DomainRestrictions(par => [0, 1])
        @test domain_restrictions(c2) == DomainRestrictions()
        @test_deprecated parameter_bounds(c1)
    end
    # test set_domain_restrictions
    @testset "set_domain_restrictions" begin
        # test force error
        rs = DomainRestrictions(par => [0, 1])
        @test_throws ErrorException set_domain_restrictions(c1, rs)
        # test normal
        @test set_domain_restrictions(c2, rs) isa Nothing
        @test domain_restrictions(c2) == rs
        @test !optimizer_model_ready(m)
        # test test error with restrictions
        rs2 = DomainRestrictions(par => [-1, 1])
        @test_throws ErrorException set_domain_restrictions(c1, rs2, force = true)
        # test deprecation
        empty!(domain_restrictions(c1).intervals)
        @test_deprecated set_parameter_bounds(c1, rs)
        @test_macro_throws ErrorException @set_parameter_bounds(c1, rs)
    end
    # test _update_restrictions
    @testset "_update_restrictions" begin
        # test normal
        rs1 = DomainRestrictions(pars[1] => [0, 1], pars[2] => [-1, 2])
        rs2 = DomainRestrictions(par => 0, pars[1] => [0.5, 1.5],
                                 pars[2] => [0, 1])
        @test isa(InfiniteOpt._update_restrictions(rs1, rs2), Nothing)
        @test length(rs1) == 3
        @test rs1[par] == IntervalDomain(0, 0)
        @test rs1[pars[1]] == IntervalDomain(0.5, 1)
        @test rs1[pars[2]] == IntervalDomain(0, 1)
        # test error
        new_rs = DomainRestrictions(par => 1)
        @test_throws ErrorException InfiniteOpt._update_restrictions(rs1, new_rs)
        new_rs = DomainRestrictions(par => -1)
        @test_throws ErrorException InfiniteOpt._update_restrictions(rs1, new_rs)
    end
    # test add_parameter_restrictions
    @testset "add_domain_restrictions" begin
        # test already has restrictions
        rs = DomainRestrictions(par => [0, 1])
        @test add_domain_restrictions(c1, copy(rs)) isa Nothing
        @test domain_restrictions(c1) == rs
        # test doesn't have restrictions
        @constraint(m, c3, inf + pt == 0)
        @test add_domain_restrictions(c3, copy(rs)) isa Nothing
        @test domain_restrictions(c3) == rs
        # test deprecation
        @test_deprecated add_parameter_bounds(c1, copy(rs))
        @test_macro_throws ErrorException @add_parameter_bounds(c1, rs)
    end
    # test delete_domain_restrictions
    @testset "delete_domain_restrictions" begin
        # test partial deletion
        rs = DomainRestrictions(pars[1] => [0, 1])
        @test add_domain_restrictions(c1, rs) isa Nothing
        @test delete_domain_restrictions(c1) isa Nothing
        @test !has_domain_restrictions(c1)
        # test already gone
        @test delete_domain_restrictions(c1) isa Nothing
        @test !has_domain_restrictions(c1)
        # test deprecation
        @test_deprecated delete_parameter_bounds(c1)
    end
end

# Test coefficient queries and modifications
@testset "Coefficient Methods" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1])
    @variable(m, inf, Infinite(par))
    @variable(m, pt, Point(inf, 0.5))
    @variable(m, x >= 0., Int)
    @constraint(m, c1, inf + x == 0., DomainRestrictions(par => [0, 1]))
    @constraint(m, c2, x * pt + x == 2.)
    @constraint(m, c3, [inf, x] in MOI.Zeros(2))
    # test set_normalized_coefficient
    @testset "JuMP.set_normalized_coefficient" begin
        @test isa(set_normalized_coefficient(c1, x, 2), Nothing)
        @test jump_function(constraint_object(c1)) == inf + 2x
        @test isa(set_normalized_coefficient(c2, inf, 2), Nothing)
        @test jump_function(constraint_object(c2)) == x * pt + x + 2inf
        @test_throws ErrorException set_normalized_coefficient(c3, x, 42)
    end
    # test normalized_coefficient
    @testset "JuMP.normalized_coefficient" begin
        @test normalized_coefficient(c1, x) == 2.
        @test normalized_coefficient(c1, pt) == 0.
        @test normalized_coefficient(c2, inf) == 2.
        @test normalized_coefficient(LowerBoundRef(x), x) == 1.
        @test normalized_coefficient(LowerBoundRef(x), pt) == 0.
        @test_throws ErrorException normalized_coefficient(c3, x)
    end
    # test _set_set_value
    @testset "_set_set_value" begin
        @test InfiniteOpt._set_set_value(MOI.EqualTo(1.), 2) == MOI.EqualTo(2.)
        @test InfiniteOpt._set_set_value(MOI.GreaterThan(1.), 2) == MOI.GreaterThan(2.)
        @test InfiniteOpt._set_set_value(MOI.LessThan(1.), 2) == MOI.LessThan(2.)
    end
    # test set_normalized_rhs
    @testset "JuMP.set_normalized_rhs" begin
        @test isa(set_normalized_rhs(c1, 42.), Nothing)
        @test constraint_object(c1).set == MOI.EqualTo(42.0)
        @test isa(set_normalized_rhs(c2, 42.), Nothing)
        @test constraint_object(c2).set == MOI.EqualTo(42.0)
        @test isa(set_normalized_rhs(LowerBoundRef(x), 42.0), Nothing)
        @test constraint_object(LowerBoundRef(x)).set == MOI.GreaterThan(42.0)
        @test_throws ErrorException set_normalized_rhs(c3, 42)
    end
    # test normalized_rhs
    @testset "JuMP.normalized_rhs" begin
        @test normalized_rhs(c1) == 42.
        @test normalized_rhs(c2) == 42.
        @test normalized_rhs(LowerBoundRef(x)) == 42.
        @test_throws ErrorException normalized_rhs(c3)
    end
    # test add_to_function_constant
    @testset "JuMP.add_to_function_constant" begin
        @test isa(add_to_function_constant(c1, 4.), Nothing)
        @test constraint_object(c1).set == MOI.EqualTo(38.0)
        @test isa(add_to_function_constant(c2, 4.), Nothing)
        @test constraint_object(c2).set == MOI.EqualTo(38.0)
        @test isa(add_to_function_constant(LowerBoundRef(x), 4.0), Nothing)
        @test constraint_object(LowerBoundRef(x)).set == MOI.GreaterThan(38.0)
        @test_throws ErrorException add_to_function_constant(c3, 42)
    end
end

# Test num_constraints extensions
@testset "JuMP.num_constraints Extensions" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1])
    @variable(m, inf >= 0, Infinite(par))
    @variable(m, pt, Point(inf, 0.5))
    @variable(m, 0 <= x <= 1, Int)
    # test search function type and set type
    @testset "Function and Set" begin
        @test num_constraints(m, GeneralVariableRef, MOI.LessThan) == 1
        @test num_constraints(m, GeneralVariableRef, MOI.GreaterThan) == 3
    end
    # test search function type
    @testset "Function" begin
        @test num_constraints(m, GeneralVariableRef) == 5
    end
    # test search set type
    @testset "Set" begin
        @test num_constraints(m, MOI.LessThan) == 1
        @test num_constraints(m, MOI.GreaterThan) == 3
    end
    # test search total
    @testset "Total" begin
        @test num_constraints(m) == 5
    end
end

# Test all_constraints extensions
@testset "JuMP.all_constraints Extensions" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1])
    @variable(m, inf >= 0, Infinite(par))
    @variable(m, pt, Point(inf, 0.5))
    @variable(m, 0 <= x <= 1, Int)
    # test search function type and set type
    @testset "Function and Set" begin
        list = [InfOptConstraintRef(m, InfOptConstraintIndex(4))]
        @test all_constraints(m, GeneralVariableRef, MOI.LessThan) == list
        list = [InfOptConstraintRef(m, InfOptConstraintIndex(1)),
                InfOptConstraintRef(m, InfOptConstraintIndex(2)),
                InfOptConstraintRef(m, InfOptConstraintIndex(3))]
        @test all_constraints(m, GeneralVariableRef, MOI.GreaterThan) == list
    end
    # test search function type
    @testset "Function" begin
        list = [InfOptConstraintRef(m, InfOptConstraintIndex(1)),
                InfOptConstraintRef(m, InfOptConstraintIndex(2)),
                InfOptConstraintRef(m, InfOptConstraintIndex(3)),
                InfOptConstraintRef(m, InfOptConstraintIndex(4)),
                InfOptConstraintRef(m, InfOptConstraintIndex(5))]
        @test all_constraints(m, GeneralVariableRef) == list
    end
    # test search set type
    @testset "Set" begin
        list = [InfOptConstraintRef(m, InfOptConstraintIndex(4))]
        @test all_constraints(m, MOI.LessThan) == list
        list = [InfOptConstraintRef(m, InfOptConstraintIndex(1)),
                InfOptConstraintRef(m, InfOptConstraintIndex(2)),
                InfOptConstraintRef(m, InfOptConstraintIndex(3))]
        @test all_constraints(m, MOI.GreaterThan) == list
    end
    # test search total
    @testset "Total" begin
        list = [InfOptConstraintRef(m, InfOptConstraintIndex(1)),
                InfOptConstraintRef(m, InfOptConstraintIndex(2)),
                InfOptConstraintRef(m, InfOptConstraintIndex(3)),
                InfOptConstraintRef(m, InfOptConstraintIndex(4)),
                InfOptConstraintRef(m, InfOptConstraintIndex(5))]
        @test all_constraints(m) == list
    end
end

# Test list_of_constraint_types
@testset "JuMP.list_of_constraint_types" begin
    # initialize model and references
    m = InfiniteModel()
    @infinite_parameter(m, par in [0, 1])
    @variable(m, inf >= 0., Infinite(par))
    @variable(m, pt, Point(inf, 0.5))
    @variable(m, 0. <= x <= 1., Int)
    expected = [(GeneralVariableRef, MOI.GreaterThan{Float64}),
                (GeneralVariableRef, MOI.LessThan{Float64}),
                (GeneralVariableRef, MOI.Integer)]
    @test isempty(setdiff(list_of_constraint_types(m), expected))
end


text:
const IMAGENET_MEANS = SVector{3,Float32}(.485, 0.456, 0.406)
const IMAGENET_STDS = SVector{3,Float32}(0.229, 0.224, 0.225)

"""
    ImageTensor{N} <: Block

Block for N+1-dimensional arrays representing an N-dimensional
image with the color channels expanded.

"""
struct ImageTensor{N} <: Block
    nchannels::Int
end

function checkblock(block::ImageTensor{N}, a::AbstractArray{T,M}) where {M,N,T}
    # Tensor has dimensionality one higher and color channels need to be the same
    return (N + 1 == M) && (size(a, M) == block.nchannels)
end

"""
    ImagePreprocessing([; kwargs...]) <: Encoding

Encodes `Image`s by converting them to a common color type `C`,
expanding the color channels and normalizing the channel values.
Additionally, apply pixel-level augmentations passed in as `augmentations`
during `Training`.

Currently works with 2D images only, but this constraint will be
removed in the future.

Encodes
- `Image{2}` -> `ImageTensor{3}`

## Keyword arguments

- `augmentations::`[`DataAugmentation.Transform`](#): Augmentation to apply to every image
    before preprocessing. See [`augs_lighting`](#)
- `buffered = true`: Whether to use inplace transformations. Reduces memory usage.
- `means::SVector = IMAGENET_MEANS`: mean value of each color channel.
- `stds::SVector = IMAGENET_STDS`: standard deviation of each color channel.
- `C::Type{<:Colorant} = RGB{N0f8}`: color type to convert images to.
- `T::Type{<:Real} = Float32`: element type of output

"""
struct ImagePreprocessing{P,N,C <: Color{P,N},T <: Number} <: Encoding
    buffered::Bool
    augmentations::DataAugmentation.Transform
    stats::Tuple{SVector{N},SVector{N}}
    tfms::Dict{Context,DataAugmentation.Transform}
end

function ImagePreprocessing(;
        means::SVector{N}=IMAGENET_MEANS,
        stds::SVector{N}=IMAGENET_STDS,
        augmentations=Identity(),
        C::Type{<:Color{U,N}}=RGB{N0f8},
        T=Float32,
        buffered=true) where {N,U}
    # TODO: tensor of type T
    stats = means, stds
    basetfm = ToEltype(C) |> ImageToTensor{T}() |> Normalize(means, stds)
    if buffered
        tfms = Dict(
            Training() => BufferedThreadsafe(augmentations |> basetfm),
            Validation() => BufferedThreadsafe(basetfm),
            # Inference transform is not buffered since it can have
            # varying sizes
            Inference() => basetfm,
        )
    else
        tfms = Dict(
            Training() => augmentations |> basetfm,
            Validation() => basetfm,
            Inference() => basetfm,
        )
    end

    return ImagePreprocessing{U,N,C,T}(buffered, augmentations, stats, tfms)
end

colorchannels(C::Type{<:Color{T,N}}) where {T,N} = N

function encodedblock(ip::ImagePreprocessing{P,M,C}, ::Image{N}) where {P,M,C,N}
    return ImageTensor{N}(colorchannels(C))
end

function encode(ip::ImagePreprocessing, context, block::Image, data)
    return copy(apply(ip.tfms[context], DataAugmentation.Image(data)) |> itemdata)
end

decodedblock(::ImagePreprocessing, ::ImageTensor{N}) where N = Image{N}()

function decode(ip::ImagePreprocessing, context, block::ImageTensor, data)
    means, stds = ip.stats
    return copy(DataAugmentation.tensortoimage(DataAugmentation.denormalize(data, means, stds)))
end

# Augmentation helper

"""
    augs_lighting([; intensity = 0.2, p = 0.75])

Helper to create a set of lighting transformations for image data. With
probability `p`, applies [`AdjustBrightness`](#)`(intensity)` and
[`AdjustContrast`](#)`(intensity)`.
"""
function augs_lighting(;intensity=0.2, p=0.75)
    return Maybe(AdjustBrightness(intensity), p) |> Maybe(AdjustContrast(intensity), p)
end

# Pretty-printing

function Base.show(io::IO, ip::ImagePreprocessing{P,N,C,T}) where {P,N,C,T}
    show(io, ShowTypeOf(ip))
    fields = (
        buffered = ShowLimit(ip.buffered, limit=80),
        augmentations = ShowLimit(ip.augmentations, limit=80),
    )
    show(io, ShowProps(fields, new_lines=true))
end


text:
using RecyclePermanent
using RNGPool
using JLD2
import Statistics: mean, var

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

@load "data/gkData.jld" ys θ
ys = ys[1:100]
model = makeGKModel(θ, ys)

setRNGs(12345)

N = 100000*model.n
rpio = RPIO{model.Tx}(N, model.n, 1)
oracle = rpEstimate(model, rpio)

trials = 10000
cs = 50:10:100

means1 = Vector{Float64}(undef, length(cs))
vars1 = Vector{Float64}(undef, length(cs))
means2 = Vector{Float64}(undef, length(cs))
vars2 = Vector{Float64}(undef, length(cs))

for k in 1:length(cs)
  println(cs[k])
  N = cs[k]*model.n
  rpio = RPIO{model.Tx}(N, model.n, 1)
  vs1 = Vector{Float64}(undef, trials)
  vs2 = Vector{Float64}(undef, trials)
  for i in 1:trials
    vs1[i] = rpEstimate(model, rpio)
    vs2[i] = biasedEstimate(model, rpio)
  end
  means1[k] = mean(exp.(vs1 .- oracle))
  vars1[k] = var(exp.(vs1 .- oracle))
  means2[k] = mean(exp.(vs2 .- oracle))
  vars2[k] = var(exp.(vs2 .- oracle))
end

println(means1)
println(vars1)
println(means2)
println(vars2)
println((means2 .- 1).^2 .+ vars2)


text:
#TODO - eventually move some of these
# into the SR call itself, rather than
# passing huge options at once.

function binopmap(op)
    if op == plus
        return +
    elseif op == mult
        return *
    elseif op == sub
        return -
    elseif op == div
        return /
    elseif op == ^
        return pow
    end
    return op
end

function unaopmap(op)
    if op == log
        return logm
    elseif op == log10
        return logm10
    elseif op == log2
        return logm2
    elseif op == sqrt
        return sqrtm
    end
    return op
end

struct Options{A<:NTuple{N,Any} where {N},B<:NTuple{M,Any} where {M}}

    binops::A
    unaops::B
    bin_constraints::Array{Tuple{Int,Int}, 1}
    una_constraints::Array{Int, 1}
    ns::Int
    parsimony::Float32
    alpha::Float32
    maxsize::Int
    maxdepth::Int
    fast_cycle::Bool
    migration::Bool
    hofMigration::Bool
    fractionReplacedHof::Float32
    shouldOptimizeConstants::Bool
    hofFile::String
    npopulations::Int
    nrestarts::Int
    perturbationFactor::Float32
    annealing::Bool
    batching::Bool
    batchSize::Int
    mutationWeights::Array{Float64, 1}
    warmupMaxsize::Int
    limitPowComplexity::Bool
    useFrequency::Bool
    npop::Int
    ncyclesperiteration::Int
    fractionReplaced::Float32
    topn::Int
    verbosity::Int
    probNegate::Float32
    nuna::Int
    nbin::Int

end

function Options(;
    binary_operators::NTuple{nbin, Any}=(div, plus, mult),
    unary_operators::NTuple{nuna, Any}=(exp, cos),
    bin_constraints=nothing,
    una_constraints=nothing,
    ns=10, #1 sampled from every ns per mutation
    topn=10, #samples to return per population
    parsimony=0.000100f0,
    alpha=0.100000f0,
    maxsize=20,
    maxdepth=nothing,
    fast_cycle=false,
    migration=true,
    hofMigration=true,
    fractionReplacedHof=0.1f0,
    shouldOptimizeConstants=true,
    hofFile=nothing,
    npopulations=nothing,
    nrestarts=3,
    perturbationFactor=1.000000f0,
    annealing=true,
    batching=false,
    batchSize=50,
    mutationWeights=[10.000000, 1.000000, 1.000000, 3.000000, 3.000000, 0.010000, 1.000000, 1.000000],
    warmupMaxsize=0,
    limitPowComplexity=false,
    useFrequency=false,
    npop=1000,
    ncyclesperiteration=300,
    fractionReplaced=0.1f0,
    verbosity=convert(Int, 1e9),
    probNegate=0.01f0,
   ) where {nuna,nbin}

    if hofFile == nothing
        hofFile = "hall_of_fame.csv" #TODO - put in date/time string here
    end

    if una_constraints == nothing
        una_constraints = [-1 for i=1:nuna]
    end
    if bin_constraints == nothing
        bin_constraints = [(-1, -1) for i=1:nbin]
    end

    if maxdepth == nothing
        maxdepth = maxsize
    end

    if npopulations == nothing
        npopulations = nworkers()
    end

    binary_operators = map(binopmap, binary_operators)
    unary_operators = map(unaopmap, unary_operators)

    mutationWeights = map((x,)->convert(Float64, x), mutationWeights)
    if length(mutationWeights) != 8
        error("Not the right number of mutation probabilities given")
    end

    Options{typeof(binary_operators),typeof(unary_operators)}(binary_operators, unary_operators, bin_constraints, una_constraints, ns, parsimony, alpha, maxsize, maxdepth, fast_cycle, migration, hofMigration, fractionReplacedHof, shouldOptimizeConstants, hofFile, npopulations, nrestarts, perturbationFactor, annealing, batching, batchSize, mutationWeights, warmupMaxsize, limitPowComplexity, useFrequency, npop, ncyclesperiteration, fractionReplaced, topn, verbosity, probNegate, nuna, nbin)
end


text:
#=
Test resynthesizer in 4D using commonly available test data.
=#

# requires pkg OMETIFF
# using Pkg; Pkg.add("OMETIFF")
# OMETIFF supports 4D in .tiff format?

# The tested item
include("../../src/resynthesizer.jl")

using ImageMetadata

# Grab the open source test data
using TestImages
# "multi" is short name
fullImage = testimage("multi");

println("fullImage is:")
println(summary(fullImage))
# 167×439×3×7

#=
See JuliaImages.
The image has named axes.
And metadata??
We only want to convert the data.
Extract: data(fullImage), not fullImage.data??
=#
imageWithAxes = data(fullImage)
println("imageWithAxes is:")
println(summary(imageWithAxes))

imageView = imageWithAxes.data
println("image is:")
println(summary(imageView))

#=
imageView is still reshape(reinterpret(...)) which throws type errors.
Its a view.
So copy it.
=#
image = copy(imageView)
println("typeof copy")
println(typeof(image))

# Create a mask
mask = falses(size(image))
# ??? The third dimension is 3 channels.  Arbitrarily take them all.
mask[3:4, 2:3, 1:3, 5:6] .= true

println("mask is:")
println(summary(mask))

function test(data, mask)
    result = resynthesize(data, mask)

    # Assert result is 3D array

    #=
    Chunk it back into the Image for display.
    =#
    #fullImage.data = result
    # imshow(result)

    # return type so not massive print object
    typeof(result)
end

#imshow requires    using ImageView, Images

max_stars_repo_name: bootchk/resynthesizerJulia

max_issues_repo_name: bootchk/resynthesizerJulia

max_forks_repo_name: bootchk/resynthesizerJulia