text:
abstract type AbstractAttention <: Layer end

struct MultiheadAttention <: AbstractAttention
    head::Int
    future::Bool
    iqproj::Dense
    ikproj::Dense
    ivproj::Dense
    oproj::Dense
    drop::Dropout
end

"""
    MultiheadAttention(head::Int, is::Int, hs::Int, os::Int; future::Bool=true, pdrop = 0.1)
Multihead dot product Attention Layer, `head` is the number of head, `is` is the input size, `hs` is the hidden size of input projection layer of each head,
`os` is the output size. When `future` is `false`, the k-th token can't see tokens at > k. `pdrop` is the dropout rate.
"""
MultiheadAttention(head::Int,
                   is::Int,
                   hs::Int,
                   os::Int;
                   future::Bool=true, pdrop = 0.1) = MultiheadAttention(head,
                                                                        future,
                                                                        Dense(input=is,output=hs*head),
                                                                        Dense(input=is,output=hs*head),
                                                                        Dense(input=is,output=hs*head),
                                                                        Dense(input=hs*head, output=os),
                                                                        Dropout(pdrop),
                                                                        )

function Base.show(io::IO, mh::MultiheadAttention)
    hs = div(size(mh.iqproj)[1], mh.head)
    is = size(mh.iqproj)[end]
    os = size(mh.oproj)[1]

    print(io, "MultiheadAttention(")
    print(io, "head=$(mh.head), ")
    print(io, "head_size=$(hs), ")
    print(io, "$(is)=>$(os)")

    print(io, ", dropout=$(mh.drop.p))")
end

function (mh::MultiheadAttention)(query,
                                  key,
                                  value;
                                  mask=nothing)
    qs = size(query)
    ks = size(key)
    vs = size(value)
    if length(qs) == 3
        #size(ipq) == (h, q_seq_len, batch)
        ipq = mh.iqproj(query)
        ipk = mh.ikproj(key)
        ipv = mh.ivproj(value)

        h = size(ipq, 1)
        hs = div(h, mh.head)

        #size(ipq) == (hs, q_seq_len, head, batch)
        ipq = permutedims(reshape(ipq, hs, mh.head, qs[2], qs[3]), [1, 3, 2, 4])
        ipk = permutedims(reshape(ipk, hs, mh.head, ks[2], ks[3]), [1, 3, 2, 4])
        ipv = permutedims(reshape(ipv, hs, mh.head, vs[2], vs[3]), [1, 3, 2, 4])

        #size(ipq) == (hs, q_seq_len, head * batch)
        ipq = reshape(ipq, hs, qs[2], :)
        ipk = reshape(ipk, hs, ks[2], :)
        ipv = reshape(ipv, hs, vs[2], :)

        atten = attention(ipq,ipk,ipv;
                          mask=mask,
                          future=mh.future,
                          dropout=mh.drop)

        atten = permutedims(reshape(atten, hs, qs[2], mh.head, qs[3]), [1, 3, 2, 4]) #size(atten) == (hs, head, ql, b)
        atten = reshape(atten, h, qs[2], qs[3]) #size(atten) == (h, ql, b)

        return  mh.oproj(atten)
    else
      ipq = mh.iqproj(query)
      ipk = mh.ikproj(key)
      ipv = mh.ivproj(value)

      h = size(ipq)[1] #h == hs * head
      hs = div(h, mh.head)

      #size(hq) == (hs, seq_len, head)
      hq = permutedims(reshape(ipq, hs, mh.head, :), [1, 3, 2])
      hk = permutedims(reshape(ipk, hs, mh.head, :), [1, 3, 2])
      hv = permutedims(reshape(ipv, hs, mh.head, :), [1, 3, 2])

      atten = attention(hq, hk, hv;
                        mask=mask,
                        future=mh.future,
                        dropout=mh.drop)

      # size(atten) == (head*hs, seq_len)
      atten = reshape(permutedims(atten, [1, 3, 2]), h, :)

      return mh.oproj(atten)
    end
end

function attention(query,
                   key,
                   value;
                   mask=nothing,
                   future::Bool = false,
                   dropout=nothing)
    T = eltype(query)
    dk = size(key, 1)
    score = bmm(key, query; transA=true)
    score = score ./ convert(T , sqrt(dk))

    s = size(score)

    if mask !== nothing
        #weird issue on @. mask = (1 - mask) * -1e9 which casue mask to be -Inf
        mask = (T(1.0) .- mask) .* T(-1e9)
        ms = size(mask)
        #score = score .+ mask; use broadcast instead of repeat mask for head
        score = reshape(reshape(score, s[1:end-1]..., :, ms[end]) .+ reshape(mask, ms[1:end-1]..., 1, ms[end]), s)
    end

    if !future
        #without ... will cause data move back to cpu
        fmask = convert(arrtype,tril!(fill!(Matrix{T}(undef,s[1:end-1]...),T(-1e9)),-1))
        #fmask = tril!(fill!(similar(score, s[1:end-1]...), convert(T, -1e9)), -1)
        score = score .+ fmask
    end

    score = softmax(score;dims=1)
    dropout !== nothing && (score = dropout(score))
    bmm(value, score) #size(return) == (dims, q_seq_len, batch)
end


text:
using PQEcolaPoint

# setup DTE -
T = 300.0
DTE = range(T, stop=0.0,step=-1.0) |> collect

# setup underlying range -
Sₒ = 76.04
IV = 0.4456
σ = Sₒ*(IV)*(√(T/365))
L = round(Sₒ - 1.96*σ)  # 95%
U = round(Sₒ + 1.96*σ)  # 95%
S = range(max.(L, 0.0),stop=U, step=0.5) |> collect

# setup IV array -
IV_array = range(0.1, stop = 1.2, step = 0.05) |> collect

# setup strike coordinates -
K = range(50.0, stop = 125.0, step = 5.0) |> collect

# build JIT lattice model -
JIT_lattice = build(CRRJITContractPremiumLatticeModel, PutContractModel, S, K, IV_array, DTE; number_of_levels = 80)

# encode the to and from points -
from_point = encode(JIT_lattice; S = 78.0, DTE = 35.0, IV = 0.4498, K = 75.0)
to_point = encode(JIT_lattice; S = Sₒ, DTE = 49.0, IV = 0.4681, K = 70.0)
p_array = move(JIT_lattice; from = from_point, to = to_point)

ext: jl


text:
using DiscreteChoice
using Test

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


text:
include("imports.jl")

@testset ExtendedTestSet "referencetype" begin
    @test referencetype(Main, nothing) == :module
    @test referencetype(Base, :π) == :const
    @test referencetype(Base, :sum) == :function
    @test referencetype(Base, :UnitRange) == :ptype
end

referencetype(Base, :sum)

Reference(Base, :sum)


text:
export
    # cudnnRNNInputMode_t
    CUDNN_LINEAR_INPUT,
    CUDNN_SKIP_INPUT,

    # cudnnDirectionMode_t
    CUDNN_UNIDIRECTIONAL,
    CUDNN_BIDIRECTIONAL,

    # cudnnRNNMode_t
    CUDNN_RNN_RELU,
    CUDNN_RNN_TANH,
    CUDNN_LSTM,
    CUDNN_GRU

type RNNDesc
    ptr::Ptr{Void}

    function RNNDesc()
        p = Ptr{Void}[0]
        cudnnCreateRNNDescriptor(p)
        desc = new(p[1])
        finalizer(desc, cudnnDestroyRNNDescriptor)
        desc
    end
end

Base.unsafe_convert(::Type{Ptr{Void}}, desc::RNNDesc) = desc.ptr

type RNN
    desc::RNNDesc
    wdesc
    w
    hxdesc
    hx
    cxdesc
    cx
end

"""
* w
* hx: (hsize, batchsize, layersize)
* cx: same as hx
"""
function RNN{T}(w::CuArray{T}, hx::CuArray{T,3}, cx::CuArray{T,3},
    droprate::Float64, dir, mode; inputmode=CUDNN_LINEAR_INPUT)

    h = handle(w)
    rnndesc = RNNDesc()
    dropdesc = DropoutDesc(h, droprate)
    hsize = size(hx, 1)
    nlayers = size(hx, 3)
    cudnnSetRNNDescriptor(rnndesc, hsize, nlayers, dropdesc, inputmode, dir, mode, datatype(T))

    xdesc = TensorDesc(CuArray{T}(1,size(hx,1),size(hx,2)))
    p = Csize_t[0]
    cudnnGetRNNParamsSize(h, rnndesc, xdesc, p, datatype(T))
    Int(p[1]) == length(w) || throw("The number of parameters is wrong: $(p[1]), $(length(w)).")
    wdesc = FilterDesc(reshape(w,1,1,length(w)))

    #=
    if mode == CUDNN_RNN_RELU || mode == CUDNN_RNN_TANH
        nids = 2
    elseif mode == CUDNN_LSTM
        nids = 8
    elseif mode == CUDNN_GRU
        nids = 6
    end
    for l = 1:(dir == CUDNN_BIDIRECTIONAL ? nlayers*2 : nlayers)
        for id = 1:nids
            p = Ptr{Void}[0]
            cudnnGetRNNLinLayerMatrixParams(h, rnndesc, l-1, xdesc, wdesc, w, id-1, C_NULL, p)
            println(UInt64(p[1]))

            p = Ptr{Void}[0]
            cudnnGetRNNLinLayerBiasParams(h, rnndesc, l-1, xdesc, wdesc, w, id-1, C_NULL, p)
            println(UInt64(p[1]))
        end
    end
    =#
    RNN(rnndesc, wdesc, w, TensorDesc(hx), hx, TensorDesc(cx), cx)
end

function (rnn::RNN){T}(xs::Vector{CuMatrix{T}})
    h = handle(xs[1])
    xdescs = map(x -> TensorDesc(reshape(x,1,size(x)...)), xs)
    xdesc = ydesc = map(x -> x.ptr, xdescs)
    y = similar(x)
    hydesc = rnn.hxdesc
    hy = similar(hx)
    cydesc = rnn.cxdesc
    cy = similar(cx)

    p = Csize_t[0]
    cudnnGetRNNWorkspaceSize(h, rnn.desc, length(xs), xdesc, p)
    workspace = CuArray{Int8}(Int(p[1]))

    p = Csize_t[0]
    cudnnGetRNNTrainingReserveSize(h, rnn.desc, length(xs), xdesc, p)
    reservespace = CuArray{Int8}(Int(p[1]))

    x = map(x -> Ptr{Void}(x), xs)
    cudnnRNNForwardTraining(h, rnn.desc, length(xs), xdesc, x, rnn.hxdesc, rnn.hx, rnn.cxdesc, rnn.cx,
        rnn.wdesc, rnn.w, ydesc, y, hydesc, hy, cydesc, cy,
        workspace, length(workspace), reservespace, length(reservespace))
end

function rnn{T}(hsize::Int, nlayers::Int, droprate::Float64, dir, mode, w::CuMatrix{T}, b::CuMatrix{T},
    xs::Vector{CuMatrix{T}}, hx::CuArray{T}, cx::CuArray{T};
    inputmode=CUDNN_LINEAR_INPUT)

    h = handle(xs[1])
    rnndesc = RNNDesc()
    dropdesc = DropoutDesc(h, droprate)
    cudnnSetRNNDescriptor(rnndesc, hsize, nlayers, dropdesc, inputmode, dir, mode, datatype(T))

    x = map(x -> Ptr{Void}(x), xs)
    xdescs = map(x -> TensorDesc(reshape(x,1,size(x)...)), xs)
    xdesc = map(x -> x.ptr, xdescs)
    hxdesc = cxdesc = hydesc = cydesc = TensorDesc(hx)
    wdesc = FilterDesc(reshape(w,1,1,length(w)+length(b)))

    y = similar(x)
    hy = similar(hx)
    cy = similar(cx)

    p = Csize_t[0]
    cudnnGetRNNWorkspaceSize(h, rnndesc, length(xs), xdesc, p)
    workspace = CuArray{Int8}(Int(p[1]))

    p = Csize_t[0]
    cudnnGetRNNTrainingReserveSize(h, rnndesc, length(xs), xdesc, p)
    reservespace = CuArray{Int8}(Int(p[1]))

    cudnnRNNForwardTraining(h, rnndesc, length(xs), xdesc, x, hxdesc, hx, cxdesc, cx,
        wdesc, w, ydesc, y, hydesc, hy, cydesc, cy,
        workspace, length(workspace), reservespace, length(reservespace))
end

function getnlinlayers(mode)
    mode == CUDNN_RNN_RELU && return 2
    mode == CUDNN_RNN_TANH && return 2
    mode == CUDNN_RNN_LSTM && return 8
    mode == CUDNN_GRU && return 6
    throw("Invalid mode: $mode.")
end


text:
#=
The following examples are taken from the source code of the delay differential equation
solver RADAR5 by Nicola Guglielmi and Ernst Hairer:
http://www.unige.ch/~hairer/radar5-v2.1.tar
=#

# OREGONATOR example
@doc raw"""
    prob_dde_RADAR5_oregonator

Delay differential equation model from chemical kinetics, given by

```math
  u_1'(t) = - k_1 A u_2(t) - k_2 u_1(t) u_2(t - \tau) + k_3 B u_1(t) - 2 k_4 u_1(t)^2,
```
```math
  u_2'(t) = - k_1 A u_2(t) - k_2 u_1(t) u_2(t - \tau) + f k_3 B u_1(t),
```

for ``t \in [0, 100.5]`` with history function

```math
  \phi_1(t) = 1e-10,
```
```math
  \phi_2(t) = 1e-5,
```

for ``t \leq 0``, where ``k_1 = 1.34``, ``k_2 = 1.6e9``, ``k_3 = 8000``, ``k_4 = 4e7``,
``k_5 = 1``, ``f = 1``, ``A = 0.06``, ``B = 0.06``, and ``\tau = 0.15``.

# References

Epstein, I. and Luo, Y. (1991). Differential delay equations in chemical kinetics. Nonlinear
models, Journal of Chemical Physics (95), pp. 244-254.
"""
const prob_dde_RADAR5_oregonator =
  let k₁ = 1.34, k₂ = 1.6e9, k₃ = 8_000, k₄ = 4e7, f = 1, A = 0.06, B = 0.06, τ = 0.15
    global function f_dde_RADAR5_oregonator!(du, u, h, p, t)
      # past value of the second component
      v = h(p, t - τ; idxs = 2)

      # precalculations
      a = - k₁ * A * u[2] - k₂ * u[1] * v
      b = k₃ * B * u[1]

      du[1] = a + b - 2 * k₄ * u[1]^2
      du[2] = a + f * b

      nothing
    end

    global function h_dde_RADAR5_oregonator(p, t; idxs::Union{Nothing,Int} = nothing)
      t ≤ 0 || error("history function is only implemented for t ≤ 0")

      if idxs === nothing
        [1e-10, 1e-5]
      elseif idxs == 2
        1e-5
      else
        error("history function is only implemented for the second component")
      end
    end

    DDEProblem(f_dde_RADAR5_oregonator!, h_dde_RADAR5_oregonator, (0.0, 100.5);
               constant_lags = [τ])
  end

# ROBERTSON example
@doc raw"""
    prob_dde_RADAR5_robertson

Delay differential equation model of a chemical reaction with steady state solution, given
by

```math
  u_1'(t) = - a u_1(t) + b u_2(t - \tau) u_3(t),
```
```math
  u_2'(t) = a u_1(t) - b u_2(t - \tau) u_3(t) - c u_2(t)^2,
```
```math
  u_3'(t) = c u_2(t)^2,
```

for ``t \in [0, 10e10]`` with history function ``\phi_1(0) = 1``, ``\phi_2(t) = 0`` for
``t \in [-\tau, 0]``, and ``\phi_3(0) = 0``, where ``a = 0.04``, ``b = 10_000``,
``c = 3e7``, and ``\tau = 0.01``.

# References

Guglielmi, N. and Hairer, E. (2001). Implementing Radau IIA methods for stiff delay
differential equations, Computing (67), pp. 1-12.
"""
const prob_dde_RADAR5_robertson =
  let a = 0.04, b = 10_000, c = 3e7, τ = 0.01
    global function f_dde_RADAR5_robertson!(du, u, h, p, t)
      # past value of the second component
      v = h(p, t - τ; idxs = 2)

      # precalculations
      x = a * u[1] - b * v * u[3]
      y = c * u[2]^2

      du[1] = - x
      du[2] = x - y
      du[3] = y

      nothing
    end

    global function h_dde_RADAR5_robertson(p, t; idxs::Union{Nothing,Int} = nothing)
      - τ ≤ t ≤ 0 || error("history function is only implemented for t ∈ [-τ, 0]")

      if idxs === nothing
        [1.0, 0.0, 0.0]
      elseif idxs == 2
        0.0
      else
        error("history function is only implemented for the second component")
      end
    end

    DDEProblem(f_dde_RADAR5_robertson!, h_dde_RADAR5_robertson, (0.0, 1e10);
               constant_lags = [τ])
  end

# WALTMAN example
@doc raw"""
    prob_dde_RADAR5_waltman

Delay differential equation model of antibody production, given by

```math
  u_1'(t) = - r u_1(t) u_2(t) - s u_1(t) u_4(t),
```
```math
  u_2'(t) = - r u_1(t) u_2(t) + \alpha r u_1(u_5(t)) u_2(u_5(t)) [t \geq t_0],
```
```math
  u_3'(t) = r u_1(t) u_2(t),
```
```math
  u_4'(t) = - s u_1(t) u_4(t) - \gamma u_4(t) + \beta r u_1(u_6(t)) u_2(u_6(t)) [t > t_1],
```
```math
  u_5'(t) = [t \geq t_0] \frac{u_1(t) u_2(t) + u_3(t)}{u_1(u_5(t)) u_2(u_5(t)) + u_3(u_5(t))},
```
```math
  u_6'(t) = [t \geq t_1] \frac{1e-12 + u_2(t) + u_3(t)}{1e-12 + u_2(u_6(t)) + u_3(u_6(t))},
```

for ``t \in [0, 300]`` with history function

```math
  \phi_1(t) = \phi_0,
```
```math
  \phi_2(t) = 1e-15,
```
```math
  \phi_3(t) = 0,
```
```math
  \phi_4(t) = 0,
```
```math  \phi_5(t) = 0,
```
```math
  \phi_6(t) = 0,
```

for ``t \leq 0``, where ``\alpha = 1.8``, ``\beta = 20``, ``\gamma = 0.002``, ``r = 5e4``,
``s = 1e5``, ``t_0 = 32``, ``t_1 = 119``, and ``\phi_0 = 0.75e-4``.

# References

Waltman, P. (1978). A threshold model of antigen-stimulated antibody production, Theoretical
Immunology (8), pp. 437-453.
"""
prob_dde_RADAR5_waltman

let α = 1.8, β = 20, γ = 0.002, r = 50_000, s = 100_000
  global function f_dde_RADAR5_waltman!(du, u, h, p, t)
    # time of switches
    t₀, t₁ = p.t₀, p.t₁

    # precalculations
    u₁u₂ = u[1] * u[2]
    a = r * u₁u₂
    b = s * u[1] * u[4]

    du[1] = - a - b
    du[3] = a

    if t < t₀
      du[2] = - a
      du[5] = 0
    else
      v = h(p, u[5])
      v₁v₂ = v[1] * v[2]
      du[2] = - a + α * r * v₁v₂
      du[5] = (u₁u₂ + u[3]) / (v₁v₂ + v[3])
    end

    if t < t₁
      du[4] = - b - γ * u[4]
      du[6] = 0
    else
      w = h(p, u[6])
      du[4] = - b - γ * u[4] + β * r * w[1] * w[2]
      du[6] = (1e-12 + u[2] + u[3]) / (1e-12 + w[2] + w[3])
    end

    nothing
  end
end

function h_dde_RADAR5_waltman(p, t)
  t ≤ 0 || error("history function is only implemented for t ≤ 0")

  [p.ϕ₀, 1e-15, 0.0, 0.0, 0.0, 0.0]
end

const prob_dde_RADAR5_waltman =
  DDEProblem(f_dde_RADAR5_waltman!, (p, t) -> h_dde_RADAR5_waltman(p, 0.0),
             h_dde_RADAR5_waltman, (0.0, 300.0), (ϕ₀ = 0.75e-4, t₀ = 32, t₁ = 119);
             dependent_lags = ((u, p, t) -> t - u[5], (u, p, t) -> t - u[6]),
             order_discontinuity_t0 = 1)
const prob_dde_RADAR5_waltman_1 = prob_dde_RADAR5_waltman

@doc raw"""
    prob_dde_RADAR5_waltman_2

Same delay differential equation as [`prob_dde_RADAR5_waltman`] with ``t_0 = 32``,
``t_1 = 111``, and ``\phi_0 = 0.5e-4``.

# References

Waltman, P. (1978). A threshold model of antigen-stimulated antibody production, Theoretical
Immunology (8), pp. 437-453.
"""
const prob_dde_RADAR5_waltman_2 =
  remake(prob_dde_RADAR5_waltman; p = (ϕ₀ = 0.5e-4, t₀ = 32, t₁ = 111))

@doc raw"""
    prob_dde_RADAR5_waltman_3

Same delay differential equation as [`prob_dde_RADAR5_waltman`] with ``t_0 = 33``,
``t_1 = 145``, and ``\phi_0 = 1e-5``.

# References

Waltman, P. (1978). A threshold model of antigen-stimulated antibody production, Theoretical
Immunology (8), pp. 437-453.
"""
const prob_dde_RADAR5_waltman_3 =
  remake(prob_dde_RADAR5_waltman; p = (ϕ₀ = 1e-5, t₀ = 33, t₁ = 145))

@doc raw"""
    prob_dde_RADAR5_waltman_4

Same delay differential equation as [`prob_dde_RADAR5_waltman`] with ``t_0 = 34``,
``t_1 = 163``, and ``\phi_0 = 0.75e-5``.

# References

Waltman, P. (1978). A threshold model of antigen-stimulated antibody production, Theoretical
Immunology (8), pp. 437-453.
"""
const prob_dde_RADAR5_waltman_4 =
  remake(prob_dde_RADAR5_waltman; p = (ϕ₀ = 0.75e-5, t₀ = 34, t₁ = 163))

@doc raw"""
    prob_dde_RADAR5_waltman_5

Same delay differential equation as [`prob_dde_RADAR5_waltman`] with ``t_0 = 35``,
``t_1 = 197``, and ``\phi_0 = 0.5e-5``.

# References

Waltman, P. (1978). A threshold model of antigen-stimulated antibody production, Theoretical
Immunology (8), pp. 437-453.
"""
const prob_dde_RADAR5_waltman_5 =
  remake(prob_dde_RADAR5_waltman; p = (ϕ₀ = 0.5e-5, t₀ = 35, t₁ = 197))


text:
abstract type AbstractTerm end
const TermOrTerms = Union{AbstractTerm, NTuple{N, AbstractTerm} where N}
const TupleTerm = NTuple{N, TermOrTerms} where N

width(::T) where {T<:AbstractTerm} =
    throw(ArgumentError("terms of type $T have undefined width"))

"""
    Term <: AbstractTerm

A placeholder for a variable in a formula where the type (and necessary data
invariants) is not yet known.  This will be converted to a
[`ContinuousTerm`](@ref) or [`CategoricalTerm`](@ref) by [`apply_schema`](@ref).

# Fields

* `sym::Symbol`: The name of the data column this term refers to.
"""
struct Term <: AbstractTerm
    sym::Symbol
end
width(::Term) =
    throw(ArgumentError("Un-typed Terms have undefined width.  " *
                        "Did you forget to call apply_schema?"))

"""
    ConstantTerm{T<:Number} <: AbstractTerm

Represents a literal number in a formula.  By default will be converted to
[`InterceptTerm`] by [`apply_schema`](@ref).

# Fields

* `n::T`: The number represented by this term.
"""
struct ConstantTerm{T<:Number} <: AbstractTerm
    n::T
end
width(::ConstantTerm) = 1

"""
    FormulaTerm{L,R} <: AbstractTerm

Represents an entire formula, with a left- and right-hand side.  These can be of
any type (captured by the type parameters).

# Fields

* `lhs::L`: The left-hand side (e.g., response)
* `rhs::R`: The right-hand side (e.g., predictors)
"""
struct FormulaTerm{L,R} <: AbstractTerm
    lhs::L
    rhs::R
end

"""
    FunctionTerm{Forig,Fanon,Names} <: AbstractTerm

Represents a call to a Julia function.  The first type parameter is the type
of the function as originally specified (e.g., `typeof(log)`), while the second
is the type of the anonymous function that will be applied element-wise to the
data table.

The `FunctionTerm` _also_ captures the arguments of the original call and parses
them _as if_ they were part of a special DSL call, applying the rules to expand
`*`, distribute `&` over `+`, and wrap symbols in `Term`s.

By storing the original function as a type parameter _and_ pessimistically
parsing the arguments as if they're part of a special DSL call, this allows
custom syntax to be supported with minimal extra effort.  Packages can dispatch
on `apply_schema(f::FunctionTerm{typeof(special_syntax)}, schema,
::Type{<:MyModel})` and pull out the arguments parsed as terms from
`f.args_parsed` to construct their own custom terms.

# Fields

* `forig::Forig`: the original function (e.g., `log`)
* `fanon::Fanon`: the generated anonymous function (e.g., `(a, b) -> log(1+a+b)`)
* `exorig::Expr`: the original expression passed to `@formula`
* `args_parsed::Vector`: the arguments of the call passed to `@formula`, each
  parsed _as if_ the call was a "special" DSL call.

# Type parameters

* `Forig`: the type of the original function (e.g., `typeof(log)`)
* `Fanon`: the type of the generated anonymous function
* `Names`: the names of the arguments to the anonymous function (as a
  `NTuple{N,Symbol}`)

# Example

```jldoctest
julia> f = @formula(y ~ log(1 + a + b))
FormulaTerm
Response:
  y(unknown)
Predictors:
  (a,b)->log(1 + a + b)

julia> typeof(f.rhs)
FunctionTerm{typeof(log),var"##1#2",(:a, :b)}

julia> f.rhs.forig(1 + 3 + 4)
2.0794415416798357

julia> f.rhs.fanon(3, 4)
2.0794415416798357

julia> modelcols(f.rhs, (a=3, b=4))
2.0794415416798357

julia> modelcols(f.rhs, (a=[3, 4], b=[4, 5]))
2-element Array{Float64,1}:
 2.0794415416798357
 2.302585092994046
```
"""
struct FunctionTerm{Forig,Fanon,Names} <: AbstractTerm
    forig::Forig
    fanon::Fanon
    exorig::Expr
    args_parsed::Vector
end
FunctionTerm(forig::Fo, fanon::Fa, names::NTuple{N,Symbol},
             exorig::Expr, args_parsed) where {Fo,Fa,N}  =
    FunctionTerm{Fo, Fa, names}(forig, fanon, exorig, args_parsed)
width(::FunctionTerm) = 1

"""
    InteractionTerm{Ts} <: AbstractTerm

Represents an _interaction_ between two or more individual terms.

Generated by combining multiple `AbstractTerm`s with `&` (which is what calls to
`&` in a `@formula` lower to)

# Fields

* `terms::Ts`: the terms that participate in the interaction.

# Example

```jldoctest
julia> d = (y = rand(9), a = 1:9, b = rand(9), c = repeat(["d","e","f"], 3));

julia> t = InteractionTerm(term.((:a, :b, :c)))
a(unknown) & b(unknown) & c(unknown)

julia> t == term(:a) & term(:b) & term(:c)
true

julia> t = apply_schema(t, schema(d))
a(continuous) & b(continuous) & c(DummyCoding:3→2)

julia> modelcols(t, d)
9×2 Array{Float64,2}:
 0.0      0.0
 1.09793  0.0
 0.0      2.6946
 0.0      0.0
 4.67649  0.0
 0.0      4.47245
 0.0      0.0
 0.64805  0.0
 0.0      6.97926

julia> modelcols(t.terms, d)
([1, 2, 3, 4, 5, 6, 7, 8, 9], [0.8865801492659497, 0.5489667874821704, 0.8981985570141182, 0.5043129521484462, 0.9352977047074365, 0.7454079139997376, 0.4898716849925324, 0.08100620947201143, 0.7754728346104993], [0.0 0.0; 1.0 0.0; … ; 1.0 0.0; 0.0 1.0])
```
"""
struct InteractionTerm{Ts} <: AbstractTerm
    terms::Ts
end
width(ts::InteractionTerm) = prod(width(t) for t in ts.terms)

"""
    InterceptTerm{HasIntercept} <: AbstractTerm

Represents the presence or (explicit) absence of an "intercept" term in a
regression model.  These terms are generated from [`ConstantTerm`](@ref)s in a
formula by `apply_schema(::ConstantTerm, schema, ::Type{<:StatisticalModel})`.
A `1` yields `InterceptTerm{true}`, and `0` or `-1` yield `InterceptTerm{false}`
(which explicitly omits an intercept for models which implicitly includes one
via the [`implicit_intercept`](@ref) trait).
"""
struct InterceptTerm{HasIntercept} <: AbstractTerm end
width(::InterceptTerm{H}) where {H} = H ? 1 : 0

# Typed terms

"""
    ContinuousTerm <: AbstractTerm

Represents a continuous variable, with a name and summary statistics.

# Fields

* `sym::Symbol`: The name of the variable
* `mean::T`: Mean
* `var::T`: Variance
* `min::T`: Minimum value
* `max::T`: Maximum value
"""
struct ContinuousTerm{T} <: AbstractTerm
    sym::Symbol
    mean::T
    var::T
    min::T
    max::T
end
width(::ContinuousTerm) = 1

"""
    CategoricalTerm{C,T,N} <: AbstractTerm

Represents a categorical term, with a name and [`ContrastsMatrix`](@ref)

# Fields

* `sym::Symbol`: The name of the variable
* `contrasts::ContrastsMatrix`: A contrasts matrix that captures the unique
  values this variable takes on and how they are mapped onto numerical
  predictors.
"""
struct CategoricalTerm{C,T,N} <: AbstractTerm
    sym::Symbol
    contrasts::ContrastsMatrix{C,T}
end
width(::CategoricalTerm{C,T,N}) where {C,T,N} = N

# constructor that computes the width based on the contrasts matrix
CategoricalTerm(sym::Symbol, contrasts::ContrastsMatrix{C,T}) where {C,T} =
    CategoricalTerm{C,T,length(contrasts.termnames)}(sym, contrasts)

"""
    MatrixTerm{Ts} <: AbstractTerm

A collection of terms that should be combined to produce a single numeric matrix.

A matrix term is created by [`apply_schema`](@ref) from a tuple of terms using
[`collect_matrix_terms`](@ref), which pulls out all the terms that are matrix
terms as determined by the trait function [`is_matrix_term`](@ref), which is
true by default for all `AbstractTerm`s.
"""
struct MatrixTerm{Ts<:TupleTerm} <: AbstractTerm
    terms::Ts
end
# wrap single terms in a tuple
MatrixTerm(t::AbstractTerm) = MatrixTerm((t, ))
width(t::MatrixTerm) = sum(width(tt) for tt in t.terms)

"""
    collect_matrix_terms(ts::TupleTerm)
    collect_matrix_terms(t::AbstractTerm) = collect_matrix_term((t, ))

Depending on whether the component terms are matrix terms (meaning they have
[`is_matrix_term(T) == true`](@ref is_matrix_term)), `collect_matrix_terms` will
return

1.  A single `MatrixTerm` (if all components are matrix terms)
2.  A tuple of the components (if none of them are matrix terms)
3.  A tuple of terms, with all matrix terms collected into a single `MatrixTerm`
    in the first element of the tuple, and the remaining non-matrix terms passed
    through unchanged.

By default all terms are matrix terms (that is,
`is_matrix_term(::Type{<:AbstractTerm}) = true`), the first case is by far the
most common.  The others are provided only for convenience when dealing with
specialized terms that can't be concatenated into a single model matrix, like
random effects terms in
[MixedModels.jl](https://github.com/dmbates/MixedModels.jl).

"""
function collect_matrix_terms(ts::TupleTerm)
    ismat = collect(is_matrix_term.(ts))
    if all(ismat)
        MatrixTerm(ts)
    elseif any(ismat)
        matterms = ts[ismat]
        (MatrixTerm(ts[ismat]), ts[.!ismat]...)
    else
        ts
    end
end
collect_matrix_terms(t::T) where {T<:AbstractTerm} =
    is_matrix_term(T) ? MatrixTerm((t, )) : t
collect_matrix_terms(t::MatrixTerm) = t

"""
    is_matrix_term(::Type{<:AbstractTerm})

Does this type of term get concatenated with other matrix terms into a single
model matrix?  This controls the behavior of the [`collect_matrix_terms`](@ref),
which collects all of its arguments for which `is_matrix_term` returns `true`
into a [`MatrixTerm`](@ref), and returns the rest unchanged.

Since all "normal" terms which describe one or more model matrix columns are
matrix terms, this defaults to `true` for any `AbstractTerm`.

An example of a non-matrix term is a random effect term in
[MixedModels.jl](https://github.com/dmbates/MixedModels.jl).
"""
is_matrix_term(::T) where {T} = is_matrix_term(T)
is_matrix_term(::Type{<:AbstractTerm}) = true

"""
    capture_call(f_orig::Function, f_anon::Function, argnames::NTuple{N,Symbol},
                 ex_orig::Expr, args_parsed::Vector{AbstractTerm})

When the [`@formula`](@ref) macro encounters a call to a function that's not
part of the DSL, it replaces the expression with a call to `capture_call` that
captures captures the original function called, an anonymous function that can
be applied row-wise to a table, and the arguments of the original call parsed as
if they were part of the formula.

The call to `capture_call` is generated by [`capture_call_ex!`](@ref), and just
calls passes its arguments to the [`FunctionTerm`](@ref) constructor.
"""
capture_call(args...) = FunctionTerm(args...)

extract_symbols(x) = Symbol[]
extract_symbols(x::Symbol) = [x]
extract_symbols(ex::Expr) =
    is_call(ex) ? mapreduce(extract_symbols, union, ex.args[2:end]) : Symbol[]

################################################################################
# showing terms

function Base.show(io::IO, mime::MIME"text/plain", term::AbstractTerm; prefix="")
    print(io, prefix, term)
end

function Base.show(io::IO, mime::MIME"text/plain", terms::TupleTerm; prefix=nothing)
    for t in terms
        show(io, mime, t; prefix=something(prefix, ""))
        # ensure that there are newlines in between each term after the first
        # if no prefix is specified
        prefix = something(prefix, '\n')
    end
end
Base.show(io::IO, terms::TupleTerm) = join(io, terms, " + ")

Base.show(io::IO, ::MIME"text/plain", t::Term; prefix="") =
    print(io, prefix, t.sym, "(unknown)")
Base.show(io::IO, t::Term) = print(io, t.sym)

Base.show(io::IO, t::ConstantTerm) = print(io, t.n)

Base.show(io::IO, t::FormulaTerm) = print(io, "$(t.lhs) ~ $(t.rhs)")
function Base.show(io::IO, mime::MIME"text/plain", t::FormulaTerm; prefix="")
    println(io, "FormulaTerm")
    print(io, "Response:")
    show(io, mime, t.lhs, prefix="\n  ")
    println(io)
    print(io, "Predictors:")
    show(io, mime, t.rhs, prefix="\n  ")
end

Base.show(io::IO, t::FunctionTerm) = print(io, ":($(t.exorig))")
function Base.show(io::IO, ::MIME"text/plain",
                   t::FunctionTerm{Fo,Fa,names};
                   prefix = "") where {Fo,Fa,names}
    print(io, prefix, "(")
    join(io, names, ",")
    print(io, ")->", t.exorig)
end

Base.show(io::IO, it::InteractionTerm) = join(io, it.terms, " & ")
function Base.show(io::IO, mime::MIME"text/plain", it::InteractionTerm; prefix="")
    for t in it.terms
        show(io, mime, t; prefix=prefix)
        prefix = " & "
    end
end

Base.show(io::IO, t::InterceptTerm{H}) where {H} = print(io, H ? "1" : "0")

Base.show(io::IO, t::ContinuousTerm) = print(io, t.sym)
Base.show(io::IO, ::MIME"text/plain", t::ContinuousTerm; prefix="") =
    print(io, prefix, t.sym, "(continuous)")

Base.show(io::IO, t::CategoricalTerm{C,T,N}) where {C,T,N} = print(io, t.sym)
Base.show(io::IO, ::MIME"text/plain", t::CategoricalTerm{C,T,N}; prefix="") where {C,T,N} =
    print(io, prefix, t.sym, "($C:$(length(t.contrasts.levels))→$N)")

Base.show(io::IO, t::MatrixTerm) = show(io, t.terms)
Base.show(io::IO, mime::MIME"text/plain", t::MatrixTerm; prefix="") =
    show(io, mime, t.terms, prefix=prefix)

################################################################################
# operators on Terms that create new terms:

Base.:~(lhs::TermOrTerms, rhs::TermOrTerms) = FormulaTerm(lhs, rhs)

Base.:&(terms::AbstractTerm...) = InteractionTerm(terms)
Base.:&(term::AbstractTerm) = term
Base.:&(it::InteractionTerm, terms::AbstractTerm...) = InteractionTerm((it.terms..., terms...))

Base.:+(terms::AbstractTerm...) = (unique(terms)..., )
Base.:+(as::TupleTerm, b::AbstractTerm) = (as..., b)
Base.:+(a::AbstractTerm, bs::TupleTerm) = (a, bs...)

################################################################################
# evaluating terms with data to generate model matrix entries

"""
    modelcols(t::AbstractTerm, data)

Create a numerical "model columns" representation of data based on an
`AbstractTerm`.  `data` can either be a whole table (a property-accessible
collection of iterable columns or iterable collection of property-accessible
rows, as defined by [Tables.jl](https://github.com/JuliaData/Tables.jl) or a
single row (in the form of a `NamedTuple` of scalar values).  Tables will be
converted to a `NamedTuple` of `Vectors` (e.g., a `Tables.ColumnTable`).
"""
function modelcols(t, d::D) where D
    Tables.istable(d) || throw(ArgumentError("Data of type $D is not a table!"))
    ## throw an error for t which don't have a more specific modelcols method defined
    ## TODO: this seems like it ought to be handled by dispatching on something
    ## like modelcols(::Any, ::NamedTuple) or modelcols(::AbstractTerm, ::NamedTuple)
    ## but that causes ambiguity errors or under-constrained modelcols methods for
    ## custom term types...
    d isa NamedTuple && throw(ArgumentError("don't know how to generate modelcols for " *
                                            "term $t. Did you forget to call apply_schema?"))
    modelcols(t, columntable(d))
end

"""
    modelcols(ts::NTuple{N, AbstractTerm}, data) where N

When a tuple of terms is provided, `modelcols` broadcasts over the individual
terms.  To create a single matrix, wrap the tuple in a [`MatrixTerm`](@ref).

# Example

```jldoctest
julia> d = (a = [1:9;], b = rand(9), c = repeat(["d","e","f"], 3));

julia> ts = apply_schema(term.((:a, :b, :c)), schema(d))
a(continuous)
b(continuous)
c(DummyCoding:3→2)

julia> cols = modelcols(ts, d)
([1, 2, 3, 4, 5, 6, 7, 8, 9], [0.7184176729016183, 0.4881665815778522, 0.7081609847641785, 0.7743011281211944, 0.584295963367869, 0.32493666547553657, 0.9894077965577408, 0.3331747574477202, 0.6532298571732302], [0.0 0.0; 1.0 0.0; … ; 1.0 0.0; 0.0 1.0])

julia> reduce(hcat, cols)
9×4 Array{Float64,2}:
 1.0  0.718418  0.0  0.0
 2.0  0.488167  1.0  0.0
 3.0  0.708161  0.0  1.0
 4.0  0.774301  0.0  0.0
 5.0  0.584296  1.0  0.0
 6.0  0.324937  0.0  1.0
 7.0  0.989408  0.0  0.0
 8.0  0.333175  1.0  0.0
 9.0  0.65323   0.0  1.0

julia> modelcols(MatrixTerm(ts), d)
9×4 Array{Float64,2}:
 1.0  0.718418  0.0  0.0
 2.0  0.488167  1.0  0.0
 3.0  0.708161  0.0  1.0
 4.0  0.774301  0.0  0.0
 5.0  0.584296  1.0  0.0
 6.0  0.324937  0.0  1.0
 7.0  0.989408  0.0  0.0
 8.0  0.333175  1.0  0.0
 9.0  0.65323   0.0  1.0
```
"""
modelcols(ts::TupleTerm, d::NamedTuple) = modelcols.(ts, Ref(d))

modelcols(t::Term, d::NamedTuple) =
    throw(ArgumentError("can't generate modelcols for un-typed term $t. " *
                        "Use apply_schema to create concrete terms first"))

# TODO: @generated to unroll the getfield stuff
modelcols(ft::FunctionTerm{Fo,Fa,Names}, d::NamedTuple) where {Fo,Fa,Names} =
    ft.fanon.(getfield.(Ref(d), Names)...)

modelcols(t::ContinuousTerm, d::NamedTuple) = copy.(d[t.sym])

modelcols(t::CategoricalTerm, d::NamedTuple) = t.contrasts[d[t.sym], :]

"""
    reshape_last_to_i(i::Int, a)

Reshape `a` so that its last dimension moves to dimension `i` (+1 if `a` is an
`AbstractMatrix`).
"""
reshape_last_to_i(i, a) = a
reshape_last_to_i(i, a::AbstractVector) = reshape(a, ones(Int, i-1)..., :)
reshape_last_to_i(i, a::AbstractMatrix) = reshape(a, size(a,1), ones(Int, i-1)..., :)

# an "inside out" kronecker-like product based on broadcasting reshaped arrays
# for a single row, some will be scalars, others possibly vectors.  for a whole
# table, some will be vectors, possibly some matrices
function kron_insideout(op::Function, args...)
    args = (reshape_last_to_i(i,a) for (i,a) in enumerate(args))
    vec(broadcast(op, args...))
end

function row_kron_insideout(op::Function, args...)
    rows = size(args[1], 1)
    args = (reshape_last_to_i(i,reshape(a, size(a,1), :)) for (i,a) in enumerate(args))
    # args = (reshape(a, size(a,1), ones(Int, i-1)..., :) for (i,a) in enumerate(args))
    reshape(broadcast(op, args...), rows, :)
end

# two options here: either special-case ColumnTable (named tuple of vectors)
# vs. vanilla NamedTuple, or reshape and use normal broadcasting
modelcols(t::InteractionTerm, d::NamedTuple) =
    kron_insideout(*, (modelcols(term, d) for term in t.terms)...)

function modelcols(t::InteractionTerm, d::ColumnTable)
    row_kron_insideout(*, (modelcols(term, d) for term in t.terms)...)
end

modelcols(t::InterceptTerm{true}, d::NamedTuple) = ones(size(first(d)))
modelcols(t::InterceptTerm{false}, d) = Matrix{Float64}(undef, size(first(d),1), 0)

modelcols(t::FormulaTerm, d::NamedTuple) = (modelcols(t.lhs,d), modelcols(t.rhs, d))

function modelcols(t::MatrixTerm, d::ColumnTable)
    mat = reduce(hcat, [modelcols(tt, d) for tt in t.terms])
    reshape(mat, size(mat, 1), :)
end

modelcols(t::MatrixTerm, d::NamedTuple) =
    reduce(vcat, [modelcols(tt, d) for tt in t.terms])

vectorize(x::Tuple) = collect(x)
vectorize(x::AbstractVector) = x
vectorize(x) = [x]

"""
    coefnames(term::AbstractTerm)

Return the name(s) of column(s) generated by a term.  Return value is either a
`String` or an iterable of `String`s.
"""
StatsBase.coefnames(t::FormulaTerm) = (coefnames(t.lhs), coefnames(t.rhs))
StatsBase.coefnames(::InterceptTerm{H}) where {H} = H ? "(Intercept)" : []
StatsBase.coefnames(t::ContinuousTerm) = string(t.sym)
StatsBase.coefnames(t::CategoricalTerm) =
    ["$(t.sym): $name" for name in t.contrasts.termnames]
StatsBase.coefnames(t::FunctionTerm) = string(t.exorig)
StatsBase.coefnames(ts::TupleTerm) = reduce(vcat, coefnames.(ts))
StatsBase.coefnames(t::MatrixTerm) = mapreduce(coefnames, vcat, t.terms)
StatsBase.coefnames(t::InteractionTerm) =
    kron_insideout((args...) -> join(args, " & "), vectorize.(coefnames.(t.terms))...)

################################################################################
# old Terms features:

hasintercept(f::FormulaTerm) = hasintercept(f.rhs)
hasintercept(t::TermOrTerms) =
    InterceptTerm{true}() ∈ terms(t) ||
    ConstantTerm(1) ∈ terms(t)
omitsintercept(f::FormulaTerm) = omitsintercept(f.rhs)
omitsintercept(t::TermOrTerms) =
    InterceptTerm{false}() ∈ terms(t) ||
    ConstantTerm(0) ∈ terms(t) ||
    ConstantTerm(-1) ∈ terms(t)

hasresponse(t) = false
hasresponse(t::FormulaTerm) =
    t.lhs !== nothing &&
    t.lhs !== ConstantTerm(0) &&
    t.lhs !== InterceptTerm{false}()

# convenience converters
"""
    term(x)

Wrap argument in an appropriate `AbstractTerm` type: `Symbol`s become `Term`s,
and `Number`s become `ConstantTerm`s.  Any `AbstractTerm`s are unchanged.

# Example

```jldoctest
julia> ts = term.((1, :a, :b))
1
a(unknown)
b(unknown)

julia> typeof(ts)
Tuple{ConstantTerm{Int64},Term,Term}
```
"""
term(n::Number) = ConstantTerm(n)
term(s::Symbol) = Term(s)
term(t::AbstractTerm) = t


text:
rmprocs(workers()...); addprocs(2); @assert nprocs() == 3
@everywhere using Test
@everywhere using ArrayChannels

function test_interleave()
   @testset "Order Maintained" begin
        A = ArrayChannel(Float64, procs(), 2, 2)
        pid = myid()
        id = A.rrid
        local remote
        @sync begin
            for i in 1:10
                @async begin
                    A[1,1] = i
                    put!(A, procs()[2])
                end
            end

            remote = @async remotecall_fetch(procs()[2], id) do rrid
                X = ac_get_from(rrid)
                v = Vector{Int64}()
                for i in 1:10
                    take!(X, pid)
                    push!(v, i)
                end
                return v
            end
        end
        for i in 1:10
           @test remote.result[i] == i
        end
   end
end


text:
module TestMOIWrapper

using Ipopt
using MathOptInterface
using Test

const MOI = MathOptInterface

const OPTIMIZER = Ipopt.Optimizer()
MOI.set(OPTIMIZER, MOI.Silent(), true)

# Without fixed_variable_treatment set, duals are not computed for variables
# that have lower_bound == upper_bound.
MOI.set(
    OPTIMIZER,
    MOI.RawParameter("fixed_variable_treatment"),
    "make_constraint",
)

# TODO(odow): add features to Ipopt so we can remove some of this caching.
const BRIDGED_OPTIMIZER = MOI.Bridges.full_bridge_optimizer(
    MOI.Utilities.CachingOptimizer(
        MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
        OPTIMIZER,
    ),
    Float64,
)

const CONFIG = MOI.Test.TestConfig(
    atol = 1e-4,
    rtol = 1e-4,
    optimal_status = MOI.LOCALLY_SOLVED,
    infeas_certificates = false,
)

const CONFIG_NO_DUAL = MOI.Test.TestConfig(
    atol = 1e-4,
    rtol = 1e-4,
    optimal_status = MOI.LOCALLY_SOLVED,
    infeas_certificates = false,
    duals = false,  # Don't check dual result!
)

function test_solvername()
    @test MOI.get(OPTIMIZER, MOI.SolverName()) == "Ipopt"
end

function test_supports_default_copy_to()
    @test MOI.Utilities.supports_default_copy_to(OPTIMIZER, false)
    @test !MOI.Utilities.supports_default_copy_to(OPTIMIZER, true)
end

function test_basicconstraint()
    # TODO(odow): these basic constraint tests are broken.
    # MOI.Test.basic_constraint_tests(OPTIMIZER, CONFIG)
end

function test_unittest()
    return MOI.Test.unittest(
        BRIDGED_OPTIMIZER,
        CONFIG,
        String[
            # VectorOfVariables-in-SecondOrderCone not supported
            "delete_soc_variables",
            # NumberOfThreads not supported
            "number_threads",
            # MOI.Integer not supported.
            "solve_integer_edge_cases",
            # ObjectiveBound not supported.
            "solve_objbound_edge_cases",
            # DualObjectiveValue not supported.
            "solve_result_index",
            # Returns NORM_LIMIT instead of DUAL_INFEASIBLE
            "solve_unbounded_model",
            # MOI.ZeroOne not supported.
            "solve_zero_one_with_bounds_1",
            "solve_zero_one_with_bounds_2",
            "solve_zero_one_with_bounds_3",
        ],
    )
end

function test_ConstraintDualStart()
    model = Ipopt.Optimizer()
    x = MOI.add_variables(model, 2)
    l = MOI.add_constraint(model, x[1], MOI.GreaterThan(1.0))
    u = MOI.add_constraint(model, x[1], MOI.LessThan(1.0))
    e = MOI.add_constraint(model, x[2], MOI.EqualTo(1.0))
    c = MOI.add_constraint(
        model,
        MOI.ScalarAffineFunction(MOI.ScalarAffineTerm.(1.0, x), 0.0),
        MOI.LessThan(1.5),
    )
    @test MOI.get(model, MOI.ConstraintDualStart(), l) === nothing
    @test MOI.get(model, MOI.ConstraintDualStart(), u) === nothing
    @test MOI.get(model, MOI.ConstraintDualStart(), e) === nothing
    @test MOI.get(model, MOI.ConstraintDualStart(), c) === nothing
    @test MOI.get(model, MOI.NLPBlockDualStart()) === nothing
    MOI.set(model, MOI.ConstraintDualStart(), l, 1.0)
    MOI.set(model, MOI.ConstraintDualStart(), u, -1.0)
    MOI.set(model, MOI.ConstraintDualStart(), e, -1.5)
    MOI.set(model, MOI.ConstraintDualStart(), c, 2.0)
    MOI.set(model, MOI.NLPBlockDualStart(), [1.0, 2.0])
    @test MOI.get(model, MOI.ConstraintDualStart(), l) == 1.0
    @test MOI.get(model, MOI.ConstraintDualStart(), u) == -1.0
    @test MOI.get(model, MOI.ConstraintDualStart(), e) == -1.5
    @test MOI.get(model, MOI.ConstraintDualStart(), c) == 2.0
    @test MOI.get(model, MOI.NLPBlockDualStart()) == [1.0, 2.0]
    MOI.set(model, MOI.ConstraintDualStart(), l, nothing)
    MOI.set(model, MOI.ConstraintDualStart(), u, nothing)
    MOI.set(model, MOI.ConstraintDualStart(), e, nothing)
    MOI.set(model, MOI.ConstraintDualStart(), c, nothing)
    MOI.set(model, MOI.NLPBlockDualStart(), nothing)
    @test MOI.get(model, MOI.ConstraintDualStart(), l) === nothing
    @test MOI.get(model, MOI.ConstraintDualStart(), u) === nothing
    @test MOI.get(model, MOI.ConstraintDualStart(), e) === nothing
    @test MOI.get(model, MOI.ConstraintDualStart(), c) === nothing
    @test MOI.get(model, MOI.NLPBlockDualStart()) === nothing
end

function test_contlinear()
    MOI.Test.contlineartest(
        BRIDGED_OPTIMIZER,
        CONFIG,
        String[
            # Tests requiring DualObjectiveValue. Tested below.
            "linear1",
            "linear2",
            "linear10",
            "linear14",
            # Tests requiring infeasibility certificates
            "linear8a",
            "linear8b",
            "linear8c",
            "linear12",
            # An INVALID_MODEL because it contains an empty 0 == 0 row.
            "linear15",
        ],
    )
    MOI.Test.linear1test(BRIDGED_OPTIMIZER, CONFIG_NO_DUAL)
    MOI.Test.linear2test(BRIDGED_OPTIMIZER, CONFIG_NO_DUAL)
    MOI.Test.linear10test(BRIDGED_OPTIMIZER, CONFIG_NO_DUAL)
    return MOI.Test.linear14test(BRIDGED_OPTIMIZER, CONFIG_NO_DUAL)
end

function test_qp()
    return MOI.Test.qptest(BRIDGED_OPTIMIZER, CONFIG)
end

function test_qcp()
    MOI.empty!(BRIDGED_OPTIMIZER)
    return MOI.Test.qcptest(BRIDGED_OPTIMIZER, CONFIG_NO_DUAL)
end

function test_nlptest()
    return MOI.Test.nlptest(OPTIMIZER, CONFIG)
end

function test_getters()
    return MOI.Test.copytest(
        MOI.instantiate(Ipopt.Optimizer, with_bridge_type = Float64),
        MOI.Utilities.Model{Float64}(),
    )
end

function test_boundsettwice()
    MOI.Test.set_lower_bound_twice(OPTIMIZER, Float64)
    return MOI.Test.set_upper_bound_twice(OPTIMIZER, Float64)
end

function test_nametest()
    return MOI.Test.nametest(BRIDGED_OPTIMIZER)
end

function test_validtest()
    return MOI.Test.validtest(BRIDGED_OPTIMIZER)
end

function test_emptytest()
    return MOI.Test.emptytest(BRIDGED_OPTIMIZER)
end

function test_solve_time()
    model = Ipopt.Optimizer()
    x = MOI.add_variable(model)
    @test isnan(MOI.get(model, MOI.SolveTime()))
    MOI.optimize!(model)
    @test MOI.get(model, MOI.SolveTime()) > 0.0
end

# Model structure for test_check_derivatives_for_naninf()
struct Issue136 <: MOI.AbstractNLPEvaluator end
MOI.initialize(::Issue136, ::Vector{Symbol}) = nothing
MOI.features_available(d::Issue136) = [:Grad, :Jac]
MOI.eval_objective(::Issue136, x) = x[1]
MOI.eval_constraint(::Issue136, g, x) = (g[1] = x[1]^(1 / 3))
MOI.eval_objective_gradient(::Issue136, grad_f, x) = (grad_f[1] = 1.0)
MOI.jacobian_structure(::Issue136) = Tuple{Int64,Int64}[(1, 1)]
function MOI.eval_constraint_jacobian(::Issue136, J, x)
    J[1] = (1 / 3) * x[1]^(1 / 3 - 1)
    return
end

function test_check_derivatives_for_naninf()
    model = Ipopt.Optimizer()
    x = MOI.add_variable(model)
    MOI.set(
        model,
        MOI.NLPBlock(),
        MOI.NLPBlockData(MOI.NLPBoundsPair.([-Inf], [0.0]), Issue136(), false),
    )
    # Failure to set check_derivatives_for_naninf="yes" may cause Ipopt to
    # segfault or return a NUMERICAL_ERROR status. Check that it is set to "yes"
    # by obtaining an INVALID_MODEL status.
    # MOI.set(model, MOI.RawParameter("check_derivatives_for_naninf"), "no")
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.INVALID_MODEL
end

function test_deprecation()
    model = Ipopt.Optimizer(print_level = 0)
    @test MOI.get(model, MOI.RawParameter("print_level")) == 0
end

function test_callback()
    model = Ipopt.Optimizer()
    MOI.set(model, MOI.RawParameter("print_level"), 0)
    x = MOI.add_variable(model)
    MOI.add_constraint(model, MOI.SingleVariable(x), MOI.GreaterThan(1.0))
    MOI.set(model, MOI.ObjectiveSense(), MOI.MIN_SENSE)
    f = MOI.ScalarAffineFunction([MOI.ScalarAffineTerm(1.0, x)], 0.5)
    MOI.set(model, MOI.ObjectiveFunction{typeof(f)}(), f)
    x_vals = Float64[]
    function my_callback(
        prob::IpoptProblem,
        alg_mod::Cint,
        iter_count::Cint,
        obj_value::Float64,
        inf_pr::Float64,
        inf_du::Float64,
        mu::Float64,
        d_norm::Float64,
        regularization_size::Float64,
        alpha_du::Float64,
        alpha_pr::Float64,
        ls_trials::Cint,
    )
        c = Ipopt.column(x)
        push!(x_vals, prob.x[c])
        @test isapprox(obj_value, 1.0 * x_vals[end] + 0.5, atol = 1e-1)
        return iter_count < 1
    end
    MOI.set(model, Ipopt.CallbackFunction(), my_callback)
    MOI.optimize!(model)
    @test MOI.get(model, MOI.TerminationStatus()) == MOI.INTERRUPTED
    @test length(x_vals) == 2
end

function test_empty_optimize()
    model = Ipopt.Optimizer()
    err = ErrorException(
        "IPOPT: Failed to construct problem because there are 0 variables. " *
        "If you intended to construct an empty problem, one work-around is " *
        "to add a variable fixed to 0.",
    )
    @test_throws err MOI.optimize!(model)
end

end  # module TestMOIWrapper

runtests(TestMOIWrapper)


text:
#___  ____       ___ _____   _   _       _ _
#|  \/  | |     |_  /  __ \ | | | |     (_| |
#| .  . | |       | | /  \/ | | | |_ __  _| |_ ___
#| |\/| | |       | | |     | | | | '_ \| | __/ _ \
#| |  | | |___/\__/ | \__/\ | |_| | | | | | || (_) |
#_______\_____\____/ \____/  _____|_|___|_____\_____ _____ _____
#| ___ \        (_)         | | \ \ / / / __  |  _  / __  |  _  |
#| |_/ _ __ ___  _  ___  ___| |_ \ V /  `' / /| |/' `' / /| |/' |
#|  __| '__/ _ \| |/ _ \/ __| __|/   \    / / |  /| | / / |  /| |
#| |  | | | (_) | |  __| (__| |_/ /^\ \ ./ /__\ |_/ ./ /__\ |_/ /
#\_|  |_|  \___/| |\___|\___|\__\/   \/ \_____/\___/\_____/\___/
#              _/ |
#             |__/
#
# This code is part of the proposal of the team "MLJC UniTo" - University of Turin
# for "ProjectX 2020" Climate Change for AI.
# The code is licensed under MIT 3.0
# Please read readme or comments for credits and further information.

# Compiler: Julia 1.5

# update NeuralPDE#quadrature_training
using NeuralPDE
using Quadrature, Cubature, Cuba
using Flux, ModelingToolkit, GalacticOptim, Optim, DiffEqFlux
using Plots
using PyPlot
using CUDA
using CUDAdrv

print("Precompiling Done")

# the example is taken from this article https://arxiv.org/abs/1910.10503
@parameters x θ
@variables p(..)
@derivatives Dx'~x
@derivatives Dxx''~x

#2D PDE
α = 0.3
β = 0.5
_σ = 0.5
# Discretization
dx = 0.05
# here we use normalization condition: dx*p(x,θ) ~ 1 in order to get a non-zero solution.
eq  = [(α - 3*β*x^2)*p(x,θ) + (α*x - β*x^3)*Dx(p(x,θ)) ~ (_σ^2/2)*Dxx(p(x,θ)),
       dx*p(x,θ) ~ 1.]

# Initial and boundary conditions
bcs = [p(-2.2,θ) ~ 0. ,p(2.2,θ) ~ 0. , p(-2.2,θ) ~ p(2.2,θ)]

# Space and time domains
domains = [x ∈ IntervalDomain(-2.2,2.2)]

# Neural network
chain = Chain(Dense(1,12,Flux.σ),Dense(12,12,Flux.σ),Dense(12,1)) |>gpu

discretization = NeuralPDE.PhysicsInformedNN(dx,
                                             chain,
                                             strategy= NeuralPDE.GridTraining()) |>gpu

pde_system = PDESystem(eq,bcs,domains,[x],[p])
prob = NeuralPDE.discretize(pde_system,discretization)

cb = function (p,l)
    println("Current loss is: $l")
    return false
end
CUDA.allowscalar(false)
res = GalacticOptim.solve(prob, ADAM(0.1); cb = cb, maxiters=8000)
phi = discretization.phi

analytic_sol_func(x) = 28.022*exp((1/(2*_σ^2))*(2*α*x^2 - β*x^4))

xs = [domain.domain.lower:dx:domain.domain.upper for domain in domains][1]
u_real  = [analytic_sol_func(x) for x in xs]
u_predict  = [first(phi(x,res.minimizer)) for x in xs]

plot(xs ,u_real, label = "analytic")
plot!(xs ,u_predict, label = "predict")


text:
using CMake
builddir = joinpath(@__DIR__, "build")
if !isdir(builddir)
    Base.Filesystem.mkdir(builddir)
end
mercurylib_dir = joinpath(@__DIR__, "src")
config_cmd = `$cmake -S$mercurylib_dir -B$builddir`
build_cmd = `$cmake --build $builddir`
run(config_cmd)
run(build_cmd)


text:
macro pricecharts()
  println("")
  reload(Pkg.dir("PriceCharts/run_tests.jl"))
end


text:
module PolicyResponse

using ClimateMARGO.Models
using ClimateMARGO.Diagnostics

export step_forward!, add_emissions_bump!

include("stepping_forward.jl")

end


text:
# AAA algorithm from the paper "The AAA Alorithm for Rational Approximation"
# by Y. Nakatsukasa, O. Sete, and L.N. Trefethen, SIAM Journal on Scientific Computing
# 2018

struct AAAapprox{T <: AbstractArray} <: BRInterp
    x::T
    f::T
    w::T
    errvec::T
end

(a::AAAapprox)(zz) = reval(zz, a.x, a.f, a.w)

# Handle function inputs as well
function aaa(Z::AbstractArray{T,1}, F::S;  tol=1e-13, mmax=100, verbose=false, clean=true) where {T, S<:Function}
    aaa(Z, F.(Z), tol=tol, mmax=mmax, verbose=verbose, clean=clean)
end

"""aaa  rational approximation of data F on set Z
        r = aaa(Z, F; tol, mmax, verbose, clean)

 Input: Z = vector of sample points
        F = vector of data values at the points in Z
        tol = relative tolerance tol, set to 1e-13 if omitted
        mmax: max type is (mmax-1, mmax-1), set to 100 if omitted
        verbose: print info while calculating default = false
        clean: detect and remove Froissart doublets default = true

 Output: r = an AAA approximant as a callable struct with fields
         z, f, w = vectors of support pts, function values, weights
         errvec = vector of errors at each step

 NOTE: changes from matlab version:
         switched order of Z and F in function signature
         added verbose and clean boolean flags
         pol, res, zer = vectors of poles, residues, zeros are now only calculated on demand
         by calling prz(z::AAAapprox)
"""
function aaa(Z::AbstractArray{T,1}, F::AbstractArray{S,1}; tol=1e-13, mmax=100,
             verbose=false, clean=true) where {S, T}
    # filter out any NaN's or Inf's in the input
    keep = isfinite.(F)
    F = F[keep]
    Z = Z[keep]

    M = length(Z)                    # number of sample points
    mmax = min(M, mmax)              # max number of support points

    reltol = tol*norm(F, Inf)
    SF = spdiagm(M, M, 0 => F)       # left scaling matrix

    F, Z = promote(F, Z)
    P = promote_type(S, T)

    J = [1:M;]
    z = P[]                          # support points
    f = P[]                          # function values at support points
    C = P[]
    w = P[]

    errvec = P[]
    R = F .- mean(F)
    @inbounds for m = 1:mmax
        j = argmax(abs.(F .- R))               # select next support point
        push!(z, Z[j])
        push!(f, F[j])
        deleteat!(J, findfirst(isequal(j), J))   # update index vector

        # next column of Cauchy matrix
        C = isempty(C) ? reshape((1 ./ (Z .- Z[j])), (M,1)) : [C (1 ./ (Z .- Z[j]))]

        Sf = diagm(f)                         # right scaling matrix
        A = SF * C - C * Sf                   # Loewner matrix
        G = svd(A[J, :])

        # A[J, :] might not have full rank, so svd can return fewer than m columns
        #w = G.V[:, m]
        w = G.V[:, end]                       # weight vector = min sing vector

        N = C * (w .* f)                      # numerator
        D = C * w
        R .= F
        R[J] .= N[J] ./ D[J]                  # rational approximation

        err = norm(F - R, Inf)
        verbose && println("Iteration: ", m, "  err: ", err)
        errvec = [errvec; err]                # max error at sample points
        err <= reltol && break                # stop if converged
    end
    r = AAAapprox(z, f, w, errvec)

    # remove Frois. doublets if desired.  We do this in place
    if clean
        pol, res, zer = prz(r)            # poles, residues, and zeros
        ii = findall(abs.(res) .< 1e-13)  # find negligible residues
        length(ii) != 0 && cleanup!(r, pol, res, zer, Z, F)
    end
    return r
end

function prz(r::AAAapprox)
    z, f, w = r.x, r.f, r.w
    m = length(w)
    B = diagm(ones(m+1))
    B[1, 1] = 0.0
    E = [0.0  transpose(w); ones(m) diagm(z)]
    pol, _ = eigen(E, B)
    pol = pol[isfinite.(pol)]
    dz = 1e-5 * exp.(2im*pi*[1:4;]/4)

    # residues
    res = r(pol .+ transpose(dz)) * dz ./ 4

    E = [0 transpose(w .* f); ones(m) diagm(z)]
    zer, _ = eigen(E, B)
    zer = zer[isfinite.(zer)]
    pol, res, zer
end

function reval(zz, z, f, w)
    # evaluate r at zz
    zv = size(zz) == () ? [zz] : vec(zz)
    CC = 1.0 ./ (zv .- transpose(z))           # Cauchy matrix
    r = (CC * (w .* f)) ./ (CC * w)            # AAA approx as vector
    r[isinf.(zv)] .= sum(f .* w) ./ sum(w)

    ii = findall(isnan.(r))                    # find values NaN = Inf/Inf if any
    @inbounds for j in ii
        if !isnan(zv[j]) && ((v = findfirst(isequal(zv[j]), z)) !== nothing)
            r[j] = f[v]  # force interpolation there
        end
    end
    r = size(zz) == () ? r[1] : reshape(r, size(zz))             # the AAA approximation
end

# Only calculate the updated z, f, and w
function cleanup!(r, pol, res, zer, Z, F)
    z, f, w = r.x, r.f, r.w
    m = length(z)
    M = length(Z)
    ii = findall(abs.(res) .< 1e-13)  # find negligible residues
    ni = length(ii)
    ni == 0 && return
    println("$ni Froissart doublets. Number of residues = ", length(res))

    @inbounds for j = 1:ni
        azp = abs.(z .- pol[ii[j]] )
        jj = findall(isequal(minimum(azp)), azp)
        deleteat!(z, jj)    # remove nearest support points
        deleteat!(f, jj)
    end

    @inbounds for j = 1:length(z)
        jj = findall(isequal(z[j]), Z)
        deleteat!(F, jj)
        deleteat!(Z, jj)
    end
    m = m - length(ii)
    SF = spdiagm(M-m, M-m, 0 => F)
    Sf = diagm(f)
    C = 1 ./ (Z .- transpose(z))
    A = SF*C - C*Sf
    G = svd(A)
    w[:] .= G.V[:, m]
    println("cleanup: ", size(z), "  ", size(f), "  ", size(w))
    return nothing
end


text:
module DeepQLearning

using Random
using StatsBase
using Printf
using Parameters
using Flux
using BSON
using POMDPModelTools
using POMDPPolicies
using POMDPLinter
using LinearAlgebra
using TensorBoardLogger: TBLogger, log_value
using EllipsisNotation

using CommonRLInterface: AbstractEnv, reset!, actions, observe, act!, terminated
using POMDPs

export DeepQLearningSolver,
       AbstractNNPolicy,
       NNPolicy,
       getnetwork,
       resetstate!,
       DQExperience,
       restore_best_model,
       PrioritizedReplayBuffer,
       EpisodeReplayBuffer,

       # helpers
       flattenbatch,
       huber_loss,
       isrecurrent,
       batch_trajectories

include("helpers.jl")
include("policy.jl")
include("exploration_policy.jl")
include("evaluation_policy.jl")
include("prioritized_experience_replay.jl")
include("episode_replay.jl")
include("dueling.jl")
include("solver.jl")

end # module DeepQLearning


text:
function reach_ASB07_decomposed!(R::Vector{<:ReachSet},
                                 Ωhat0::Vector{<:LazySet},
                                 U::Union{ConstantInput, Nothing},
                                 ϕ::AbstractMatrix{IntervalArithmetic.Interval{NUM}},
                                 N::Int,
                                 δ::Float64,
                                 max_order::Int,
                                 n::Int,
                                 partition,
                                 blocks) where {NUM}
    # initial reach set
    t0, t1 = zero(δ), δ
    R[1] = ReachSet(CartesianProductArray(Ωhat0[blocks]), t0, t1)

    b = length(blocks)
    Rₖ_array = Vector{LazySet{NUM}}(undef, b)
    ϕpowerk = copy(ϕ)

    if U != nothing
        Whatk = Vector{Zonotope{NUM}}(undef, b)
        inputs = next_set(U)
        @inbounds for i in 1:b
            bi = partition[blocks[i]]
            Whatk[i] = linear_map(proj(bi, n), inputs)
        end
    end

    k = 2
    while k <= N
        for i in 1:b
            bi = partition[blocks[i]]
            Rₖ_bi = ZeroSet(length(bi))
            for (j, bj) in enumerate(partition)
                block = IntervalMatrix(ϕpowerk[bi, bj])
                if !iszero(block)
                    Rₖ_bij = overapproximate(block * Ωhat0[j], Zonotope)
                    Rₖ_bi = minkowski_sum(Rₖ_bi, Rₖ_bij)
                end
            end
            if U != nothing
                Rₖ_bi = minkowski_sum(Rₖ_bi, Whatk[i])
            end
            Rₖ_bi = reduce_order(Rₖ_bi, max_order)  # reduce order
            Rₖ_array[i] = Rₖ_bi
        end
        Rₖ = CartesianProductArray(copy(Rₖ_array))

        # store reach set
        t0 = t1
        t1 += δ
        R[k] = ReachSet(Rₖ, t0, t1)

        if U != nothing
            for i in 1:b
                bi = partition[blocks[i]]
                block_row = IntervalMatrix(row(ϕpowerk, bi))
                W_new = overapproximate(block_row * inputs, Zonotope)
                Whatk[i] = minkowski_sum(Whatk[i], W_new)
            end
        end

        ϕpowerk *= ϕ

        k += 1
    end
    return R
end


text:
if !("." in LOAD_PATH)
    push!(LOAD_PATH, ".")
end
import Thermodynamics
import RootSolvers
import CLIMAParameters
import Coverage
import Plots

# Packages to monitor

mod_dir(x) = dirname(dirname(pathof(x)))
all_dirs_to_monitor = [
    ".",
    mod_dir(Thermodynamics),
    mod_dir(RootSolvers),
    mod_dir(CLIMAParameters),
]

# https://github.com/jheinen/GR.jl/issues/278#issuecomment-587090846
ENV["GKSwstype"] = "nul"

constructors = ["ρeq", "pθq", "pTq"]

allocs = Dict()
for constructor in constructors
    ENV["ALLOCATION_CONSTRUCTOR"] = constructor
    run(`julia --project=test/ --track-allocation=all perf/alloc_per_constructor.jl`)

    allocs[constructor] = Coverage.analyze_malloc(all_dirs_to_monitor)

    # Clean up files
    for d in all_dirs_to_monitor
        all_files = [
            joinpath(root, f)
            for (root, dirs, files) in Base.Filesystem.walkdir(d)
            for f in files
        ]
        all_mem_files = filter(x -> endswith(x, ".mem"), all_files)
        for f in all_mem_files
            rm(f)
        end
    end
end

@info "Post-processing allocations"

function plot_allocs(constructor, allocs_per_case, n_unique_bytes)
    p = Plots.plot()
    @info "Allocations for $constructor"

    function filename_only(fn)
        fn = first(split(fn, ".jl")) * ".jl"
        splitby = "central/scratch/climaci/thermodynamics-ci/"
        if occursin(splitby, fn)
            fn = last(split(fn, splitby))
        end
        splitby = "/thermodynamics-ci/"
        if occursin(splitby, fn)
            fn = joinpath("Thermodynamics.jl", last(split(fn, splitby)))
        end
        splitby = "/depot/cpu/packages/"
        if occursin(splitby, fn)
            fn = last(split(fn, splitby))
        end
        return fn
    end
    function compile_pkg(fn, linenumber)
        c1 = endswith(filename_only(fn), "Thermodynamics.jl")
        c2 = linenumber == 1
        return c1 && c2
    end

    filter!(x -> x.bytes ≠ 0, allocs_per_case)
    filter!(x -> !compile_pkg(x.filename, x.linenumber), allocs_per_case)

    for alloc in allocs_per_case
        println(alloc)
    end
    println("Number of allocating sites: $(length(allocs_per_case))")
    case_bytes = getproperty.(allocs_per_case, :bytes)[end:-1:1]
    case_filename = getproperty.(allocs_per_case, :filename)[end:-1:1]
    case_linenumber = getproperty.(allocs_per_case, :linenumber)[end:-1:1]
    all_bytes = Int[]
    filenames = String[]
    linenumbers = Int[]
    loc_ids = String[]
    for (bytes, filename, linenumber) in
        zip(case_bytes, case_filename, case_linenumber)
        compile_pkg(filename, linenumber) && continue # Skip loading module
        loc_id = "$(filename_only(filename))" * "$linenumber"
        if !(bytes in all_bytes) && !(loc_id in loc_ids)
            push!(all_bytes, bytes)
            push!(filenames, filename)
            push!(linenumbers, linenumber)
            push!(loc_ids, loc_id)
            if length(all_bytes) ≥ n_unique_bytes
                break
            end
        end
    end

    all_bytes = all_bytes ./ 10^3
    if isempty(all_bytes)
        @info "$constructor: 0 allocations! 🎉"
        return nothing
    end
    max_bytes = maximum(all_bytes)
    @info "$constructor: $all_bytes"
    xtick_name(filename, linenumber) = "$filename, line number: $linenumber"
    markershape = (:square, :hexagon, :circle, :star, :utriangle, :dtriangle)
    for (bytes, filename, linenumber) in zip(all_bytes, filenames, linenumbers)
        Plots.plot!(
            [0],
            [bytes];
            seriestype = :scatter,
            label = xtick_name(filename_only(filename), linenumber),
            markershape = markershape[1],
            markersize = 1 + bytes / max_bytes * 10,
        )
        markershape = (markershape[end], markershape[1:(end - 1)]...)
    end
    p1 = Plots.plot!(ylabel = "Allocations (KB)", title = constructor)
    subset_allocs_per_case = collect(Base.Iterators.Reverse(allocs_per_case))
    p2 = Plots.plot(
        1:length(subset_allocs_per_case),
        getproperty.(subset_allocs_per_case, :bytes) ./ 1000;
        xlabel = "i-th allocating line (truncated and sorted)",
        ylabel = "Allocations (KB)",
        markershape = :circle,
    )
    Plots.plot(p1, p2, layout = Plots.grid(2, 1))
    Plots.savefig(joinpath(folder, "allocations_$constructor.png"))
end

folder = "perf/allocations_output"
mkpath(folder)

@info "Allocated bytes for constructor broadcasted across tested profiles"
for constructor in constructors
    plot_allocs(constructor, allocs[constructor], 10)
end


text:
export AbstractEnvironmentModel, AbstractSampleBasedModel, AbstractDistributionBasedModel

"""
Describe how to model a reinforcement learning environment.

See also [`AbstractDistributionBasedModel`](@ref), [`AbstractSampleBasedModel`](@ref).
"""
abstract type AbstractEnvironmentModel end

"""
A collection of models that can be used to sample transitions.
"""
abstract type AbstractSampleBasedModel <: AbstractEnvironmentModel end

"""
A collection of models that can get the distribution given a state and an actioi.
"""
abstract type AbstractDistributionBasedModel <: AbstractEnvironmentModel end


text:
include("../lib/quantum_utils.jl")
using DifferentialEquations
using Random

#######################################################################################
#
# define belavkin equation problem with an ordinary algorithm
#   T: final time
#   ψ0_qo: initial ket
#   A,B,C,D: tensors for Hamiltonian
#   vec_EF: vector of tensors defining Lindblad operators
#
function Belavkin_problem(T::Float64,ψ0::Ket,A::Array{Float64},B1::Array{Float64},B2::Array{Float64},C::Array{Float64},D::Array{Float64},vec_EF::Array{Array{Complex{Float64}},2};seed_input=0)

    Nfocks = ψ0.basis.shape .-1
    N = length(Nfocks)
    as_qo = a_vectors(Nfocks)
    adags_qo = dagger.(as_qo);
    z_qo = vcat((as_qo+adags_qo),(as_qo-adags_qo)/im)/√2;
    zz_qo = [(zi*zj+zj*zi)/2 for zi in z_qo, zj in z_qo];
    Ls = get_operator(vec_EF,z_qo)
    ψ = ground_state(Nfocks)
    dψ = ground_state(Nfocks)

    A′= Array{Complex{Float64}}(undef,2N);
    B′= Array{Complex{Float64}}(undef,2N,2N);
    C′= Array{Complex{Float64}}(undef,2N,2N,2N);
    D′= Array{Complex{Float64}}(undef,2N,2N,2N,2N);

    function f_det(du,u,p,t)
        convert_realvec2complexvec!(ψ.data,u)
        μ₁ = [real(expect(zi_qo,ψ)) for zi_qo in z_qo]
        μ₂ = [real(expect(zzij_qo,ψ)) for zzij_qo in zz_qo]
        B = B1 + B2 * t/T
        get_effective_Hamiltonian_tensors!(A′,B′,C′,D′,A,B,C,D,vec_EF,μ₁,μ₂)
        H = get_operator([A′,B′,C′,D′],z_qo)
        mul!(dψ.data,-im*H.data,ψ.data)
        dif_norm = real(dagger(ψ) * dψ)
        dψ.data -= dif_norm*ψ.data
        convert_complexvec2realvec!(du,dψ.data)
    end

    function f_sto(du,u,p,t)
        convert_realvec2complexvec!(ψ.data,u)
        for i in 1:length(Ls)
            mul!(dψ.data,Ls[i].data,ψ.data)
            dif_norm = real(dagger(ψ) * dψ)
            dψ.data -= dif_norm*ψ.data
            convert_complexvec2realvec!(view(du,:,i),dψ.data)
        end
    end

    u0 = Array{Float64}(undef,2*length(ψ0.data))
    convert_complexvec2realvec!(u0,ψ0.data)

    if seed_input==0
        return SDEProblem(f_det,f_sto,u0,(0.0,T),noise_rate_prototype=zeros(2*length(ψ0.data),size(vec_EF)[1]))
    else
        return SDEProblem(f_det,f_sto,u0,(0.0,T),noise_rate_prototype=zeros(2*length(ψ0.data),size(vec_EF)[1]),seed=seed_input)
    end
end

#######################################################################################
#
# define Schrodinger equation problem with an ordinary algorithm
#   T: final time
#   ψ0_qo: initial ket
#   A,B,C,D: tensors for Hamiltonian
#
#
function Schrodinger_problem(T::Float64,ψ0_qo::Ket,A::Array{Float64},B1::Array{Float64},B2::Array{Float64},C::Array{Float64},D::Array{Float64})

    Nfocks = ψ0_qo.basis.shape .-1
    N = length(Nfocks)
    as_qo = a_vectors(Nfocks)
    adags_qo = dagger.(as_qo);
    z_qo = vcat((as_qo+adags_qo),(as_qo-adags_qo)/im)/√2;
    zz_qo = [(zi*zj+zj*zi)/2 for zi in z_qo, zj in z_qo];

    ψ = Array{Complex{Float64}}(undef,length(ψ0_qo.data))
    dψ = Array{Complex{Float64}}(undef,length(ψ0_qo.data))

    function f_det(du,u,p,t)
        convert_realvec2complexvec!(ψ,u)
        B = B1 + B2 * t/T
        H = get_operator([A,B,C,D],z_qo)
        mul!(dψ,-im*H.data,ψ)
        convert_complexvec2realvec!(du,dψ)
    end

    u0 = Array{Float64}(undef,2*length(ψ0_qo.data))
    convert_complexvec2realvec!(u0,ψ0_qo.data)
    return ODEProblem(f_det,u0,(0.0,T))
end

function get_hist_ϕ(sol::Any,Nfocks::Array{Int64,1},Nsave::Int,T::Float64)
    N = length(Nfocks)

    ϕ_hist = Array{Any}(undef,Nsave)

    for i in 1:Nsave
        ϕ_hist[i] = ground_state(Nfocks)
        state_vec = sol(T*i/Nsave)
        ϕ_hist[i].data[:] .= state_vec[1:length(ϕ_hist[i].data)] + im * state_vec[1+length(ϕ_hist[i].data):end]
    end
    return ϕ_hist
end

max_stars_repo_name: ash-y/QuantumCOMFrame

max_issues_repo_name: ash-y/QuantumCOMFrame

max_forks_repo_name: ash-y/QuantumCOMFrame


text:
"""
    ConfigTypes

Module containing CLIMA configuration types.
"""
module ConfigTypes

export CLIMAConfigType,
       AtmosLESConfigType,
       AtmosGCMConfigType,
       OceanBoxGCMConfigType

abstract type CLIMAConfigType end
struct AtmosLESConfigType <: CLIMAConfigType end
struct AtmosGCMConfigType <: CLIMAConfigType end
struct OceanBoxGCMConfigType <: CLIMAConfigType end

end


text:
module Spider

using CommonMark

include("util.jl")
include("run_spider.jl")

include("plugins/bibtex.jl")
include("plugins/toc.jl")
include("plugins/mathjax.jl")
include("plugins/wikilinks.jl")
include("plugins/arxivlinks.jl")
include("plugins/backlinks.jl")
include("plugins/github_edit_link.jl")

end


text:
module JGE

using GLFW, VulkanCore.LibVulkan
import Base: reverse

include("utils/containers.jl")
include("utils/structures.jl")

include("core/glfw.jl")
include("core/keymaps.jl")

include("vulkan/base_extensions.jl")
include("vulkan/defaults.jl")
include("vulkan/common.jl")
include("vulkan/validation.jl")
include("vulkan/init.jl")
include("vulkan/presentation.jl")
include("vulkan/queues.jl")
include("vulkan/logical_device.jl")

function main(;context::AbstractContext = VulkanContext())
    glfw = true
    debug = true
        # run_window((1000, 720); key_callback = main, context = context)
    window = create_window((1000, 720), "JGE", context)
    cinst, physical_device = initialize(; glfw = glfw, debug = debug)
    surface_ptr = create_surface(cinst, window, glfw = glfw)
    @info "Surface created"
    avail_qf = available_queue_families(physical_device)

        # # no queue family is found with this predicate
        # select_queue_predicate = qf->(Bool(qf.queueFlags & VK_QUEUE_GRAPHICS_BIT) && has_presentation_support(physical_device, findfirst(q->q == qf, avail_qf), surface_ptr))
        # queue_family_index = select_queue_family(select_queue_predicate, avail_qf)

    queue_family_index = select_queue_family(physical_device)
    cqueue_info = create_queue_info(queue_family_index)
    cqueue_infos = Container([cqueue_info[]], cqueue_info.deps)
    @info "Creating logical device"
    cdevice = create_logical_device(physical_device, cqueue_infos)
    @info "Everything ready!"
    run_window(window, window_loop(context, ()->nothing), key_callback = main_keymap)
    return nothing
    return cinst, cdevice # keep them alive until the end
end

function main_gl(; context::AbstractContext = VulkanContext())
    run_window((1000, 720), context = context, key_callback = main_keymap)
end

export main, main_gl, OpenGLContext, VulkanContext, Container, unsafe_pointer, initialize

end # module


text:
using Oceananigans
using Oceananigans.Units
using Oceananigans.Simulations: reset!
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid, GridFittedBottom
using Printf
using JLD2
using GLMakie
using SeawaterPolynomials.TEOS10
using Oceananigans.Operators: Δx, Δy, ζ₃ᶠᶠᶜ

using Oceananigans: fields
using Oceananigans.Grids: ynode
using Oceananigans.Utils: with_tracers
using Oceananigans.TurbulenceClosures: validate_closure

arch = CPU()
filename = "idealized_near_global_one_degree"

include("one_degree_artifacts.jl")
# bathymetry_path = download_bathymetry() # not needed because we uploaded to repo
bathymetry = jldopen(bathymetry_path)["bathymetry"]

include("one_degree_interface_heights.jl")
z = one_degree_interface_heights()

underlying_grid = LatitudeLongitudeGrid(arch;
                                        size = (360, 150, 48),
                                        halo = (4, 4, 4),
                                        latitude = (-75, 75),
                                        z,
                                        longitude = (-180, 180),
                                        precompute_metrics = true)

grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bathymetry))

@show grid

#####
##### Closures
#####

include("variable_biharmonic_diffusion_coefficient.jl") # defines VariableBiharmonicDiffusionCoefficient
vitd = VerticallyImplicitTimeDiscretization()
background_vertical_diffusivity = VerticalScalarDiffusivity(vitd, ν=1e-2, κ=1e-4)
dynamic_vertical_diffusivity = RiBasedVerticalDiffusivity()

function idealized_one_degree_closure(; νh = (100kilometers)^2 / 1day,
                                        κ_skew = 0,
                                        κ_symmetric = κ_skew,
                                        biharmonic_time_scale = 1day)

    ν₄ = VariableBiharmonicDiffusionCoefficient(biharmonic_time_scale)
    biharmonic_viscosity = HorizontalScalarBiharmonicDiffusivity(ν=ν₄, discrete_form=true)

    gent_mcwilliams_diffusivity = IsopycnalSkewSymmetricDiffusivity(κ_skew = κ_skew,
                                                                    κ_symmetric = κ_symmetric,
                                                                    slope_limiter = FluxTapering(1e-2))

    horizontal_diffusivity = HorizontalScalarDiffusivity(ν=νh, κ=νh)

    closure_tuple = (gent_mcwilliams_diffusivity,
                     biharmonic_viscosity,
                     dynamic_vertical_diffusivity,
                     horizontal_diffusivity,
                     background_vertical_diffusivity)

    closure_tuple = with_tracers(tuple(:T), closure_tuple)
    closure_tuple = validate_closure(closure_tuple)

    return closure_tuple
end

free_surface = ImplicitFreeSurface(solver_method=:HeptadiagonalIterativeSolver)
equation_of_state = LinearEquationOfState()
buoyancy = SeawaterBuoyancy(; equation_of_state, constant_salinity=35.0)

@inline T_reference(φ) = max(0.0, 30.0 * cos(1.2 * π * φ / 180))

@inline function T_relaxation(i, j, grid, clock, fields, tᵣ)
    φ = ynode(Center(), j, grid)
    return @inbounds 1 / tᵣ * (fields.T[i, j, grid.Nz] - T_reference(φ))
end

T_top_bc = FluxBoundaryCondition(T_relaxation, discrete_form=true, parameters=30days)
T_bcs = FieldBoundaryConditions(top=T_top_bc)

@inline surface_stress_x(φ, τ₀, τˢ, τᴺ) = τ₀ * (1 - exp(-φ^2 / 200)) - (τ₀ + τˢ) * exp(-(φ + 50)^2 / 200) -
                                                                       (τ₀ + τᴺ) * exp(-(φ - 50)^2 / 200)

@inline function surface_stress_x(i, j, grid, clock, fields, p)
    φ = ynode(Center(), j, grid)
    return surface_stress_x(φ, p.τ₀, p.τˢ, p.τᴺ)
end

u_top_bc = FluxBoundaryCondition(surface_stress_x, discrete_form=true, parameters=(τ₀=6e-5, τˢ=2e-4, τᴺ=5e-5))
u_bcs = FieldBoundaryConditions(top=u_top_bc)

model = HydrostaticFreeSurfaceModel(; grid, free_surface, buoyancy,
                                    momentum_advection = VectorInvariant(),
                                    coriolis = HydrostaticSphericalCoriolis(),
                                    tracers = :T,
                                    closure = idealized_one_degree_closure(),
                                    boundary_conditions = (u=u_bcs, T=T_bcs),
                                    tracer_advection = WENO5(grid=underlying_grid))

@show model

function T_initial(λ, φ, z)
    H_stratification = 2000
    T_bottom = 0.0
    T_surface = T_reference(φ)
    dTdz = (T_surface - T_bottom) / H_stratification
    return max(T_bottom, T_surface + z * dTdz)
end

set!(model, T=T_initial)

simulation = Simulation(model; Δt=10minutes, stop_time=10years)

wall_clock = Ref(time_ns())
function progress(sim)
    elapsed = 1e-9 * (time_ns() - wall_clock[])

    u, v, w = sim.model.velocities
    T = sim.model.tracers.T
    η = sim.model.free_surface.η
    umax = maximum(abs, u)
    vmax = maximum(abs, v)
    wmax = maximum(abs, w)
    ηmax = maximum(abs, η)
    Tmax = maximum(T)
    Tmin = minimum(T)

    msg1 = @sprintf("Iteration: %d, time: %s, wall time: %s", iteration(sim), prettytime(sim), prettytime(elapsed))
    msg2 = @sprintf("├── max(u): (%.2e, %.2e, %.2e) m s⁻¹", umax, vmax, wmax)
    msg3 = @sprintf("├── extrema(T): (%.2f, %.2f) ᵒC", Tmin, Tmax)
    msg4 = @sprintf("└── max|η|: %.2e m", ηmax)

    @info string(msg1, '\n', msg2, '\n', msg3, '\n', msg4)

    wall_clock[] = time_ns()

    return nothing
end

simulation.callbacks[:p] = Callback(progress, IterationInterval(10))

## Spin up simulation then reset.
#run!(simulation)
#reset!(simulation)
#simulation.stop_time = 30days
#model.closure = idealized_one_degree_closure(κ_skew=1e3)

u, v, w = model.velocities
KE = @at (Center, Center, Center) u^2 + v^2
ζ = KernelFunctionOperation{Face, Face, Center}(ζ₃ᶠᶠᶜ, grid, computed_dependencies=(u, v))

simulation.output_writers[:surface] = JLD2OutputWriter(model, merge(fields(model), (; KE, ζ)),
                                                       schedule = TimeInterval(1day),
                                                       filename = filename * "_surface.jld2",
                                                       indices = (:, :, grid.Nz),
                                                       overwrite_existing = true)

run!(simulation)

ut = FieldTimeSeries(filename * "_surface.jld2", "u")
vt = FieldTimeSeries(filename * "_surface.jld2", "v")
Tt = FieldTimeSeries(filename * "_surface.jld2", "T")
Kt = FieldTimeSeries(filename * "_surface.jld2", "KE")
Zt = FieldTimeSeries(filename * "_surface.jld2", "ζ")
t = ut.times
Nt = length(t)

fig = Figure(resolution=(1800, 900))
axk = Axis(fig[2, 2], xlabel="Longitude", ylabel="Latitude")
axz = Axis(fig[2, 3], xlabel="Longitude", ylabel="Latitude")
slider = Slider(fig[3, 2:3], range=1:Nt, startvalue=1)
n = slider.value

title = @lift string("Surface kinetic energy at ", prettytime(t[$n]))
Label(fig[1, 1], title, tellwidth=false)

KEⁿ = @lift interior(Kt[$n], :, :, 1) ./ 2
ζⁿ = @lift interior(Zt[$n], :, :, 1)

hmk = heatmap!(axk, KEⁿ, colorrange=(0, 0.1))
hmz = heatmap!(axz, ζⁿ, colorrange=(-1e-4, 1e-4))

Colorbar(fig[2, 1], hmk, label="Surface kinetic energy (m² s⁻²)")
Colorbar(fig[2, 4], hmz, label="Surface vertical vorticity (s⁻¹)")

display(fig)

record(fig, filename * ".mp4", 1:Nt, framerate=12) do nn
    n[] = nn
end


text:
using Monadic
using ExprParsers
import Base: foreach, map, eltype

"""
  foreach(f, functor)

Like map, but destroying the context. I.e. makes only sense for functions with side effects

map, flatmap and foreach should have same semantics.

Note that this does not work for all Functors (e.g. not for Callables), however is handy in many cases.
This is also the reason why we don't use the default fallback to map, as this may make no sense for your custom Functor.
"""

"""
    @syntax_foreach begin
      # Vectors behaves like nested for loops within @syntax_foreach
      a = [1, 2, 3]
      b = [10, 20]
      @pure a + b
    end
    # [[11, 21], [12, 22], [13, 23]]

This is a variant of the monadic syntax which uses `foreach` for both map_like and flatmap_like.
See `Monadic.@monadic` for more details.
"""
macro syntax_foreach(block::Expr)
  block = macroexpand(__module__, block)
  esc(monadic(:(TypeClasses.foreach), :(TypeClasses.foreach), block))
end
macro syntax_foreach(wrapper, block::Expr)
  block = macroexpand(__module__, block)
  esc(monadic(:(TypeClasses.foreach), :(TypeClasses.foreach), wrapper, block))
end

# Functor
# =======

"""
  map(f, functor)

The core functionality of a functor, applying a normal function "into" the context defined by the functor.
Think of map and vector as best examples.
"""

"""
    @syntax_map begin
      # Vectors behave similar to nested for loops within @syntax_map
      a = [1, 2, 3]
      b = [10, 20]
      @pure a + b
    end
    # [[11, 21], [12, 22], [13, 23]]

This is a variant of the monadic syntax which uses `map` for both map_like and flatmap_like.
See `Monadic.@monadic` for more details.
"""
macro syntax_map(block::Expr)
  block = macroexpand(__module__, block)
  esc(monadic(:(TypeClasses.map), :(TypeClasses.map), block))
end
macro syntax_map(wrapper, block::Expr)
  esc(monadic(:(TypeClasses.map), :(TypeClasses.map), wrapper, block))
end

# Applicative
# ===========

"""
  pure(T::Type, a)

wraps value a into container T
"""
function pure end

"""
  ap(f::F1, a::F2) -> F3

Apply function in container `F1` to element in container `F2`, returning results in the same container `F3`.
The default implementation uses `flatmap` and `map`, so that in general only those two need to be defined.
"""
function ap(f, a)
  default_ap_having_map_flatmap(f, a)
end

# basic functionality
# -------------------

"""
    @mapn f(a, b, c, d)

translates to

    mapn(f, a, b, c, d)
"""
macro mapn(call_expr)
  parsed = parse_expr(EP.Call(), call_expr)
  @assert isempty(parsed.kwargs) "mapn does not support keyword arguments, but found $(parsed.kwargs)."
  callee = if isempty(parsed.curlies)
    parsed.name
  else
    :($(parsed.name){($(parsed.curlies...))})
  end
  esc(:(TypeClasses.mapn($callee, $(parsed.args...))))
end

# generic implementation, looping over the arguments
# - - - - - - - - - - - - - - - - - - - - - - - - - -

# because of world problem we cannot call eval in curry (see e.g. https://discourse.julialang.org/t/running-in-world-age-x-while-current-world-is-y-errors/5871/3)
# hence we define it logically by collecting all args
function curry(func, n::Int)
  function h(x, prev_xs)
    new_xs = tuple(prev_xs..., x)
    if length(new_xs) == n
      func(new_xs...)
    else
      x -> h(x, new_xs)
    end
  end
  x -> h(x, ())
end

# Note that the try to make this a generative function failed because of erros like
# """ERROR: The function body AST defined by this @generated function is not pure. This likely means it contains a closure or comprehension."""
"""
    mapn(f, a1::F{T1}, a2::F{T2}, a3::F{T3}, ...) -> F{T}

Apply a function over applicative contexts instead of plain values.
Similar to `Base.map`, however sometimes the semantic differs slightly. `TypeClasses.mapn` always follows the semantics of `TypeClasses.flatmap` if defined.

E.g. for `Base.Vector` the `Base.map` function zips the inputs and checks for same length. On the other hand `TypeClasses.mapn` combines all combinations of inputs
instead of the zip (which conforms with the semantics of flattening nested Vectors).
"""
function mapn(func, args...)
  n = length(args)
  if n == 0
    error("mapn is not defined for 0 args as the type T needed for pure(T, x) is not known")
  else
    cfunc = curry(func, n)
    Base.foldl((f, a) -> ap(f, a), args[2:end]; init = map(cfunc, args[1]))
  end
end

# specifically compiled versions
# - - - - - - - - - - - - - - -

# We compile specific versions for specific numbers of arguments via code generation.
# Tests showed that this leads way more optimal code.
# Note, that it is not possible to do this with generated functions, as the resulting expression needs to curry
# the function, which requires closures, however they are not supported within generated function bodies.

function _create_curried_expr(func, N::Int)
  args_symbols = Symbol.(:a, 1:N)
  final_call = :($func($(args_symbols...)))
  foldr(args_symbols, init=final_call) do a, subfunc
    :($a -> $subfunc)
  end
end
function _create_ap_expr(curried_func, args_symbols)
  @assert length(args_symbols) >= 1
  first_applied = :(map($curried_func, $(args_symbols[1])))
  foldl(args_symbols[2:end], init=first_applied) do partially_applied, a
    :(ap($partially_applied, $a))
  end
end

for N in 1:128  # 128 is an arbitrary number
  func_name = :func
  args_symbols = Symbol.(:a, 1:N)
  curried_func_name = :curried_func
  curried_func_expr = _create_curried_expr(func_name, N)
  ap_expr = _create_ap_expr(curried_func_name, args_symbols)
  @eval function mapn($func_name, $(args_symbols...))
    $curried_func_name = $curried_func_expr
    $ap_expr
  end
end

# common applications of mapn
# - - - - - - - - - - - - - -

"""
    tupled(Option(1), Option(2), Option(3)) == Option((1,2,3))
    tupled(Option(1), None, Option(3)) == None

Combine several Applicative contexts by building up a Tuple
"""
function tupled(applicatives...)
  mapn(tuple, applicatives...)
end

# default implementations
# -----------------------

# default implementations collide to often with other traits definitions and rather do harm than good
# hence we just give a default definition which can be used for custom definitions
default_map_having_ap_pure(f, a::A) where {A} = ap(pure(A, f), a)

# generic Monoid implementations for Applicative
# ----------------------------------------------

# it is tempting to overload `neutral`, `combine` and so forth directly,
# however dispatching on eltype is not allowed if used for different semantics, as the element-type is an unstable property of a return value, so it should not lead to different results the one or the other way.

function neutral_applicative(T::Type)
  pure(T, neutral(eltype(T)))
end

function combine_applicative(a, b)
  mapn(⊕, a, b)
end

function orelse_applicative(a, b)
  mapn(orelse, a, b)
end

# Monad
# =====

# we use flatmap instead of flatten as the default function, because of the following reasons:
# - flatmap has a similar signature as map and ap, and hence should fit more intuitively into the overal picture for beginners
# - flatmap is what is most often used in e.g. `@syntax_flatmap`, and hence this should optimal code
# - flatten seams to have more information about the nested types, however as julia's typeinference is incomplete
#     and may not correctly infer subtypes, it is actually quite tricky to dispatch correctly on nested typevariables.
#     Concretely you want the functionality to not differ whether typeinference worked or not, and hence you have
#     To deal with Vector{Any} and similar missing typeinformation on the typeparameters. You can fix the types,
#     but you always have to be careful that you do so.
#     And of course this approach only works for containers where you actually have a proper `eltype`.
#     A last point to raise is that not dispatching on nested Functors prevents the possibility of flattening
#     multiple different Functors. This is partly, however the same result can be achieved much better using
#     ExtensibleEffects.
#     Hence we don't see any real use of complex nested dispatch any longer and recommend rather not to use it
#     because of the unintuitive complexity.

"""
    flatmap(function_returning_A, a::A)::A

`flatmap` applies a function to a container and immediately flattens it out.
While map would give you `A{A{...}}`, flatmap gives you a plain `A{...}`, without any nesting.

If you define your own versions of flatmap, the recommendation is to apply a `Base.convert` after applying `f`.
This makes sure your flatmap is typesafe, and actually enables sound interactions with other types which may be
convertable to your A.

E.g. for Vector the implementation looks as follows:
```
TypeClasses.flatmap(f, v::Vector) = vcat((convert(Vector, f(x)) for x in v)...)
```
"""
function flatmap end

# basic functionality
# -------------------

"""
    flatten(::A{A})::A = flatmap(identity, a)

`flatten` gets rid of one level of nesting. Has a default fallback to use `flatmap`.
"""
flatten(a) = flatmap(identity, a)

"""
    a ↠ b = flatmap(_ -> b, a) # \\twoheadrightarrow

A convenience operator for monads which just applies the second monad within the first one.

The operator ↠ (\\twoheadrightarrow) is choosen because the other favourite ≫ (\\gg) which would have been in accordance
with [haskell unicode syntax for monads](https://hackage.haskell.org/package/base-unicode-symbols-0.2.4.2/docs/Control-Monad-Unicode.html)
is unfortunately parsed as boolean comparison with extra semantics which leads to errors with non-boolean applications.

↠ is just the most similar looking other operator which does not have this restriction and is right-associative.

"""
a ↠ b = flatmap(_ -> b, a)

# right associative
↠(a, b, more...) = a ↠ b ↠ foldr(↠, more)  # more is non-empty, as the case `a ↠ b` is always defined

"""
    @syntax_flatmap begin
      # Vector behave similar to nested for loops within @syntax_flatmap
      a = [1, 2, 3]
      b = [10, 20]
      @pure a + b
    end
    # [11, 21, 12, 22, 13, 23]

This is the standard monadic syntax which uses `map` for map_like and `flatmap` for flatmap_like.
See `Monadic.@monadic` for more details.
"""
macro syntax_flatmap(block::Expr)
  block = macroexpand(__module__, block)
  esc(monadic(:(TypeClasses.map), :(TypeClasses.flatmap), block))
end
macro syntax_flatmap(wrapper, block::Expr)
  block = macroexpand(__module__, block)
  esc(monadic(:(TypeClasses.map), :(TypeClasses.flatmap), wrapper, block))
end

# Now can define the fallback definition for ap with monad style
# ==============================================================

# we don't overwrite `ap` directly because it can easily conflict with custom definitions
# instead it is offered as a default definition which you can use for your custom type

# for monadic code we actually don't need Pure
default_ap_having_map_flatmap(f, a) = @syntax_flatmap begin
  f = f
  a = a
  @pure f(a)
end


text:
# This file was generated by the Julia Swagger Code Generator
# Do not modify this file directly. Modify the swagger specification instead.

mutable struct UpdateKeyVaultMetaInfo <: SwaggerModel
    encryptionKeyVersion::Any # spec type: Union{ Nothing, String } # spec name: encryptionKeyVersion

    function UpdateKeyVaultMetaInfo(;encryptionKeyVersion=nothing)
        o = new()
        validate_property(UpdateKeyVaultMetaInfo, Symbol("encryptionKeyVersion"), encryptionKeyVersion)
        setfield!(o, Symbol("encryptionKeyVersion"), encryptionKeyVersion)
        o
    end
end # type UpdateKeyVaultMetaInfo

const _property_map_UpdateKeyVaultMetaInfo = Dict{Symbol,Symbol}(Symbol("encryptionKeyVersion")=>Symbol("encryptionKeyVersion"))
const _property_types_UpdateKeyVaultMetaInfo = Dict{Symbol,String}(Symbol("encryptionKeyVersion")=>"String")
Base.propertynames(::Type{ UpdateKeyVaultMetaInfo }) = collect(keys(_property_map_UpdateKeyVaultMetaInfo))
Swagger.property_type(::Type{ UpdateKeyVaultMetaInfo }, name::Symbol) = Union{Nothing,eval(Meta.parse(_property_types_UpdateKeyVaultMetaInfo[name]))}
Swagger.field_name(::Type{ UpdateKeyVaultMetaInfo }, property_name::Symbol) =  _property_map_UpdateKeyVaultMetaInfo[property_name]

function check_required(o::UpdateKeyVaultMetaInfo)
    true
end

function validate_property(::Type{ UpdateKeyVaultMetaInfo }, name::Symbol, val)
end


text:
# functions for doing parallel computations that are precisely the reverse mode
# of the regular parallel communication
# This can be suboptimal because it requires doing 2 parallel communications
# one after another (one for q and another one to return q_bar), rather than
# doing one big communication.
# The idea behind the functions in this file is:
#
# Primal: q_local + q_remote -> res_local
# Reverse: q_local + q_remote + res_bar_local -> q_bar_local + q_bar_remote
#
# The fact that q_bar_remote is computed necessitates the second communication

"""
  This function combines the functionality of finishExchangeData for the
  solution and exchangeData for `q_bar`.  This function does 4 things:

    1. waits to receive the solution from the remote process (typically started
       with [`startSolutionExchange`](@ref)
    2. calls `calc_func` to populate `send_data_bar.q_recv` with `q_bar`
    3. sends `q_bar` to remote process
    4. posts receives for `q_bar` from remote processes

  Users should call [`finishSolutionExchange_rev`](@ref) to complete
  the receives for `q_bar`.

  It is safe to call this function even if `shared_data` has already been
  waited on.

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * shared_data: the vector of `SharedFaceData` objects for the solution
   * shared_data_bar: the vector of `SharedFaceData` objects for `q_bar`
   * calc_func: function that takes an element of `shared_data` and populates
                the *receive* buffer of the corresponding element of
                `shared_data_bar`.  See below

  **Keyword Arguments**

   * wait: if true, waits for the `shared_data_bar` sends and receives to finish
           before exiting.  Default false

  `calc_func` should have the signature

  ```
    calc_func(mesh, sbp, eqn, opts, data::SharedFaceData,
              data_bar::SharedFaceData)
  ```

  where `data` is one element of `shared_data` and `data_bar` is the
  corresponding element of `shared_data_bar`.  The receive buffer of
  `data` will have been populated before `calc_func` is called.  `calc_func`
  must populate the *receive* buffer of `data_bar` with `q_bar` values.
  Note that this is opposite the usual usage of `SharedFaceData` objects.
  After `calc_func` returns, the data in `data_bar.q_recv` will be sent
  to the remote MPI process.  The receive buffer will be zeroed out
  before being passed to `calc_func`.

  Currently `shared_data_bar` must have a `parallel_data` setting of PARALLEL_DATA_ELEMENT
"""
function exchangeData_rev(mesh::AbstractMesh, sbp::AbstractOperator,
                      eqn::AbstractSolutionData, opts,
                      shared_data::Vector{SharedFaceData{T}},
                      shared_data_bar::Vector{SharedFaceData{T}},
                      calc_func::Function;
                      wait=false) where T

  npeers = length(shared_data)

  # bail out early if there is no communication to do
  # not sure if the rest of this function runs correctly if npeers == 0
  if npeers == 0
    return nothing
  end

  # this should already have happened.  If it hasn't something else has
  # gone wrong in the solver.  Throw an exception
  assertSendsWaited(shared_data_bar)

  # post the receives first
  for i=1:npeers
    Irecv_rev!(shared_data_bar[i])
  end

  # verify the sends are consistent
  val = assertReceivesConsistent(shared_data)
  assertReceivesConsistent(shared_data_bar)

  for i=1:npeers
    # wait on a primal communication to finish
    if val == 0  # requests not previously waited on
      eqn.params.time.t_wait += @elapsed idx = waitAnyReceive(shared_data)
    else
      idx = i
    end

    data_i = shared_data[idx]
    data_bar_i = shared_data_bar[idx]
    tag = data_bar_i.tag

    if data_i.pdata == PARALLEL_DATA_FACE && val == 0
      # permute the received nodes to be in the elementR orientation
      permuteinterface!(mesh.sbpface, data_i.interfaces, data_i.q_recv)
    end

    # wait on the previous bar send if it hasn't been waited on yet
    # this should have completed long ago
    if !data_bar_i.recv_waited
      waitReceive(data_bar_i)
    end
    fill!(data_bar_i.q_recv, 0)

    # call calc_func to populate bar q_recv
    calc_func(mesh, sbp, eqn, opts, data_i, data_bar_i)

    # post the bar send
    Isend_rev(data_bar_i)
  end

  if wait
    waitAllSends(shared_data)
    waitAllSends(shared_data_bar)
    waitAllReceives(shared_data)
    waitAllReceives(shared_data_bar)
  end

  return nothing
end

"""
  This function completes the receive of the [`shared_data_bar`](@ref)
  started by [`exchangeData_rev`](@ref).

  It is safe to call this function even if the `shared_data_bar` objects
  have been waited on already, however calling this function more than once
  will result in `populate_buffer_rev` being called multipled times, which is
  likely not what the user wants.

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * shared_data_bar: vector of [`SharedFaceData`](@ref) objects.  Must be
                      the same vector passed to the `shared_data_bar` argument
                      of `exchangeData_rev`.
   * populate_buffer_rev: function that will be called after every receive.
                          This function is often the reverse mode of
                          the `populate_buffer` argument to `exchangeData`

  `populate_buffer_rev` must have the signature

  ```
    populate_buffer_rev(mesh, sbp, eqn, opts, data_bar::SharedFaceData)
  ```

  And should take the data in the *send* buffer of the `SharedFaceData` object
  and do some computation with it, typically updating `eqn.q_bar`.  The reason
  the send buffer should be used is that this is a reverse mode function, which
  reverses the typical rules of send and receive.

  The parallel_data setting of the `shared_data_bar` objects must be
  `PARALLEL_DATA_ELEMENT`.
"""
function finishExchangeData_rev(mesh, sbp, eqn, opts,
                            shared_data_bar::Vector{SharedFaceData{T}},
                            populate_buffer_rev::Function) where T

  npeers = length(shared_data_bar)
  val = assertSendsConsistent(shared_data_bar)

  for i=1:npeers
    if val == 0  # request have not been waited on previously
      eqn.params.time.t_wait += @elapsed idx = waitAnySend(shared_data_bar)
    else
      idx = i
    end

    data_idx = shared_data_bar[idx]

    # to support PARALLEL_DATA_FACE: add permutation function here

    populate_buffer_rev(mesh, sbp, eqn, opts, data_idx)
  end

  # wait on all receives to complete.  This is not strictly necessary, but
  # avoids a problem with exchangeData_rev2 where it expects the receives
  # to already have been waited on
  waitAllReceives(shared_data_bar)

  return nothing
end

"""
  This is a thin wrapper around [`finishExchangeData_rev`](@ref) for the
  common case of receiving `q_bar` from other processes and adding it to
  `eqn.q_bar.

  **Inputs**

   * mesh
   * sbp
   * eqn: communication for `eqn.shared_data_bar` completed
   * opts
"""
function finishSolutionBarExchange(mesh::AbstractMesh, sbp::AbstractOperator,
                                   eqn::AbstractSolutionData, opts)

  if mesh.npeers == 0
    return nothing
  end

  pdata = eqn.shared_data_bar[1].pdata  # assume all are same
  if pdata == PARALLEL_DATA_FACE
    error("PARALLEL_DATA_FACE not supported")
  elseif pdata == PARALLEL_DATA_ELEMENT
    populate_buffer_rev = getSendDataElement_rev
  else
    throw(ErrorException("unsupported parallel_type = $(getParallelDataString(pdata))"))
  end

  finishExchangeData_rev(mesh, sbp, eqn, opts, eqn.shared_data_bar, populate_buffer_rev)

  return nothing
end

"""
  Reverse mode of [`getSendDataElement`](@ref).  Takes data from the send
  buffer of `data` and adds it to `eqn.q_bar`
"""
function getSendDataElement_rev(mesh::AbstractMesh, sbp::AbstractOperator,
                         eqn::AbstractSolutionData, opts, data::SharedFaceData)

  # copy data into send buffer
  idx = data.peeridx
  local_els = mesh.local_element_lists[idx]
  send_buff = data.q_send
  for j=1:length(local_els)
    el_j = local_els[j]
    for k=1:size(eqn.q, 2)
      for p=1:size(eqn.q, 1)
        eqn.q_bar[p, k, el_j] += send_buff[p, k, j]
      end
    end
  end

  return nothing
end