text:
searchdir(path,key) = filter(x->occursin(key,x), readdir(path))
filename(store::FileDataStore, entry::E) where E <: AbstractBigDataEntry = "$(store.folder)/$(entry.storeKey).dat"
function readentry(store::FileDataStore, entry::E) where E <: AbstractBigDataEntry
    open(filename(store, entry)) do f
        return read(f)
    end
end
function writeentry(store::FileDataStore, entry::E, data::Vector{UInt8}) where E <: AbstractBigDataEntry
    open(filename(store, entry), "w") do f
        write(f, data)
    end
end

function getBigData(store::FileDataStore, entry::E)::Union{Vector{UInt8}, Nothing} where {E <: AbstractBigDataEntry}
    length(searchdir(store.folder, String(entry.storeKey)*".dat")) !=1 && (@warn "Could not find unique file for key '$(entry.storeKey)'."; return nothing)
    return readentry(store, entry)
end

function addBigData!(store::FileDataStore, entry::E, data::Vector{UInt8})::Vector{UInt8} where {E <: AbstractBigDataEntry}
    length(searchdir(store.folder, String(entry.storeKey)*".dat")) !=0 && @warn "Key '$(entry.storeKey)' already exists, overwriting!."
    writeentry(store, entry, data)
    # Update timestamp
    entry.lastUpdatedTimestamp = now()
    return getBigData(store, entry)
end

function updateBigData!(store::FileDataStore, entry::E, data::Vector{UInt8})::Union{Vector{UInt8}, Nothing} where {E <: AbstractBigDataEntry}
    length(searchdir(store.folder, String(entry.storeKey)*".dat")) !=1 && (@warn "Could not find unique file for key '$(entry.storeKey)'."; return nothing)
    writeentry(store, entry, data)
    # Update timestamp
    entry.lastUpdatedTimestamp = now()
    return getBigData(store, entry)
end

function deleteBigData!(store::FileDataStore, entry::E)::Vector{UInt8} where {E <: AbstractBigDataEntry}
    data = getBigData(store, entry)
    data == nothing && return nothing
    rm(filename(store, entry))
    return data
end

# TODO: Manifest file
#
# function listStoreEntries(store::FileDataStore)::Vector{E} where {E <: AbstractBigDataEntry}
#     return collect(values(store.entries))
# end


text:
module PackageExample

export foo, bar

foo(x::T, y::T) where T <: Real = x + y - 5
bar(z::Float64) = foo(sqrt(z), z)

end # module

ext: jl


text:
# http://web.mit.edu/puzzle/www/2013/coinheist.com/rubik/clockwork_orange/index.html
@testset "rubik clockwork orange meta" begin
    f = best_feature(["armoredrecon", "hypapante", "commemorativebats", "derricktruck", "brownrot", "attorneysgeneral", "sacrosanct", "impromptu"], 1)
    @test f.description == "has at least 2 letters repeated exactly 2 times each"
end


text:
using Base.Test

using TALib
using TALib: basepath

using DataFrames

@testset "level 2 dataframes" begin

    epsilon = 1e-10
    eps_price = 1e-6

    filename = joinpath(basepath(), "ford_2012.csv")
    df = readtable(filename)
    df[:Date] = Date(df[:Date])

    indic = MA(df)
    @test indic[end, :Real] ≈ 11.546 atol=eps_price

    indic = MA(df, price=:Open)
    @test indic[end, :Real] ≈ 11.4613333 atol=eps_price

    indic = BBANDS(df)
    @test indic[end, :UpperBand] ≈ 13.13191 atol=eps_price  # UpperBand
    @test indic[end, :MiddleBand] ≈ 12.75400 atol=eps_price  # MiddleBand
    @test indic[end, :LowerBand] ≈ 12.37609 atol=eps_price  # LowerBand

end


text:
@everywhere function encode_viterbi!{T <: AbstractFloat}(
  CODES::SharedMatrix{Int16},  # out. Where to save the result
  X::Matrix{T},                # in. d-by-n matrix to encode
  C::Vector{Matrix{T}},        # in. m-long vector with d-by-h codebooks
  binaries::Vector{Matrix{T}}, # in. Binary terms
  IDX::UnitRange{Int64})       # in. Index to save the result

  # Get unaries
  unaries = get_unaries( X, C );

  h, n = size( unaries[1] );
  m    = length( binaries ) + 1;

  # We need a matrix to keep track of the min and argmin
  mincost = zeros(T, h, m );
  minidx  = zeros(Integer, h, m );

  # Allocate memory for brute-forcing each pair
  cost = zeros( T, h );

  U = zeros(T, h, m)

  minv = typemax(T);
  mini = 1;

  uidx = 1;
  @inbounds for idx = IDX # Loop over the datapoints

    # Put all the unaries of this item together
    for i = 1:m
      ui = unaries[i]
      @simd for j = 1:h
        U[j,i] = ui[j,uidx];
      end
    end

    # Forward pass
    for i = 1:(m-1) # Loop over states

      # If this is not the first iteration, add the precomputed costs
      if i > 1
        @simd for j = 1:h
          U[j,i] += mincost[j,i-1];
        end
      end

      bb = binaries[i];
      for j = 1:h # Loop over the cost of going to j
        @simd for k = 1:h # Loop over the cost of coming from k
          ucost   =  U[k, i]; # Pay the unary of coming from k
          bcost   = bb[k, j]; # Pay the binary of going from j-k
          cost[k] = ucost + bcost;
        end

        # Writing my own findmin
        minv = cost[1]
        mini = 1
        for k = 2:h
          costi = cost[k]
          if costi < minv #|| minv!=minv
            minv = costi
            mini = k
          end
        end

        mincost[j, i] = minv;
         minidx[j, i] = mini;
      end
    end

    @simd for j = 1:h
      U[j,m] += mincost[j,m-1];
    end

    _, mini = findmin( U[:,m] );

    # Backward trace
    backpath = [ mini ];
    for i = (m-1):-1:1
      push!( backpath, minidx[ backpath[end], i ] );
    end

    # Save the inferred code
    CODES[:, idx] = reverse( backpath );

    uidx = uidx + 1;
  end # for idx = IDX
end

"Function to call that encodes a dataset using dynamic programming"
function encoding_viterbi(
  X::Matrix{Float32},         # d-by-n matrix. Data to encode
  C::Vector{Matrix{Float32}}, # m-long vector with d-by-h codebooks
  V::Bool=false)              # whether to print progress

  d, n = size( X );
  m    = length( C );

  # Compute binary tables
  binaries = Vector{Matrix{Float32}}(m-1);
  for i = 1:(m-1)
    binaries[i] = 2 * C[i]' * C[i+1];
  end
  CODES = SharedArray{Int16,2}((m, n));

  if nworkers() == 1
    encode_viterbi!( CODES, X, C, binaries, 1:n );
  else
    paridx = splitarray( 1:n, nworkers() );
    @sync begin
      for (i,wpid) in enumerate(workers())
        @async begin
          Xw = X[:,paridx[i]]
          remotecall_wait(encode_viterbi!, wpid, CODES, Xw, C, binaries, paridx[i]);
        end
      end
    end
  end

  return sdata(CODES);
end

max_issues_repo_name: xiaosu-zhu/local-search-quantization

max_forks_repo_name: xiaosu-zhu/local-search-quantization


text:
#---------------------------------------------------------
# # Optimization of a simple phosphorus cycling model
#---------------------------------------------------------

#md # !!! tip
#md #     This example is also available as a Jupyter notebook:
#md #     [`P_model_2_tracers.ipynb`](@__NBVIEWER_ROOT_URL__/examples/generated/P_model_2_tracers.ipynb)

#md # !!! warn
#md #     This example is not finished yet.
#md #     It lacks a number of things, including explanatory text, references, and more.
#md #     However, it can already serve as an example of model optimization.

#nb # > **Warning!**
#nb # > This example is not finished yet.
#nb # > It lacks a number of things, including explanatory text, references, and more.
#nb # > However, it can already serve as an example of model optimization.

#---------------------------------------------------------
# ## Tracer equations
#---------------------------------------------------------

# We consider a simple model for the cycling of phosphorus with 2 state variables consisting of phosphate (PO₄) AKA dissolved inorganic phosphorus (DIP) and particulate organic phosphorus (POP).
# The dissolved phases are transported by advection and diffusion whereas the particulate phase sinks rapidly down the water column without any appreciable transport by the circulation.
#
# The governing equations that couple the 3D concentration fields of DIP and POP, denoted $x_\mathsf{DIP}$ and $x_\mathsf{POP}$, respectively, are:
#
# $$\left[\frac{\partial}{\partial t} + \nabla \cdot (\boldsymbol{u} + \mathbf{K}\cdot\nabla )\right] x_\mathsf{DIP} = -U(x_\mathsf{DIP}) + R(x_\mathsf{POP}),$$
#
# and
#
# $$\left[\frac{\partial}{\partial t} + \nabla \cdot \boldsymbol{w}\right] x_\mathsf{POP} = U(x_\mathsf{DIP}) - R(x_\mathsf{POP}).$$

# The 3D vector $\boldsymbol{u}$ is the fluid velocity and $\mathbf{K}$ is the eddy-diffusion tensor.
# Thus, $\nabla \cdot \left[ \boldsymbol{u} - \mathbf{K} \cdot \nabla \right]$ is a differential operator that represents the transport by the ocean circulation.
# The 3D vector $\boldsymbol{w}$ is the particulate sinking velocity, so that the operator $\nabla \cdot \boldsymbol{w}$ acting on $x_\mathsf{POP}$ is the divergence of the flux of POP.
# We will assume that it increases linearly with depth, i.e., that
#
# $$\boldsymbol{w}(\boldsymbol{r}) = \boldsymbol{w}_0 + \boldsymbol{w}' z,$$
#
# where $\boldsymbol{w}_0$ and $\boldsymbol{w}'$ are simple scalars, independent of location $\boldsymbol{r}$.
# ($z$ represents the depth of location $\boldsymbol{r}=(x,y,z)$.)

# The function $U$ represents the biological uptake of DIP by phytoplankton, which we model as
#
# $$U(x_\mathsf{DIP}) = \frac{x_\mathsf{DIP}}{\tau} \, \frac{x_\mathsf{DIP}}{x_\mathsf{DIP} + k} \, (z < z_0),$$
#
# where the timescale, $\tau$, sets the maximum specific DIP uptake rate.
# The first term in the product scales linearly with DIP, which is a nice property to have in general.
# It also makes it easier to reason about the extreme cases.
# However, here the uptake is not linear:
# the second term in the product is a hyperbolic term (after Monod or Michaelis-Menten) and reduces the slope of the uptake function when DIP is low.
# When DIP is very low, the slope is close to 0.
# But when DIP is very high, the slope is close to $1/τ$.
# The last term is to be understood as a condition, i.e., that the uptake happens only in the euphotic layer, where there is sufficient light for photosynthesis to occur.
# The depth $z_0$ is the depth of the euphotic layer.

# The function $R$ defines the remineralization rate of POP, which converts POP back into DIP.
# For the remineralization, we simply use a linear rate constant, i.e.,
#
# $$R(x_\mathsf{POP}) = \kappa \, x_\mathsf{POP}.$$

#---------------------------------------------------------
# ## Translation into AIBECS code
#---------------------------------------------------------

# ### Transport matrices

# We start by telling Julia we want to use the AIBECS

using AIBECS

# We load the circulation and grid of the OCIM0.1 via

wet3D, grd, T_OCIM = OCIM0.load() ;

# where the operator $\nabla \cdot \left[ \boldsymbol{u} - \mathbf{K} \cdot \nabla \right]$ is represented by the sparse matrix `T_OCIM`.
# In AIBECS, all the transport matrices must be declared as functions of the parameters `p`, so we simply define

T_DIP(p) = T_OCIM

# For the sinking of particles, we are going to create a sparse matrix that depends on the parameters, too, which define how fast particles sink.

T_POP(p) = buildPFD(grd, wet3D, sinking_speed = w(p))

# for which we need to define the sinking speed `w(p)` as a function of the parameters `p`.
# Following the assumption that it increases linearly with depth, we write it as

w(p) = p.w₀ .+ p.w′ * z

# For this to work, we must create a vector of depths, `z`, which is simply done via

iwet = findall(wet3D)
z = ustrip.(grd.depth_3D[iwet])

# (We must strip the units via `ustrip` because units do not percolate through the functionality of AIBECS flawlessly at this early stage of development.)
# We combine the transport operators for DIP and POP into a single tuple via

T_all = (T_DIP, T_POP) ;

# ### Local sources and sinks

# ##### Geological Restoring

# Because AIBECS will solve for the steady state solution directly without time-stepping the goverining equations to equilibrium, we don't have any opportunity to specify any intial conditions.
# Initial conditions are how the total amount of conserved elements get specified in most global biogeochemical modelels.
# Thus to specify the total inventory of P in AIBECS we add a very weak resporing term to, e.g., the DIP equation.
# The time-scale for this restoring term is chosen to be very long compared to the timescale with which the ocean circulation homogenizes a tracer.
# Because of this long timescale we call it the geological restoring term, but geochemists who work on geological processes don't like that name!
# In any event the long timescale allows us to prescribe the total inventory of P in a way that yields the same solution we would have gotten had we time-stepped the model to steady-state with the total inventory prescribed by the initial condition.
# We define the geological restoring simply as

geores(x, p) = (p.DIPgeo .- x) / p.τgeo

# ##### Uptake (DIP → POP)

# For the phosphate uptake, we must take some precautions to avoid the pole of the hyperbolic term.
# That is, we ensure that uptake only occurs when DIP is positive via the `relu` function

relu(x) = (x .≥ 0) .* x

# Then we simply write the uptake as a Julia function, following the definition of $U$:

function uptake(DIP, p)
    τ, k, z₀ = p.τ, p.k, p.z₀
    DIP⁺ = relu(DIP)
    return 1/τ * DIP⁺.^2 ./ (DIP⁺ .+ k) .* (z .≤ z₀)
end

# where we have "unpacked" the parameters to make the code clearer and as close to the mathematical equation as possible.

# ##### Remineralization (POP → DIP)

remineralization(POP, p) = p.κ * POP

# ##### Net sources and sinks

# We lump the sources and sinks into `sms` functions (for Sources Minus Sinks):

function sms_DIP(DIP, POP, p)
    return -uptake(DIP, p) + remineralization(POP, p) + geores(DIP, p)
end
function sms_POP(DIP, POP, p)
    return  uptake(DIP, p) - remineralization(POP, p)
end

# and create a tuple of those just like for the transport matrices via

sms_all = (sms_DIP, sms_POP) # bundles all the source-sink functions in a tuple

# For the sake of this notebook, however, we write more efficient BGC functions for the sources and sinks by using "in place" functions that mutate a preallocated vector to avoid allocations.
# `G_DIP!` replaces `sms_DIP`:

function G_DIP!(dDIP, DIP, POP, p)
    τ, k, z₀, κ, DIPgeo, τgeo = p.τ, p.k, p.z₀, p.κ, p.DIPgeo, p.τgeo
    dDIP .= @. -(DIP ≥ 0) / τ * DIP^2 / (DIP + k) * (z ≤ z₀) + κ * POP + (DIPgeo - DIP) / τgeo
    return dDIP
end

# Note the use of the `@.` macro to "fuse" the loops together, which is faster than looping for each operation.
# And `G_POP!` replaces `sms_POP`:

function G_POP!(dPOP, DIP, POP, p)
    τ, k, z₀, κ = p.τ, p.k, p.z₀, p.κ
    dPOP .= @. (DIP ≥ 0) / τ * DIP^2 / (DIP + k) * (z ≤ z₀) - κ * POP
    return dPOP
end

# which we bundle into

Gs = (G_DIP!, G_POP!)

# ### Parameters

# We now define and build the parameters.
# In AIBECS, this is done by first creating a table of parameters, using the API designed for it.
# First, we create an empty table:

t = empty_parameter_table()    # empty table of parameters

# Then, we add every parameter that we need

add_parameter!(t, :DIPgeo, 2.12u"mmol/m^3",
               optimizable = true,
               variance_obs = ustrip(upreferred(0.1 * 2.17u"mmol/m^3"))^2,
               description = "Mean PO₄ concentration")
add_parameter!(t, :τgeo, 1.0u"Myr",
               description = "Geological restoring timescale")
add_parameter!(t, :k, 6.62u"μmol/m^3",
               optimizable = true,
               description = "Half-saturation constant (Michaelis-Menten)")
add_parameter!(t, :z₀, 80.0u"m",
               description = "Depth of the euphotic layer base")
add_parameter!(t, :w₀, 0.64u"m/d",
               optimizable = true,
               description = "Sinking velocity at surface")
add_parameter!(t, :w′, 0.13u"1/d",
               optimizable = true,
               description = "Vertical gradient of sinking velocity")
add_parameter!(t, :κ, 0.19u"1/d",
               optimizable = true,
               description = "Remineralization rate constant (POP to DIP)")
add_parameter!(t, :τ, 236.52u"d",
               optimizable = true,
               description = "Maximum uptake rate timescale")

# These values of these parameters have been optimized for the same model but embedded in the OCIM1 circulation and should yield a pretty good fit to observations.
# However, because the circulation has changed from the OCIM1 to the OCIM0.1 used here, the best parameters have probably changed, too.
# Thus, we will re-optimize these parameters below.
# But first, we will just try to find the steady-state solution with these initial parameter values.

# We generate a new type for the parameters, via

initialize_Parameters_type(t, "Pcycle_Parameters")   # Generate the parameter type

# Finally, we can instantiate the parameter vector in a single line:

p = Pcycle_Parameters()

# where we have used the name of the new type that we created just earlier.

#---------------------------------------------------------
# ## Finding the steady-state solution
#---------------------------------------------------------

# ### Generate state function and Jacobian

# We generate the in-place state function `F!` and its Jacobian in a single line
# and the out-of-place `F` for use by the solver.
# (Using `dx = similar(x)` creates a vector of the same type and size as `x` but without spending time copying its values, which is OK because we always fill the entire `dx` in calls to `g_DIP!` and G_POP!`.)

nb = length(iwet)
F!, ∇ₓF = inplace_state_function_and_Jacobian(T_all, Gs, nb)
F(x::Vector{Tx}, p::Pcycle_Parameters{Tp}) where {Tx,Tp} = F!(Vector{promote_type(Tx,Tp)}(undef,length(x)),x,p)

# ### Solve for the steady-state

# We start from an initial guess,

x = p.DIPgeo * ones(2nb) # initial iterate

# Create an instance of the steady-state problem

prob = SteadyStateProblem(F, ∇ₓF, x, p)

# and solve it

s = solve(prob, CTKAlg()).u

# ### Plotting

# For plotting, we first unpack the state

DIP, POP = state_to_tracers(s, nb, 2)

# We will plot the concentration of DIP at a given depth horizon

iz = findfirst(grd.depth .> 2000u"m")
iz, grd.depth[iz]

#-

DIP_3D = rearrange_into_3Darray(DIP, wet3D)
DIP_2D = DIP_3D[:,:,iz] * ustrip(1.0u"mol/m^3" |> u"mmol/m^3")
lat, lon = ustrip.(grd.lat), ustrip.(grd.lon)

# Create the Cartopy canvas, make the data cyclic in order for the plot to look good in Cartopy, and plot the filled contour via `contourf`

ENV["MPLBACKEND"]="qt5agg"
using PyPlot, PyCall
clf()
ccrs = pyimport("cartopy.crs")
cfeature = pyimport("cartopy.feature")
ax = subplot(projection = ccrs.EqualEarth(central_longitude=-155.0))
ax.add_feature(cfeature.COASTLINE, edgecolor="#000000") # black coast lines
ax.add_feature(cfeature.LAND, facecolor="#CCCCCC")      # gray land
lon_cyc = [lon; 360+lon[1]]
DIP_2D_cyc = hcat(DIP_2D, DIP_2D[:,1])
plt = contourf(lon_cyc, lat, DIP_2D_cyc, levels=0:0.2:3.6, transform=ccrs.PlateCarree(), zorder=-1)
colorbar(plt, orientation="horizontal");
title("PO₄ at $(string(round(typeof(1u"m"),grd.depth[iz]))) depth using the OCIM0.1 circulation")
gcf()

#---------------------------------------------------------
# ## Optimizing the model parameters
#---------------------------------------------------------

# To optimize the model, we will use data from the World Ocean Atlas (WOA).
# Specifically, we will use the observed PO4 climatology, which the WOA produces every few years.
# Thanks to the [WorldOceanAtlasTools](https://github.com/briochemc/WorldOceanAtlasTools.jl) package combined with some AIBECS functionality, this is done a just a few commands:

using WorldOceanAtlasTools
μDIPobs3D, σ²DIPobs3D = WorldOceanAtlasTools.fit_to_grid(grd, 2018, "phosphate", "annual", "1°", "an")
μDIPobs, σ²DIPobs = μDIPobs3D[iwet], σ²DIPobs3D[iwet]
μx = (μDIPobs, missing)
σ²x = (σ²DIPobs, missing)

# We define some hyper parameters for the weight of the mismatch of DIP, POP, and the parameters with observations / their prior estimates.

ωs = [1.0, 0.0] # the weight for the mismatch (weight of POP = 0)
ωp = 1e-4       # the weight for the parameters prior estimates

# We then generate the objective function and some derivatives using the PO₄ mean and variance

v = ustrip.(vector_of_volumes(wet3D, grd))
f   =   generate_objective(ωs, μx, σ²x, v, ωp, mean_obs(p), variance_obs(p))
∇ₓf = generate_∇ₓobjective(ωs, μx, σ²x, v, ωp, mean_obs(p), variance_obs(p))
∇ₚf = generate_∇ₚobjective(ωs, μx, σ²x, v, ωp, mean_obs(p), variance_obs(p))

# For the optimization, we will use the F-1 algorithm from the [F1Method](https://github.com/briochemc/F1Method.jl) package.

using F1Method

# For the F-1 algorithm, we must initialize the memory buffer

mem = F1Method.initialize_mem(s, p)

# We then define the objective, gradient, and Hessian using the F1Method package API

objective(p) = F1Method.objective(f, F, ∇ₓF, mem, p, CTKAlg())
gradient(p) = F1Method.gradient(f, F, ∇ₓf, ∇ₓF, mem, p, CTKAlg())
hessian(p) = F1Method.hessian(f, F, ∇ₓf, ∇ₓF, mem, p, CTKAlg())

# Because we want the parameters to remain positive, we define the change of variables
# The change variable from vector `λ = log(p)` to Parameters `p`:

λ2p(λ) = AIBECS.opt_para(exp.(λ))

# and the corresponding derivatives:

∇λ2p(λ) = exp.(λ)
∇²λ2p(λ) = exp.(λ)

# Reverse change of variables from Parameters p to vector λ

p2λ(p) = log.(optvec(p))
λ = p2λ(p)

obj(λ) = objective(λ2p(λ))
using LinearAlgebra
grad(λ) = gradient(λ2p(λ)) * Diagonal(∇λ2p(λ))
function hess(λ)
    ∇p = Diagonal(∇λ2p(λ)) # for variable change
    ∇²p = Diagonal(∇²λ2p(λ)) # for variable change
    G = vec(gradient(λ2p(λ)))
    H = hessian(λ2p(λ))
    return ∇p * H * ∇p + Diagonal(G) * ∇²p
end

# Use optim to optimize

using Optim

# Add storage for use with Optim

m = length(p)
grad(s, λ) = s[1:m] .= vec(grad(λ))
hess(s, λ) = s[1:m,1:m] .= hess(λ)

# Set options of Optim

opt = Optim.Options(store_trace = false, show_trace = true, extended_trace = false)

# The starting parameters are

p

# Run optimization

results = optimize(obj, grad, hess, λ, NewtonTrustRegion(), opt)

# Because we started from an other optimized set of parameters, this should only take a few iterations.
# The optimized parameters are

p_optimized = λ2p(results.minimizer)

# After optimizing these parameters, we expect the fit to be better.
# Let's have a look at the optimized steady-state solution at the same depth as earlier, and compare it to the observed DIP field.
# First, let's calculate the steady-state with the optimized parameters.

prob_optimized = SteadyStateProblem(F, ∇ₓF, s, p_optimized)
s_optimized = solve(prob_optimized, CTKAlg()).u

# Like earlier, we fill in a 3D array of `NaN`s with the values we want to plot.
# Here, we fill the arrays with the fractional difference bewteen modeled and observed DIP:

DIPold, _ = state_to_tracers(s, nb, 2)
DIPnew, _ = state_to_tracers(s_optimized, nb, 2)
δDIPold_3D = fill(NaN, size(grd))
δDIPold_3D[iwet] .= 100(DIPold - μDIPobs) ./ μDIPobs
δDIPnew_3D = fill(NaN, size(grd))
δDIPnew_3D[iwet] .= 100(DIPnew - μDIPobs) ./ μDIPobs

# And take a slice at depth index `iz`

δDIPold_2D = δDIPold_3D[:,:,iz]
δDIPnew_2D = δDIPnew_3D[:,:,iz]

# We make those 2D slices cyclic along the longitude for Cartopy

δDIPold_cyc = hcat(δDIPold_2D, δDIPold_2D[:,1])
δDIPnew_cyc = hcat(δDIPnew_2D, δDIPnew_2D[:,1])

# And we plot the mismatch of the old, OCIM1-optimized field

figure()
δDIPlevels = -20:2:20
ax = subplot(projection = ccrs.EqualEarth(central_longitude=-155.0))
ax.add_feature(cfeature.COASTLINE, edgecolor="#000000") # black coast lines
ax.add_feature(cfeature.LAND, facecolor="#CCCCCC")      # gray land
plt2 = contourf(lon_cyc, lat, δDIPold_cyc, cmap="PiYG_r", levels=δDIPlevels, transform=ccrs.PlateCarree(), zorder=-1, extend="both")
cbar2 = colorbar(plt2, orientation="horizontal", extend="both")
cbar2.set_label("δDIP / DIP [%]")
title("old DIP mismatch at $(string(round(typeof(1u"m"),grd.depth[iz]))) depth using the OCIM0.1 circulation")
gcf()

# We can see a significant negative bias in the North Atlantic.
# This is probably due to having swapped the circulation from OCIM1 to OCIM0.1.
# Let's plot the OCIM-0.1-optimized mismatch

figure()
ax = subplot(projection = ccrs.EqualEarth(central_longitude=-155.0))
ax.add_feature(cfeature.COASTLINE, edgecolor="#000000") # black coast lines
ax.add_feature(cfeature.LAND, facecolor="#CCCCCC")      # gray land
plt3 = contourf(lon_cyc, lat, δDIPnew_cyc, cmap="PiYG_r", levels=δDIPlevels, transform=ccrs.PlateCarree(), zorder=-1, extend="both")
cbar3 = colorbar(plt3, orientation="horizontal", extend="both")
cbar3.set_label("δDIP / DIP [%]")
title("new DIP mismatch at $(string(round(typeof(1u"m"),grd.depth[iz]))) depth using the OCIM0.1 circulation")
gcf()

# This shows that most of the negative bias has been corrected.
# (The green patch in the North Atlantic has dissipated.)
# Of course, this optimization was "easy", because there were only 2 tracers, few parameters, and we started from a pretty good solution.
# However, this type of optimization can really improve the skill of models, even current "state-of-the-art" ones.
# Such optimizations turn average models into better tools for cutting-edge oceanography research.

#md # !!! note
#md #     If you find any bugs or ideas you want to test and struggle with the AIBECS, don't hesitate to raise issues on the GitHub [AIBECS](https://github.com/briochemc/AIBECS.jl) repository or [contact me](www.bpasquier.com)!

#nb # > **Note**
#nb # > If you find any bugs or ideas you want to test and struggle with the AIBECS, don't hesitate to raise issues on the GitHub [AIBECS](https://github.com/briochemc/AIBECS.jl) repository or [contact me](www.bpasquier.com)!


text:
using FinancialMonteCarlo, Test;
@show "Test Parameters"
S0 = 100.0;
K = 100.0;
r = 0.02;
Tneg = -1.0;
d = 0.01;
T = -Tneg;
Nsim = 10000;
Nstep = 30;
mc1 = MonteCarloConfiguration(Nsim, Nstep, FinancialMonteCarlo.AntitheticMC());
rfCurve2 = ZeroRate([0.00, 0.02], T);
sigma = 0.2;

McConfig = MonteCarloConfiguration(Nsim, Nstep);
rfCurve = ZeroRate(r);

@show "Test Brownian Motion Parameters"
drift = 0.0
@test_throws(ErrorException, BrownianMotion(-sigma, drift))
@test_throws(ErrorException, BrownianMotion(-sigma, rfCurve2.r))


text:
@testset "777.swap-adjacent-in-lr-string.jl" begin
    @test can_transform("RXXLRXRXL", "XRLXXRRLX") == true
    @test can_transform("XXXL", "LXXL") == false
    @test can_transform("XXXL", "LXX") == false
    @test can_transform("LXXX", "LXXX") == true
end


text:
####
#### Robust online (co-)variance estimators.
####

## Variance estimator

abstract type VarEstimator{T} end

# NOTE: this naive variance estimator is used only in testing
mutable struct NaiveVar{E<:AbstractVecOrMat{<:AbstractFloat},T<:AbstractVector{E}} <: VarEstimator{T}
    n :: Int
    S :: T
end

NaiveVar(::Type{T}, ::Tuple{Int}) where {T<:AbstractFloat} = NaiveVar(0, Vector{Vector{T}}())
NaiveVar(::Type{T}, ::Tuple{Int,Int}) where {T<:AbstractFloat} = NaiveVar(0, Vector{Matrix{T}}())

function add_sample!(nc::NaiveVar, s::AbstractVecOrMat)
    nc.n += 1
    push!(nc.S, s)
end

function reset!(nc::NaiveVar{T}) where {T}
    nc.n = 0
    nc.S = T()
end

function get_var(nc::NaiveVar)
    @assert nc.n >= 2 "Cannot get variance with only one sample"
    return Statistics.var(nc.S)
end

# Ref： https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_var_estimator.hpp
mutable struct WelfordVar{T<:AbstractVecOrMat{<:AbstractFloat}} <: VarEstimator{T}
    n :: Int
    μ :: T
    M :: T
    # TODO: implement temporary `δ` as `WelfordCov`
end

WelfordVar(::Type{T}, sz::Union{Tuple{Int},Tuple{Int,Int}}) where {T<:AbstractFloat} = WelfordVar(0, zeros(T, sz), zeros(T, sz))

function reset!(wv::WelfordVar{<:AbstractVecOrMat{T}}) where {T<:AbstractFloat}
    wv.n = 0
    wv.μ .= zero(T)
    wv.M .= zero(T)
end

function add_sample!(wv::WelfordVar, s::AbstractVecOrMat)
    wv.n += 1
    @unpack μ, M, n = wv
    for i in eachindex(s)
        δ = s[i] - μ[i]
        μ[i] += δ / n
        M[i] += δ * (s[i] - μ[i])
    end
end

# https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/var_adaptation.hpp
function get_var(wv::WelfordVar{<:AbstractVecOrMat{T}}) where {T<:AbstractFloat}
    n, M = T(wv.n), wv.M
    @assert n >= 2 "Cannot get covariance with only one sample"
    return (n / ((n + 5) * (n - 1))) .* M .+ T(1e-3) * (5 / (n + 5))
end

## Covariance estimator

abstract type CovEstimator{T} end

# NOTE: This naive covariance estimator is used only in testing.
mutable struct NaiveCov{T<:AbstractVector{<:AbstractVector{<:AbstractFloat}}} <: CovEstimator{T}
    n :: Int
    S :: T
end

NaiveCov(::Type{T}, sz::Union{Tuple{Int}}) where {T<:AbstractFloat} = NaiveCov(0, Vector{Vector{T}}())

function add_sample!(nc::NaiveCov, s::AbstractVector)
    nc.n += 1
    push!(nc.S, s)
end

function get_cov(nc::NaiveCov)
    @assert nc.n >= 2 "Cannot get covariance with only one sample"
    return Statistics.cov(nc.S)
end

function reset!(nc::NaiveCov{T}) where {T}
    nc.n = 0
    nc.S = T()
end

# Ref: https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_covar_estimator.hpp
mutable struct WelfordCov{T<:AbstractFloat} <: CovEstimator{T}
    n :: Int
    μ :: Vector{T}
    M :: Matrix{T}
    δ :: Vector{T} # temporary
end

function WelfordCov(::Type{T}, sz::Tuple{Int}) where {T<:AbstractFloat}
    d, = sz
    WelfordCov(0, zeros(T, d), zeros(T, d, d), zeros(T, d))
end

function reset!(wc::WelfordCov{T}) where {T<:AbstractFloat}
    wc.n = 0
    wc.μ .= zero(T)
    wc.M .= zero(T)
end

function add_sample!(wc::WelfordCov, s::AbstractVector)
    wc.n += 1
    @unpack δ, μ, n, M = wc
    δ .= s .- μ
    μ .+= δ ./ n
    M .+= (s .- μ) .* δ'
end

# Ref: https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/covar_adaptation.hpp
function get_cov(wc::WelfordCov{T}) where {T<:AbstractFloat}
    n, M = T(wc.n), wc.M
    @assert n >= 2 "Cannot get variance with only one sample"
    return (n / ((n + 5) * (n - 1))) .* M + T(1e-3) * (5 / (n + 5)) * LinearAlgebra.I
end

###
### Preconditioning matrix adaptors.
###

abstract type AbstractPreconditioner <: AbstractAdaptor end

initialize!(adaptor::T, n_adapts::Int) where {T<:AbstractPreconditioner} = nothing
finalize!(adaptor::T) where {T<:AbstractPreconditioner} = nothing

####
#### Preconditioning matrix adaption implementation.
####

# Unit
struct UnitPreconditioner{T} <: AbstractPreconditioner end
Base.show(io::IO, ::UnitPreconditioner) = print(io, "UnitPreconditioner")

UnitPreconditioner(::Type{T}=Float64) where {T} = UnitPreconditioner{T}()

string(::UnitPreconditioner) = "I"
reset!(::UnitPreconditioner) = nothing
getM⁻¹(dpc::UnitPreconditioner{T}) where {T<:AbstractFloat} = UniformScaling{T}(one(T))
adapt!(
    ::UnitPreconditioner,
    ::AbstractVecOrMat{<:AbstractFloat},
    ::AbstractScalarOrVec{<:AbstractFloat},
    is_update::Bool=true
) = nothing

mutable struct DiagPreconditioner{
    AT<:AbstractVecOrMat{<:AbstractFloat},
    TEst<:VarEstimator{AT}
} <: AbstractPreconditioner
    n_min   :: Int
    ve      :: TEst
    var     :: AT
end
Base.show(io::IO, ::DiagPreconditioner) = print(io, "DiagPreconditioner")

# Diagonal
DiagPreconditioner(sz::Tuple{Vararg{Int}}, n_min::Int=10) = DiagPreconditioner(Float64, sz, n_min)
function DiagPreconditioner(
    ::Type{T},
    sz::Tuple{Vararg{Int}};
    n_min::Int=10,
    var=ones(T, sz)
) where {T}
    ve = WelfordVar(T, sz)
    return DiagPreconditioner(n_min, ve, var)
end

string(dpc::DiagPreconditioner) = string(dpc.var)
reset!(dpc::DiagPreconditioner) = reset!(dpc.ve)
getM⁻¹(dpc::DiagPreconditioner) = dpc.var

function adapt!(
    dpc::DiagPreconditioner,
    θ::AbstractVecOrMat{<:AbstractFloat},
    α::AbstractScalarOrVec{<:AbstractFloat},
    is_update::Bool=true
)
    resize!(dpc, θ)
    add_sample!(dpc.ve, θ)
    if dpc.ve.n >= dpc.n_min && is_update
        # Can be made inplace
        dpc.var .= get_var(dpc.ve)
    end
end

# Dense
mutable struct DensePreconditioner{
    T<:AbstractFloat,
    TEst<:CovEstimator{T}
} <: AbstractPreconditioner
    n_min :: Int
    ce    :: TEst
    covar :: Matrix{T}
end
Base.show(io::IO, ::DensePreconditioner) = print(io, "DensePreconditioner")

DensePreconditioner(sz::Tuple{Int}, n_min::Int=10) = DensePreconditioner(Float64, sz, n_min)
function DensePreconditioner(
    ::Type{T},
    sz::Tuple{Int};
    n_min::Int=10,
    covar=LinearAlgebra.diagm(0 => ones(T, sz))
) where {T}
    ce = WelfordCov(T, sz)
    # TODO: take use of the line below when we have an interface to set which covariance estimator to use
    # ce = NaiveCov(T)
    return DensePreconditioner(n_min, ce, covar)
end

string(dpc::DensePreconditioner) = string(LinearAlgebra.diag(dpc.covar))
reset!(dpc::DensePreconditioner) = reset!(dpc.ce)
getM⁻¹(dpc::DensePreconditioner) = dpc.covar

function adapt!(
    dpc::DensePreconditioner,
    θ::AbstractVecOrMat{<:AbstractFloat},
    α::AbstractScalarOrVec{<:AbstractFloat},
    is_update::Bool=true
)
    resize!(dpc, θ)
    add_sample!(dpc.ce, θ)
    if dpc.ce.n >= dpc.n_min && is_update
        # Can be made inplace
        dpc.covar .= get_cov(dpc.ce)
    end
end

# Resize pre-conditioner if necessary.
Base.resize!(
    pc::UnitPreconditioner,
    θ::AbstractVecOrMat
) = nothing

function Base.resize!(
    dpc::DiagPreconditioner,
    θ::AbstractVecOrMat{T}
) where {T<:AbstractFloat}
    if size(θ) != size(dpc.var)
        @assert dpc.ve.n == 0 "Cannot resize a var estimator when it contains samples."
        dpc.ve = WelfordVar(T, size(θ))
        dpc.var = ones(T, size(θ))
    end
end

function Base.resize!(
    dpc::DensePreconditioner,
    θ::AbstractVecOrMat{T}
) where {T<:AbstractFloat}
    if length(θ) != size(dpc.covar,1)
        @assert dpc.ce.n == 0 "Cannot resize a var estimator when it contains samples."
        dpc.ce = WelfordCov(T, size(θ))
        dpc.covar = LinearAlgebra.diagm(0 => ones(T, length(θ)))
    end
end

####
#### Preconditioning mass matrix.
####

abstract type AbstractMetric end

_string_M⁻¹(mat::AbstractMatrix, n_chars::Int=32) = _string_M⁻¹(LinearAlgebra.diag(mat), n_chars)
function _string_M⁻¹(vec::AbstractVector, n_chars::Int=32)
    s_vec = string(vec)
    l = length(s_vec)
    s_dots = " ...]"
    n_diag_chars = n_chars - length(s_dots)
    return s_vec[1:min(n_diag_chars,end)] * (l > n_diag_chars ? s_dots : "")
end

struct UnitEuclideanMetric{T,A<:Union{Tuple{Int},Tuple{Int,Int}}} <: AbstractMetric
    M⁻¹::UniformScaling{T}
    size::A
end

UnitEuclideanMetric(::Type{T}, sz) where {T} = UnitEuclideanMetric(UniformScaling{T}(one(T)), sz)
UnitEuclideanMetric(sz) = UnitEuclideanMetric(Float64, sz)
UnitEuclideanMetric(::Type{T}, dim::Int) where {T} = UnitEuclideanMetric(UniformScaling{T}(one(T)), (dim,))
UnitEuclideanMetric(dim::Int) = UnitEuclideanMetric(Float64, (dim,))

renew(ue::UnitEuclideanMetric, M⁻¹) = UnitEuclideanMetric(M⁻¹, ue.size)

Base.size(e::UnitEuclideanMetric) = e.size
Base.size(e::UnitEuclideanMetric, dim::Int) = e.size[dim]
Base.show(io::IO, uem::UnitEuclideanMetric) = print(io, "UnitEuclideanMetric($(_string_M⁻¹(ones(uem.size))))")

struct DiagEuclideanMetric{T,A<:AbstractVecOrMat{T}} <: AbstractMetric
    # Diagnal of the inverse of the mass matrix
    M⁻¹     ::  A
    # Sqare root of the inverse of the mass matrix
    sqrtM⁻¹ ::  A
    # Pre-allocation for intermediate variables
    _temp   ::  A
end

function DiagEuclideanMetric(M⁻¹::AbstractVecOrMat{T}) where {T<:AbstractFloat}
    return DiagEuclideanMetric(M⁻¹, sqrt.(M⁻¹), similar(M⁻¹))
end
DiagEuclideanMetric(::Type{T}, sz) where {T} = DiagEuclideanMetric(ones(T, sz...))
DiagEuclideanMetric(sz) = DiagEuclideanMetric(Float64, sz)
DiagEuclideanMetric(::Type{T}, dim::Int) where {T} = DiagEuclideanMetric(ones(T, dim))
DiagEuclideanMetric(dim::Int) = DiagEuclideanMetric(Float64, dim)

renew(ue::DiagEuclideanMetric, M⁻¹) = DiagEuclideanMetric(M⁻¹)

Base.size(e::DiagEuclideanMetric, dim...) = size(e.M⁻¹, dim...)
Base.show(io::IO, dem::DiagEuclideanMetric) = print(io, "DiagEuclideanMetric($(_string_M⁻¹(dem.M⁻¹)))")

struct DenseEuclideanMetric{
    T,
    AV<:AbstractVecOrMat{T},
    AM<:Union{AbstractMatrix{T},AbstractArray{T,3}},
    TcholM⁻¹<:UpperTriangular{T},
} <: AbstractMetric
    # Inverse of the mass matrix
    M⁻¹::AM
    # U of the Cholesky decomposition of the mass matrix
    cholM⁻¹::TcholM⁻¹
    # Pre-allocation for intermediate variables
    _temp::AV
end

# TODO: make dense mass matrix support matrix-mode parallel
function DenseEuclideanMetric(M⁻¹::Union{AbstractMatrix{T},AbstractArray{T,3}}) where {T<:AbstractFloat}
    _temp = Vector{T}(undef, Base.front(size(M⁻¹)))
    return DenseEuclideanMetric(M⁻¹, cholesky(Symmetric(M⁻¹)).U, _temp)
end
DenseEuclideanMetric(::Type{T}, D::Int) where {T} = DenseEuclideanMetric(Matrix{T}(I, D, D))
DenseEuclideanMetric(D::Int) = DenseEuclideanMetric(Float64, D)
DenseEuclideanMetric(::Type{T}, sz::Tuple{Int}) where {T} = DenseEuclideanMetric(Matrix{T}(I, first(sz), first(sz)))
DenseEuclideanMetric(sz::Tuple{Int}) = DenseEuclideanMetric(Float64, D)

renew(ue::DenseEuclideanMetric, M⁻¹) = DenseEuclideanMetric(M⁻¹)

Base.size(e::DenseEuclideanMetric, dim...) = size(e._temp, dim...)
Base.show(io::IO, dem::DenseEuclideanMetric) = print(io, "DenseEuclideanMetric(diag=$(_string_M⁻¹(dem.M⁻¹)))")

# `rand` functions for `metric` types.

function _rand(
    rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
    metric::UnitEuclideanMetric{T}
) where {T}
    r = randn(rng, T, size(metric)...)
    return r
end

function _rand(
    rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
    metric::DiagEuclideanMetric{T}
) where {T}
    r = randn(rng, T, size(metric)...)
    r ./= metric.sqrtM⁻¹
    return r
end

function _rand(
    rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
    metric::DenseEuclideanMetric{T}
) where {T}
    r = randn(rng, T, size(metric)...)
    ldiv!(metric.cholM⁻¹, r)
    return r
end

Base.rand(rng::AbstractRNG, metric::AbstractMetric) = _rand(rng, metric)    # this disambiguity is required by Random.rand
Base.rand(rng::AbstractVector{<:AbstractRNG}, metric::AbstractMetric) = _rand(rng, metric)
Base.rand(metric::AbstractMetric) = rand(GLOBAL_RNG, metric)

####
#### Preconditioner constructors
####

Preconditioner(m::UnitEuclideanMetric{T}) where {T} = UnitPreconditioner(T)
Preconditioner(m::DiagEuclideanMetric{T}) where {T} = DiagPreconditioner(T, size(m); var=copy(m.M⁻¹))
Preconditioner(m::DenseEuclideanMetric{T}) where {T} = DensePreconditioner(T, size(m); covar=copy(m.M⁻¹))

Preconditioner(m::Type{TM}, sz::Tuple{Vararg{Int}}=(2,)) where {TM<:AbstractMetric} = Preconditioner(Float64, m, sz)
Preconditioner(::Type{T}, ::Type{TM}, sz::Tuple{Vararg{Int}}=(2,)) where {T<:AbstractFloat, TM<:AbstractMetric} = Preconditioner(TM(T, sz))


text:
type StorePutKey{T} <: ResourceKey
  priority :: Int
  id :: UInt
  item :: T
end

type StoreGetKey <: ResourceKey
  priority :: Int
  id :: UInt
  filter :: Function
end

type Store{T} <: AbstractResource
  capacity :: UInt
  items :: Set{T}
  seid :: UInt
  put_queue :: PriorityQueue{Event, StorePutKey{T}}
  get_queue :: PriorityQueue{Event, StoreGetKey}
  function Store(capacity::UInt=typemax(UInt))
    sto = new()
    sto.capacity = capacity
    sto.items = Set{T}()
    sto.seid = zero(UInt)
    sto.put_queue = PriorityQueue(Event, StorePutKey{T})
    sto.get_queue = PriorityQueue(Event, StoreGetKey)
    return sto
  end
end

function Store(t::Type)
  Store{t}()
end

function put{T}(sim::Simulation, sto::Store{T}, item::T; priority::Int=0) :: Event
  put_ev = Event()
  sto.put_queue[put_ev] = StorePutKey(priority, sto.seid+=one(UInt), item)
  append_callback(put_ev, trigger_get, sto)
  trigger_put(sim, put_ev, sto)
  return put_ev
end

function get_any_item{T}(::T) :: Bool
  return true
end

function get{T}(sim::Simulation, sto::Store{T}, filter::Function=get_any_item; priority::Int=0) :: Event
  get_ev = Event()
  sto.get_queue[get_ev] = StoreGetKey(priority, sto.seid+=one(UInt), filter)
  append_callback(get_ev, trigger_put, sto)
  trigger_get(sim, get_ev, sto)
  return get_ev
end

function do_put{T}(sim::Simulation, sto::Store{T}, put_ev::Event, key::StorePutKey{T}) :: Bool
  if length(sto.items) < sto.capacity
    push!(sto.items, key.item)
    schedule(sim, put_ev)
  end
  return false
end

function do_get{T}(sim::Simulation, sto::Store{T}, get_ev::Event, key::StoreGetKey) :: Bool
  for item in sto.items
    if key.filter(item)
      delete!(sto.items, item)
      schedule(sim, get_ev, value=item)
      break
    end
  end
  return true
end


text:
#=
ellipse_im.jl
Copyright 2019-03-05, Jeff Fessler, University of Michigan
=#

export ellipse_im, ellipse_im_params

# using MIRT: image_geom, disk_phantom_params, downsample2, rotate2d

"""
    phantom = ellipse_im(ig, params ;
	rot=0, oversample=1, hu_scale=1, replace=false)

Generate ellipse phantom image from parameters:

`[x_center y_center x_radius y_radius angle_degrees amplitude]`

in
- `ig`		from `image_geom()`
- `params`	`[ne 6]` ellipse parameters

# Arguments
- `rot`			rotate ellipses by this amount [degrees]
- `oversample`	oversampling factor, for grayscale boundaries
- `hu_scale`	use 1000 to scale shepp-logan to HU
- `replace`		replace ellipse values if true, else add
- `how`			`:fast` is the only option

out
- `phantom`		[nx ny]	image (Float32)

note: `op ellipse` in aspire with `nsub=3` is `oversample=4 = 2^(3-1)` here

"""
function ellipse_im(
	ig::ImageGeom,
	params::AbstractMatrix{<:Real} ;
	rot::Real = 0,
	oversample::Int = 1,
	hu_scale::Real = 1,
	replace::Bool = false,
	how::Symbol = :fast, # todo
)

	if size(params,2) != 6
		throw("bad ellipse parameter vector size")
	end

	params[:,6] .*= hu_scale

#=
	if how === :fast && oversample == 1
		warning("ignoring :fast option for oversample=1")
		how = :slow
	end
=#

	if oversample > 1
		ig = ig.over(oversample)
	end
	args = (ig.nx, ig.ny, params, ig.dx, ig.dy, ig.offset_x, ig.offset_y,
		rot, oversample, replace)

	if how === :fast
		phantom = ellipse_im_fast(args...)
#	elseif how === :slow
#		phantom = ellipse_im_slow(args...)
	else
		throw("bad how $how")
	end

	if oversample > 1
		phantom = downsample2(phantom, oversample)
	end

	return phantom
end

#=
%
% ellipse_im_slow()
% brute force fine grid - can use lots of memory
%
function phantom = ellipse_im_slow(nx, ny, params, dx, dy, ...
	#offset_x, offset_y, rot, over, replace)

if size(params,2) ~= 6
	fail('bad ellipse parameter vector size')
end

% optional rotation of ellipse parameters
if rot ~= 0
	th = deg2rad(rot)
	cx = params[:,1)
	cy = params[:,2)
	params[:,1) = cx * cos(th) + cy * sin(th)
	params[:,2) = -cx * sin(th) + cy * cos(th)
	params[:,5) = params[:,5) + rot
	clear cx cy th
end

wx = (nx*over-1)/2 + offset_x * over
wy = (ny*over-1)/2 + offset_y * over
xx = ((0:nx*over-1) - wx) / over * dx
yy = ((0:ny*over-1) - wy) / over * dy
[xx yy] = ndgrid(xx, yy); % fine grid, equally spaced

phantom = zeros(nx*over, ny*over, 'single'); % fine array

ticker reset
ne = nrow(params)
for ie = 1:ne
	ticker(mfilename, ie, ne)

	ell = params[ie, :)
	cx = ell(1);	rx = ell(3)
	cy = ell(2);	ry = ell(4)
	theta = deg2rad(ell(5))
	[xr yr] = rot2(xx-cx, yy-cy, theta)
	tmp = (xr / rx).^2 + (yr / ry).^2 <= 1

	if replace
		phantom(tmp > 0) = ell(6)
	else
		phantom = phantom + ell(6) * tmp
	end
end

phantom = downsample2(phantom, over)
=#

"""
phantom = ellipse_im_fast()
"""
function ellipse_im_fast(nx, ny, params_in, dx, dy,
		offset_x, offset_y, rot, over, replace)

	params = copy(params_in)

	# optional rotation
	if rot != 0
		th = rot * 180/pi
		cx = params[:,1]
		cy = params[:,2]
		params[:,1] = cx * cos(th) + cy * sin(th)
		params[:,2] = -cx * sin(th) + cy * cos(th)
		params[:,5] .+= rot
	end

	phantom = zeros(Float32, nx, ny)

	wx = (nx-1)/2 + offset_x
	wy = (ny-1)/2 + offset_y
	x1 = ((0:nx-1) .- wx) * dx
	y1 = ((0:ny-1) .- wy) * dy
	(xx, yy) = ndgrid(x1, y1)

	hx = abs(dx) / 2
	hy = abs(dy) / 2

	for ie = 1:size(params,1)

		ell = Float32.(params[ie, :])
		cx = ell[1]; rx = ell[3]
		cy = ell[2]; ry = ell[4]
		theta = ell[5] * Float32(pi/180)
		value = Float32(ell[6])

		xs = xx .- cx # shift per ellipse center
		ys = yy .- cy

		# coordinates of "outer" corner of each pixel, relative to ellipse center
		xo = xs + sign.(xs) * hx
		yo = ys + sign.(ys) * hy

		# voxels that are entirely inside the ellipse:
		(xr, yr) = rotate2d(xo, yo, theta)
		is_inside = (xr / rx).^2 + (yr / ry).^2 .<= 1

		if replace
			phantom[is_inside] .= value
		else
			phantom += value * is_inside
		end

	end # ie loop

#@show typeof(phantom)
	return phantom
end # ellipse_im_fast()

"""
    phantom = ellipse_im(nx, dx, params ; kwarg...)

square image of size `nx` by `nx`,
specifying pixel size `dx` and ellipse `params`
"""
function ellipse_im(nx::Int, dx::Real, params ; kwarg...)
	ig = image_geom(nx=nx, dx=1)
	return ellipse_im(ig, params ; kwarg...)
end

"""
    phantom = ellipse_im(nx::Int, params ; kwarg...)

square image of size `nx` by `nx` with
pixel size `dx=1` and ellipse `params`
"""
function ellipse_im(nx::Int, params ; kwarg...)
	return ellipse_im(nx, 1., params ; kwarg...)
end

"""
    phantom = ellipse_im(nx::Int ; ny::Int=nx, dx::Real=1, kwarg...)

image of size `nx` by `ny` (default `nx`) with specified `dx` (default 1),
defaults to `:shepplogan_emis`
"""
function ellipse_im(nx::Int ; ny::Int=nx, dx::Real=1, kwarg...)
	ig = image_geom(nx=nx, ny=ny, dx=dx)
	return ellipse_im(ig, :shepplogan_emis ; kwarg...)
end

"""
    phantom = ellipse_im(nx::Int, ny::Int ; kwarg...)

`:shepplogan_emis` of size `nx` by `ny`
"""
function ellipse_im(nx::Int, ny::Int ; kwarg...)
	return ellipse_im(nx, ny=ny, dx=1. ; kwarg...)
end

"""
    phantom = ellipse_im(ig, code ; kwarg...)

`code = :shepplogan | :shepplogan_emis | :shepplogan_brainweb | :southpark`
"""
function ellipse_im(ig::ImageGeom, params::Symbol ; kwarg...)
	params = ellipse_im_params(ig, params)
	return ellipse_im(ig, params ; kwarg...)
end

"""
    phantom = ellipse_im(ig ; kwarg...)

`:shepplogan` (default) for given image geometry `ig`
"""
function ellipse_im(ig::ImageGeom ; kwarg...)
	return ellipse_im(ig, :shepplogan; kwarg...)
end

"""
    params = ellipse_im_params(ig::ImageGeom, params::Symbol)

`code = :shepplogan | :shepplogan_emis | :shepplogan_brainweb | :southpark | :disks`
"""
function ellipse_im_params(ig::ImageGeom, params::Symbol)
	if params === :disks
		params = disk_phantom_params(fov=ig.fovs[1])
	elseif params === :shepplogan || params === :kak
		params = shepp_logan_parameters(ig.fovs..., case=:kak)
	elseif params === :shepplogan_emis || params === :emis
		params = shepp_logan_parameters(ig.fovs..., case=:emis)
	elseif params === :shepplogan_brainweb || params === :brainweb
		params = shepp_logan_parameters(ig.fovs..., case=:brainweb)
	elseif params === :southpark
		params = south_park_parameters(fov=ig.fovs[1])
	else
		throw("bad phantom symbol $params")
	end
	return params
end

"""
params = shepp_logan_parameters(xfov, yfov ; case::Symbol)

parameters from Kak and Slaney text, p. 255

the first four columns are unitless "fractions of field of view"
"""
function shepp_logan_parameters(xfov::Real, yfov::Real ; case::Symbol=:kak)
	params = [
	0		0		0.92	0.69	90	2
	0		-0.0184	0.874	0.6624	90	-0.98
	0.22	0		0.31	0.11	72	-0.02
	-0.22	0		0.41	0.16	108	-0.02
	0		0.35	0.25	0.21	90	0.01
	0		0.1		0.046	0.046	0	0.01
	0		-0.1	0.046	0.046	0	0.01
	-0.08	-0.605	0.046	0.023	0	0.01
	0		-0.605	0.023	0.023	0	0.01
	0.06	-0.605	0.046	0.023	90	0.01]

	params[:,[1,3]] .*= xfov/2
	params[:,[2,4]] .*= yfov/2

	if case === :emis
		params[:,6] = [1, 1, -2, 2, 3, 4, 5, 6, 1, 1]
	elseif case === :brainweb
		params[:,6] = [1, 0, 2, 3, 4, 5, 6, 7, 8, 9] # brainweb uses index 1-10
	elseif case != :kak
		throw("bad phantom case $case")
	end

	return Float32.(params)
end

"""
param = south_park_parameters( ; fov::Real = 100)
"""
function south_park_parameters( ; fov::Real = 100)
	xell = [
		0. 0 85 115 0 100
		0 -60 30 20 0 -80 # mouth
		30 20 25 35 30 20 # eyes
		-30 20 25 35 -30 20
		35 25 7 7 0 -100 # pupils
		-15 25 7 7 0 -100
		0 75 60 15 0 -50] # hat
	xell[:,1:4] .*= fov/256
	return Float32.(xell)
end


text:
function walkingPulseFrame(leddata, audio_samp, raw_spec)
    if !("walking_offset" in keys(leddata.state))
        leddata.state["walking_offset"] = 0
    end
    if ("walking_pulse_spec" in keys(leddata.state))
        spec = leddata.state["walking_pulse_spec"]
        spec, maxspec = binFFT(raw_spec, 25)
    else
        spec,maxspec = binFFT(raw_spec, 25)
        leddata.state["walking_pulse_spec"] = spec
    end

    spec./=fft_scale[]
    spec[isnan.(spec)] = 0
    spec .= abs.(spec)
    pulse_length = 15
    default_color = colorant"black"
    pulse_coef = min(spec[1]+spec[2]+mean(spec[3:5])/3, 1)
    pulse_coef = pulse_coef>0.05 ? pulse_coef : 0
    pulse_shape = [i%pulse_length/4==0 ? abs(rand()*0.5*pulse_coef) : abs(sin(pi/4*i)*pulse_coef) for i in linspace(0,4,pulse_length)]
    pulse_shape[isnan.(pulse_shape)] = 0
    pulse = zeros(UInt8, pulse_length, 3)
    pulse_shape .= (pulse_shape).^0.25
    for i in 1:size(pulse,1)
        pulse[i, :] .= hsl_to_rgb(240/360, 0.95*pulse_shape[i]+0.05, 0.45*pulse_shape[i]+0.05)
    end
    for channel in leddata.channels
        setindex!(channel, default_color, 1:length(channel))
        numPulses = floor(Int, length(channel)/30)
        starts = (1:floor(Int, length(channel)/numPulses ):length(channel))[1:end-1] + round(Int, leddata.state["walking_offset"]*length(channel)/100)
        for start in starts
            assgn = (start:(start+size(pulse,1)-1)).%length(channel) + 1
            for i in eachindex(assgn)
                for j in 1:3
                    channel[assgn[i],j] = pulse[i,j]
                end
            end
        end
        update!(channel)
    end

    if leddata.state["walking_offset"] == 400
        leddata.state["walking_offset"] -= 100
    end
    leddata.state["walking_offset"] += 0.5
end


text:
include("../src/algorithm.jl")

function unit_test_findZerosWithSubdivision(a, b, f, expected)
    pass = 1
    box = findZerosWithSubdivision(a, b, f)
    re1 = real.(box[1][1])
    im1 = imag.(box[1][1])
    re2 = real.(box[1][2])
    im2 = imag.(box[1][2])

    if (re1 < real.(expected) < re2) && (im2 < imag.(expected) < im1)
        pass = 0
    end

    return pass
end

function test_singlezeros(a, b, f, err=0.01, increment=err/100, dif=0.5)
    pass = 0
    try boxes = findZerosWithSubdivision(a, b, f, err, increment, dif)
        # Ensure every box contains exactly one zero and is within error range
        for b in boxes
            if (countZeros(b[1], b[2], f, increment, dif) != 1 ) || getError(b[1], b[2]) > err
                pass = -1
            end
        end
        return pass
    # When an exception is thrown, ensure that the box returned contains zeros
    catch e
        box = e.biggerBox
        if countZeros(box[1], box[2], f, increment, dif) <= 0
            pass = 1
        end
        return pass
    end

end

function test_multiplezeros(a, b, f, err=0.01, increment=err/100, dif=0.5)
    pass = 0
    try boxes = findZerosWithSubdivision(a, b, f, err, increment, dif)
        # Ensure every box contains at least one zero and is within error range
        for b in boxes
            if (countZeros(b[1], b[2], f, increment, dif) < 1 ) || getError(b[1], b[2]) > err
                pass = -1
            end
        end
        return pass

    # When an exception is thrown, ensure that the box returned contains zeros
    catch e
        box = e.biggerBox
        if countZeros(box[1], box[2], f, increment, dif) <= 0
            pass = 1
        end
        return pass
    end

end

function randomised_countzeros(n)
    pass = 0
    for i in 1:n
        roots = rand(-9:.1:9,6)
        function generatedPolynomial(x)
            return (x-roots[1] - roots[2]*im)*(x-roots[3] - roots[4]*im)*(x-roots[5] - roots[6]*im)
        end
        randomPolynomial = generatedPolynomial
        if countZeros(-10 + 10im,  10 - 10im, randomPolynomial, 0.001, 0.5) != 3
            print(roots)
            pass = 1
        end
    end
    return pass
end

function randomised_subdivision(n)
    passes = []
    for i in 1:n
        roots = rand(-5:.001:5,6)
        function generatedPolynomial(x)
            return (x-roots[1] - roots[2]*im)*(x-roots[3] - roots[4]*im)*(x-roots[5] - roots[6]*im)
        end
        randomPolynomial = generatedPolynomial
        push!(passes, test_singlezeros(-10 + 10im, 10 - 10im, randomPolynomial))
    end

    return passes
end


text:
"""
    AbstractNLPEvaluator

AbstractNLPEvaluator implements the bridge between the
problem formulation (see [`AbstractFormulation`](@ref)) and the optimization
solver. Once the problem formulation bridged, the evaluator allows
to evaluate:
- the objective;
- the gradient of the objective;
- the constraints;
- the Jacobian of the constraints;
- the Jacobian-vector and transpose-Jacobian vector products of the constraints;
- the Hessian of the objective;
- the Hessian of the Lagrangian.

"""
abstract type AbstractNLPEvaluator end

abstract type AbstractNLPAttribute end

"""
    Variables <: AbstractNLPAttribute end

Attribute corresponding to the optimization variables attached
to a given [`AbstractNLPEvaluator`](@ref).
"""
struct Variables <: AbstractNLPAttribute end

"""
    Constraints <: AbstractNLPAttribute end

Attribute corresponding to the constraints  attached
to a given [`AbstractNLPEvaluator`](@ref).
"""
struct Constraints <: AbstractNLPAttribute end

"""
    AutoDiffBackend <: AbstractNLPAttribute end

Attribute corresponding to the autodiff backend used
inside the [`AbstractNLPEvaluator`](@ref).
"""
struct AutoDiffBackend <: AbstractNLPAttribute end

"""
    n_variables(nlp::AbstractNLPEvaluator)
Get the number of variables in the problem.
"""
function n_variables end

"""
    n_constraints(nlp::AbstractNLPEvaluator)
Get the number of constraints in the problem.
"""
function n_constraints end

# Callbacks
"""
    objective(nlp::AbstractNLPEvaluator, u)::Float64

Evaluate the objective at given variable `u`.
"""
function objective end

"""
    gradient!(nlp::AbstractNLPEvaluator, g, u)

Evaluate the gradient of the objective, at given variable `u`.
Store the result inplace in the vector `g`.

## Note
The vector `g` should have the same dimension as `u`.

"""
function gradient! end

"""
    constraint!(nlp::AbstractNLPEvaluator, cons, u)

Evaluate the constraints of the problem at given variable `u`. Store
the result inplace, in the vector `cons`.

## Note
The vector `cons` should have the same dimension as the result
returned by `n_constraints(nlp)`.

"""
function constraint! end

"""
    jacobian_structure!(nlp::AbstractNLPEvaluator, rows, cols)

Return the sparsity pattern of the Jacobian matrix. Stores
the results inplace, in the vectors `rows` and `cols` (whose
dimension should match the number of non-zero in the Jacobian matrix).

"""
function jacobian_structure! end

"""
    jacobian!(nlp::AbstractNLPEvaluator, jac, u)

Evaluate the Jacobian of the constraints, at variable `u`.
Store the result inplace, in the `m x n` dense matrix `jac`.

"""
function jacobian! end

"""
    jprod!(nlp::AbstractNLPEvaluator, jv, u, v)

Evaluate the Jacobian-vector product ``J v`` of the constraints.
The vector `jv` is modified inplace.

Let `(n, m) = n_variables(nlp), n_constraints(nlp)`.

* `u` is a vector with dimension `n`
* `v` is a vector with dimension `n`
* `jv` is a vector with dimension `m`

"""
function jprod! end

"""
    jtprod!(nlp::AbstractNLPEvaluator, jv, u, v)

Evaluate the transpose Jacobian-vector product ``J^{T} v`` of the constraints.
The vector `jv` is modified inplace.

Let `(n, m) = n_variables(nlp), n_constraints(nlp)`.

* `u` is a vector with dimension `n`
* `v` is a vector with dimension `m`
* `jv` is a vector with dimension `n`

"""
function jtprod! end

"""
    ojtprod!(nlp::AbstractNLPEvaluator, jv, u, σ, v)

Evaluate the transpose Jacobian-vector product `J' * [σ ; v]`,
with `J` the Jacobian of the vector `[f(x); h(x)]`.
`f(x)` is the current objective and `h(x)` constraints.
The vector `jv` is modified inplace.

Let `(n, m) = n_variables(nlp), n_constraints(nlp)`.

* `jv` is a vector with dimension `n`
* `u` is a vector with dimension `n`
* `σ` is a scalar
* `v` is a vector with dimension `m`

"""
function ojtprod! end

"""
    hessian!(nlp::AbstractNLPEvaluator, H, u)

Evaluate the Hessian `∇²f(u)` of the objective function `f(u)`.
Store the result inplace, in the `n x n` dense matrix `H`.

"""
function hessian! end

"""
    hessprod!(nlp::AbstractNLPEvaluator, hessvec, u, v)

Evaluate the Hessian-vector product `∇²f(u) * v` of the objective
evaluated at variable `u`.
Store the result inplace, in the vector `hessvec`.

## Note
The vector `hessprod` should have the same length as `u`.

"""
function hessprod! end

@doc raw"""
    hessian_lagrangian_penalty_prod!(nlp::AbstractNLPEvaluator, hessvec, u, y, σ, v, d)

Evaluate the Hessian-vector product of the Lagrangian
function ``L(u, y) = f(u) + \sum_i y_i c_i(u) + \frac{1}{2} d_i c_i(u)^2`` with a vector `v`:
```math
∇²L(u, y) ⋅ v  = σ ∇²f(u) ⋅ v + \sum_i (y_i + d_i) ∇²c_i(u) ⋅ v + \sum_i d_i ∇c_i(u)^T ∇c_i(u)
```

Store the result inplace, in the vector `hessvec`.

### Arguments

* `hessvec` is a `AbstractVector` with dimension `n`, which is modified inplace.
* `u` is a `AbstractVector` with dimension `n`, storing the current variable.
* `y` is a `AbstractVector` with dimension `n`, storing the current constraints' multipliers
* `σ` is a scalar
* `v` is a vector with dimension `n`.
* `d` is a vector with dimension `m`.
"""
function hessian_lagrangian_penalty_prod! end

# Utilities
"""
    primal_infeasibility(nlp::AbstractNLPEvaluator, u)

Return primal infeasibility associated to current model `nlp` evaluated
at variable `u`.

"""
function primal_infeasibility end

"""
    primal_infeasibility!(nlp::AbstractNLPEvaluator, cons, u)

Return primal infeasibility associated to current model `nlp` evaluated
at variable `u`. Modify vector `cons` inplace.

"""
function primal_infeasibility! end

"Return `true` if the problem is constrained, `false` otherwise."
is_constrained(nlp::AbstractNLPEvaluator) = n_constraints(nlp) > 0

"""
    constraints_type(nlp::AbstractNLPEvaluator)

Return the type of the non-linear constraints of the evaluator `nlp`,
as a `Symbol`. Result could be `:inequality` if problem has only
inequality constraints, `:equality` if problem has only equality
constraints, or `:mixed` if problem has both types of constraints.
"""
function constraints_type end

"Check if Hessian of objective is implemented."
has_hessian(nlp::AbstractNLPEvaluator) = false
"Check if Hessian of Lagrangian is implemented."
has_hessian_lagrangian(nlp::AbstractNLPEvaluator) = false

"""
    reset!(nlp::AbstractNLPEvaluator)

Reset evaluator `nlp` to default configuration.

"""
function reset! end

include("common.jl")

# Basic evaluators
include("reduced_evaluator.jl")
include("slack_evaluator.jl")
include("feasibility_evaluator.jl")
include("proxal_evaluators.jl")

# Penalty evaluators
include("penalty.jl")
include("auglag.jl")

# Bridge with MOI
include("MOI_wrapper.jl")
include("optimizers.jl")


text:
module TestExamples2DAPE

using Test
using Trixi

include("test_trixi.jl")

# pathof(Trixi) returns /path/to/Trixi/src/Trixi.jl, dirname gives the parent directory
EXAMPLES_DIR = joinpath(pathof(Trixi) |> dirname |> dirname, "examples", "tree_2d_dgsem")

@testset "Acoustic Perturbation" begin
  @trixi_testset "elixir_ape_convergence.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_ape_convergence.jl"),
      l2 = [0.0019921138796370834, 0.002090394698052287, 0.0006091925854593805, 0.0, 0.0, 0.0, 0.0],
      linf = [0.00769282588065634, 0.008276649669227254, 0.004196479023954813, 0.0, 0.0, 0.0, 0.0])
  end

  @trixi_testset "elixir_ape_gauss.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_ape_gauss.jl"),
      l2 = [0.08005276517890283, 0.08005276517890268, 0.4187202920734123, 0.0, 0.0, 0.0, 0.0],
      linf = [0.17261097190220992, 0.17261097190220973, 1.13601894068238, 0.0, 0.0, 0.0, 0.0])
  end

  @trixi_testset "elixir_ape_gaussian_source.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_ape_gaussian_source.jl"),
      l2 = [0.004296394903650806, 0.004241280404758938, 0.006269684906035964, 0.0, 0.0, 0.0, 0.0],
      linf = [0.03970270697049378, 0.04151096349298151, 0.0640019829058819, 0.0, 0.0, 0.0, 0.0])
  end

  @trixi_testset "elixir_ape_gauss_wall.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_ape_gauss_wall.jl"),
      l2 = [0.019419398248465843, 0.019510701017551826, 0.04818246051887614,
            7.382060834820337e-17, 0.0, 1.4764121669640674e-16, 1.4764121669640674e-16],
      linf = [0.18193631937316496, 0.1877464607867628, 1.0355388011792845,
              2.220446049250313e-16, 0.0, 4.440892098500626e-16, 4.440892098500626e-16])
  end

  @trixi_testset "elixir_ape_monopole.jl" begin
    @test_trixi_include(joinpath(EXAMPLES_DIR, "elixir_ape_monopole.jl"),
      l2 = [0.006816790293009947, 0.0065068948357351625, 0.008724512056168938,
            0.0009894398191644543, 0.0, 7.144325530679576e-17, 7.144325530679576e-17],
      linf = [1.000633375007386, 0.5599788929862504, 0.5738432957070382,
              0.015590137026938428, 0.0, 2.220446049250313e-16, 2.220446049250313e-16],
      maxiters=50)
  end
end

end # module


text:
# # Parallelizable reduction using `@reduce`

# !!! warning
#
#     This page is still work-in-progress.

using FLoops

using LiterateTest                                                     #src
using Test                                                             #src

# ## [`@reduce() do ... end` syntax](@id ref-reduce-do)

@test begin
    @floop for x in 1:10
        y = 2x
        @reduce() do (acc; y)
            acc += y
        end
    end
    acc
end == 110

# ### Argument symbols must be unique within a `@reduce` block

@testset_error try
    @dedent @eval begin
        @floop for x in 1:10
            @reduce() do (a; x), (b; x)
                a += x
                b *= x
            end
        end
    end
catch err
    @test err isa Exception
end

#src TODO: @reduce(a += x, b *= x) should work
#src
#src # Note that `op=` syntax does not have this restriction:
#src
#src @test begin
#src     @floop for x in 1:10
#src         @reduce(a += x, b *= x)
#src     end
#src     (a, b)
#src end == (55, 3628800)

# The argument should be manually duplicated when using the same
# variable that would be merged into multiple accumulators:

@test begin
    @floop for x in 1:10
        y = x
        @reduce() do (a; x), (b; y)
            a += x
            b *= y
        end
    end
    (a, b)
end == (55, 3628800)

# If two accumulators do not interact as in the case above, it is
# recommended to use two `@reduce() do` blocks to clarify that they
# are independent reductions:

@test begin
    @floop for x in 1:10
        @reduce() do (a; x)
            a += x
        end
        @reduce() do (b; x)
            b *= x
        end
    end
    (a, b)
end == (55, 3628800)


text:
module CUBLAS

using ..APIUtils

using ..CUDA
using ..CUDA: CUstream, cuComplex, cuDoubleComplex, libraryPropertyType, cudaDataType
using ..CUDA: libcublas, unsafe_free!, @retry_reclaim, isdebug, @sync

using GPUArrays

using LinearAlgebra

using BFloat16s

using CEnum

using Memoize

using DataStructures

# core library
include("libcublas_common.jl")
include("error.jl")
include("libcublas.jl")
include("libcublas_deprecated.jl")

# low-level wrappers
include("util.jl")
include("wrappers.jl")

# high-level integrations
include("linalg.jl")

# thread cache for task-local library handles
const thread_handles = Vector{Union{Nothing,cublasHandle_t}}()
const thread_xt_handles = Vector{Union{Nothing,cublasXtHandle_t}}()

# cache for created, but unused handles
const old_handles = DefaultDict{CuContext,Vector{cublasHandle_t}}(()->cublasHandle_t[])
const old_xt_handles = DefaultDict{Vector{CuContext},Vector{cublasHandle_t}}(()->cublasHandle_t[])

function math_mode!(handle, mode)
    flags = 0

    # https://github.com/facebookresearch/faiss/issues/1385
    if version(handle) > v"11"
        flags = CUBLAS_MATH_DISALLOW_REDUCED_PRECISION_REDUCTION
    end

    flags |= if mode == CUDA.PEDANTIC_MATH
        # prevent use of tensor cores
        if version(handle) < v"11"
            CUBLAS_DEFAULT_MATH
        else
            CUBLAS_PEDANTIC_MATH
        end
    elseif mode == CUDA.DEFAULT_MATH
        # use tensor cores, but don't reduce precision
        if version(handle) < v"11"
            CUBLAS_TENSOR_OP_MATH
        else
            CUBLAS_DEFAULT_MATH
        end
    elseif mode == CUDA.FAST_MATH
        # we'll additionally select a compute-mode with reduced precision whenever possible
        if version(handle) < v"11"
            CUBLAS_TENSOR_OP_MATH
        else
            CUBLAS_TF32_TENSOR_OP_MATH
        end
    end

    cublasSetMathMode(handle, flags)

    return
end

function handle()
    CUDA.detect_state_changes()
    tid = Threads.threadid()
    if @inbounds thread_handles[tid] === nothing
        ctx = context()
        thread_handles[tid] = get!(task_local_storage(), (:CUBLAS, ctx)) do
            handle = if isempty(old_handles[ctx])
                cublasCreate()
                # FIXME: use cublasSetWorkspace? cublasSetStream reset it.
            else
                pop!(old_handles[ctx])
            end

            finalizer(current_task()) do task
                push!(old_handles[ctx], handle)
            end
            # TODO: cublasDestroy to preserve memory, or at exit?

            cublasSetStream_v2(handle, stream())

            math_mode!(handle, CUDA.math_mode())

            handle
        end
    end
    something(@inbounds thread_handles[tid])
end

function xt_handle()
    CUDA.detect_state_changes()
    tid = Threads.threadid()
    if @inbounds thread_xt_handles[tid] === nothing
        ctxs = [context(dev) for dev in devices()]
        thread_xt_handles[tid] = get!(task_local_storage(), (:CUBLASxt, ctxs)) do
            handle = if isempty(old_xt_handles[ctxs])
                cublasXtCreate()
            else
                pop!(old_xt_handles[ctxs])
            end

            finalizer(current_task()) do task
                push!(old_xt_handles[ctxs], handle)
            end
            # TODO: cublasXtDestroy to preserve memory, or at exit?

            # select all devices
            devs = convert.(Cint, devices())
            cublasXtDeviceSelect(handle, length(devs), devs)

            handle
        end
    end
    something(@inbounds thread_xt_handles[tid])
end

@inline function set_stream(stream::CuStream)
    ctx = context()
    tls = task_local_storage()
    handle = get(tls, (:CUBLAS, ctx), nothing)
    if handle !== nothing
        cublasSetStream_v2(handle, stream)
    end
    return
end

function __init__()
    resize!(thread_handles, Threads.nthreads())
    fill!(thread_handles, nothing)

    resize!(thread_xt_handles, Threads.nthreads())
    fill!(thread_xt_handles, nothing)

    CUDA.atdevicereset() do dev
        fill!(thread_xt_handles, nothing)
    end

    CUDA.atdeviceswitch() do
        tid = Threads.threadid()
        thread_handles[tid] = nothing
    end

    CUDA.attaskswitch() do
        tid = Threads.threadid()
        thread_handles[tid] = nothing
        thread_xt_handles[tid] = nothing
    end
end

function log_message(cstr)
    # NOTE: we can't `@debug` these messages, because the logging function is called for
    #       every line... also not sure what the i!/I! prefixes mean (info?)
    # NOTE: turns out we even can't `print` these messages, as cublasXt callbacks might
    #       happen from a different thread! we could strdup + uv_async_send, but I couldn't
    #       find an easy way to attach data to that
    # TODO: use a pre-allocated lock-free global message buffer?
    len = ccall(:strlen, Csize_t, (Cstring,), cstr)
    ccall(:write, Cint, (Cint, Cstring, Csize_t), 0, cstr, len)
    return
end

function __runtime_init__()
    # enable library logging when launched with JULIA_DEBUG=CUBLAS
    if isdebug(:init, CUBLAS)
        callback = @cfunction(log_message, Nothing, (Cstring,))
        cublasSetLoggerCallback(callback)
    end
end

end


text:
using Kalman
using Base.Test

s = UnscentedState([1.0, 0.0],4*eye(2),0.1,2.0,0.0)
su,wm,wc = sigma(s)

@test s.x == [1.0,0.0]
@test s.p == [4.0 0; 0 4.0]
@test_approx_eq su[:,5] [1.0,-sqrt(0.08)]

xt = zeros(s.x)
pt = zeros(s.p)

for i in 1:size(su,2)
    xt += wm[i] * su[:,i]
    pt += wc[i] * (su[:,i]-s.x)*(su[:,i]-s.x)'
end

@test_approx_eq xt s.x
@test_approx_eq pt s.p

dt = 0.01

function fu(x)
    x1 = zeros(x)
    x1[1] = x[1] + dt*x[2]
    x1[2] = x[2] - dt*x[1]
    x1
end

fm = AdditiveUnscentedModel(fu,1e-10*eye(2))

hu(x) = [x[1]]

zm = AdditiveUnscentedObservationModel(hu,[0.1]')

kf = AdditiveUnscentedKalmanFilter(s,fm,zm)


text:
# This file was a part of Julia. License is MIT: http://julialang.org/license

using Git
using Test

include("gitutils.jl")

@test Git.version() >= v"1.7.3"

mktempdir() do dir
    cd(dir) do
        run(`$gitcmd init -q`)
        run(`$gitcmd config user.name "Julia Tester"`)
        run(`$gitcmd config user.email test@julialang.org`)
        run(`$gitcmd commit -q --allow-empty -m "initial empty commit"`)
        git_verify(Dict(), Dict(), Dict())

        # each path can have one of these content in each of head, index, work
        # for a total of length(contents)^3 = 4^3 = 64 combinations.
        # each path can be in any of these 64 "superpositions" before & after
        # for a total of 64^2 = 4096 files needed to test all transitions
        # between before and after superpositions of git repo states.

        contents = [nothing, "foo", "bar", Dict{Any,Any}("baz"=>"qux")]
        b = length(contents)
        states = [Dict([(string(k, base=b, pad=6), contents[rem(div(k,b^p),b)+1]) for k=0:(b^3)^2-1]) for p=0:5]

        git_setup(states[1:3]...)
        try Git.transact() do
            git_setup(states[4:6]...)
            throw(nothing)
        end catch x
            x === nothing || rethrow()
        end
        git_verify(states[1:3]...)
    end
end


text:
# This file is auto-generated by AWSMetadata.jl
using AWS
using AWS.AWSServices: fis
using AWS.Compat
using AWS.UUIDs

"""
    create_experiment_template(actions, client_token, description, role_arn, stop_conditions)
    create_experiment_template(actions, client_token, description, role_arn, stop_conditions, params::Dict{String,<:Any})

Creates an experiment template.  To create a template, specify the following information:
  Targets: A target can be a specific resource in your AWS environment, or one or more
resources that match criteria that you specify, for example, resources that have specific
tags.    Actions: The actions to carry out on the target. You can specify multiple actions,
the duration of each action, and when to start each action during an experiment.    Stop
conditions: If a stop condition is triggered while an experiment is running, the experiment
is automatically stopped. You can define a stop condition as a CloudWatch alarm.   For more
information, see the AWS Fault Injection Simulator User Guide.

# Arguments
- `actions`: The actions for the experiment.
- `client_token`: Unique, case-sensitive identifier that you provide to ensure the
  idempotency of the request.
- `description`: A description for the experiment template. Can contain up to 64 letters
  (A-Z and a-z).
- `role_arn`: The Amazon Resource Name (ARN) of an IAM role that grants the AWS FIS service
  permission to perform service actions on your behalf.
- `stop_conditions`: The stop conditions.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"tags"`: The tags to apply to the experiment template.
- `"targets"`: The targets for the experiment.
"""
create_experiment_template(actions, clientToken, description, roleArn, stopConditions; aws_config::AbstractAWSConfig=global_aws_config()) = fis("POST", "/experimentTemplates", Dict{String, Any}("actions"=>actions, "clientToken"=>clientToken, "description"=>description, "roleArn"=>roleArn, "stopConditions"=>stopConditions); aws_config=aws_config)
create_experiment_template(actions, clientToken, description, roleArn, stopConditions, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = fis("POST", "/experimentTemplates", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("actions"=>actions, "clientToken"=>clientToken, "description"=>description, "roleArn"=>roleArn, "stopConditions"=>stopConditions), params)); aws_config=aws_config)

"""
    delete_experiment_template(id)
    delete_experiment_template(id, params::Dict{String,<:Any})

Deletes the specified experiment template.

# Arguments
- `id`: The ID of the experiment template.

"""
delete_experiment_template(id; aws_config::AbstractAWSConfig=global_aws_config()) = fis("DELETE", "/experimentTemplates/$(id)"; aws_config=aws_config)
delete_experiment_template(id, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = fis("DELETE", "/experimentTemplates/$(id)", params; aws_config=aws_config)

"""
    get_action(id)
    get_action(id, params::Dict{String,<:Any})

Gets information about the specified AWS FIS action.

# Arguments
- `id`: The ID of the action.

"""
get_action(id; aws_config::AbstractAWSConfig=global_aws_config()) = fis("GET", "/actions/$(id)"; aws_config=aws_config)
get_action(id, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = fis("GET", "/actions/$(id)", params; aws_config=aws_config)

"""
    get_experiment(id)
    get_experiment(id, params::Dict{String,<:Any})

Gets information about the specified experiment.

# Arguments
- `id`: The ID of the experiment.

"""
get_experiment(id; aws_config::AbstractAWSConfig=global_aws_config()) = fis("GET", "/experiments/$(id)"; aws_config=aws_config)
get_experiment(id, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = fis("GET", "/experiments/$(id)", params; aws_config=aws_config)

"""
    get_experiment_template(id)
    get_experiment_template(id, params::Dict{String,<:Any})

Gets information about the specified experiment template.

# Arguments
- `id`: The ID of the experiment template.

"""
get_experiment_template(id; aws_config::AbstractAWSConfig=global_aws_config()) = fis("GET", "/experimentTemplates/$(id)"; aws_config=aws_config)
get_experiment_template(id, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = fis("GET", "/experimentTemplates/$(id)", params; aws_config=aws_config)

"""
    list_actions()
    list_actions(params::Dict{String,<:Any})

Lists the available AWS FIS actions.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"maxResults"`: The maximum number of results to return with a single call. To retrieve
  the remaining results, make another call with the returned nextToken value.
- `"nextToken"`: The token for the next page of results.
"""
list_actions(; aws_config::AbstractAWSConfig=global_aws_config()) = fis("GET", "/actions"; aws_config=aws_config)
list_actions(params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = fis("GET", "/actions", params; aws_config=aws_config)

"""
    list_experiment_templates()
    list_experiment_templates(params::Dict{String,<:Any})

Lists your experiment templates.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"maxResults"`: The maximum number of results to return with a single call. To retrieve
  the remaining results, make another call with the returned nextToken value.
- `"nextToken"`: The token for the next page of results.
"""
list_experiment_templates(; aws_config::AbstractAWSConfig=global_aws_config()) = fis("GET", "/experimentTemplates"; aws_config=aws_config)
list_experiment_templates(params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = fis("GET", "/experimentTemplates", params; aws_config=aws_config)

"""
    list_experiments()
    list_experiments(params::Dict{String,<:Any})

Lists your experiments.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"maxResults"`: The maximum number of results to return with a single call. To retrieve
  the remaining results, make another call with the returned nextToken value.
- `"nextToken"`: The token for the next page of results.
"""
list_experiments(; aws_config::AbstractAWSConfig=global_aws_config()) = fis("GET", "/experiments"; aws_config=aws_config)
list_experiments(params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = fis("GET", "/experiments", params; aws_config=aws_config)

"""
    list_tags_for_resource(resource_arn)
    list_tags_for_resource(resource_arn, params::Dict{String,<:Any})

Lists the tags for the specified resource.

# Arguments
- `resource_arn`: The Amazon Resource Name (ARN) of the resource.

"""
list_tags_for_resource(resourceArn; aws_config::AbstractAWSConfig=global_aws_config()) = fis("GET", "/tags/$(resourceArn)"; aws_config=aws_config)
list_tags_for_resource(resourceArn, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = fis("GET", "/tags/$(resourceArn)", params; aws_config=aws_config)

"""
    start_experiment(client_token, experiment_template_id)
    start_experiment(client_token, experiment_template_id, params::Dict{String,<:Any})

Starts running an experiment from the specified experiment template.

# Arguments
- `client_token`: Unique, case-sensitive identifier that you provide to ensure the
  idempotency of the request.
- `experiment_template_id`: The ID of the experiment template.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"tags"`: The tags to apply to the experiment.
"""
start_experiment(clientToken, experimentTemplateId; aws_config::AbstractAWSConfig=global_aws_config()) = fis("POST", "/experiments", Dict{String, Any}("clientToken"=>clientToken, "experimentTemplateId"=>experimentTemplateId); aws_config=aws_config)
start_experiment(clientToken, experimentTemplateId, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = fis("POST", "/experiments", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("clientToken"=>clientToken, "experimentTemplateId"=>experimentTemplateId), params)); aws_config=aws_config)

"""
    stop_experiment(id)
    stop_experiment(id, params::Dict{String,<:Any})

Stops the specified experiment.

# Arguments
- `id`: The ID of the experiment.

"""
stop_experiment(id; aws_config::AbstractAWSConfig=global_aws_config()) = fis("DELETE", "/experiments/$(id)"; aws_config=aws_config)
stop_experiment(id, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = fis("DELETE", "/experiments/$(id)", params; aws_config=aws_config)

"""
    tag_resource(resource_arn, tags)
    tag_resource(resource_arn, tags, params::Dict{String,<:Any})

Applies the specified tags to the specified resource.

# Arguments
- `resource_arn`: The Amazon Resource Name (ARN) of the resource.
- `tags`: The tags for the resource.

"""
tag_resource(resourceArn, tags; aws_config::AbstractAWSConfig=global_aws_config()) = fis("POST", "/tags/$(resourceArn)", Dict{String, Any}("tags"=>tags); aws_config=aws_config)
tag_resource(resourceArn, tags, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = fis("POST", "/tags/$(resourceArn)", Dict{String, Any}(mergewith(_merge, Dict{String, Any}("tags"=>tags), params)); aws_config=aws_config)

"""
    untag_resource(resource_arn)
    untag_resource(resource_arn, params::Dict{String,<:Any})

Removes the specified tags from the specified resource.

# Arguments
- `resource_arn`: The Amazon Resource Name (ARN) of the resource.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"tagKeys"`: The tag keys to remove.
"""
untag_resource(resourceArn; aws_config::AbstractAWSConfig=global_aws_config()) = fis("DELETE", "/tags/$(resourceArn)"; aws_config=aws_config)
untag_resource(resourceArn, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = fis("DELETE", "/tags/$(resourceArn)", params; aws_config=aws_config)

"""
    update_experiment_template(id)
    update_experiment_template(id, params::Dict{String,<:Any})

Updates the specified experiment template.

# Arguments
- `id`: The ID of the experiment template.

# Optional Parameters
Optional parameters can be passed as a `params::Dict{String,<:Any}`. Valid keys are:
- `"actions"`: The actions for the experiment.
- `"description"`: A description for the template.
- `"roleArn"`: The Amazon Resource Name (ARN) of an IAM role that grants the AWS FIS
  service permission to perform service actions on your behalf.
- `"stopConditions"`: The stop conditions for the experiment.
- `"targets"`: The targets for the experiment.
"""
update_experiment_template(id; aws_config::AbstractAWSConfig=global_aws_config()) = fis("PATCH", "/experimentTemplates/$(id)"; aws_config=aws_config)
update_experiment_template(id, params::AbstractDict{String}; aws_config::AbstractAWSConfig=global_aws_config()) = fis("PATCH", "/experimentTemplates/$(id)", params; aws_config=aws_config)


text:
using ..ParameterDistributions
using ..DataContainers
using ..Localizers

using Random
using Statistics
using Distributions
using LinearAlgebra
using DocStringExtensions

export EnsembleKalmanProcess
export get_u, get_g
export get_u_prior, get_u_final, get_u_mean_final, get_g_final, get_N_iterations, get_error
export compute_error!
export update_ensemble!
export sample_empirical_gaussian, split_indices_by_success
export SampleSuccGauss, IgnoreFailures, FailureHandler

abstract type Process end
#specific Processes and their exports are included after the general definitions

# Failure handlers
abstract type FailureHandlingMethod end

"Failure handling method that ignores forward model failures"
struct IgnoreFailures <: FailureHandlingMethod end

""""
    SampleSuccGauss <: FailureHandlingMethod

Failure handling method that substitutes failed ensemble members by new samples from
the empirical Gaussian distribution defined by the updated successful ensemble.
"""
struct SampleSuccGauss <: FailureHandlingMethod end

"""
    FailureHandler{P <: Process, FM <: FailureHandlingMethod}

Structure defining the failure handler method used in the EnsembleKalmanProcess.

# Fields

$(TYPEDFIELDS)

# Constructors

$(METHODLIST)
"""
struct FailureHandler{P <: Process, FM <: FailureHandlingMethod}
    "Failsafe algorithmic update equation"
    failsafe_update::Function
end

## begin general constructor and function definitions

"""
    EnsembleKalmanProcess{FT <: AbstractFloat, IT <: Int, P <: Process}

Structure that is used in Ensemble Kalman processes.

# Fields

$(TYPEDFIELDS)

# Constructors

    EnsembleKalmanProcess(
        params::AbstractMatrix{FT},
        obs_mean,
        obs_noise_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}},
        process::P;
        Δt = FT(1),
        rng::AbstractRNG = Random.GLOBAL_RNG
    ) where {FT <: AbstractFloat, P <: Process}

Inputs:

 - `params`         :: Initial parameter ensemble
 - `obs_mean`       :: Vector of observations
 - `obs_noise_cov`  :: Noise covariance associated with the observations `obs_mean`
 - `process`        :: Algorithm used to evolve the ensemble
 - `Δt`             :: Initial time step or learning rate
 - `rng`            :: Random number generator

$(METHODLIST)
"""
struct EnsembleKalmanProcess{FT <: AbstractFloat, IT <: Int, P <: Process}
    "array of stores for parameters (`u`), each of size [`N_par × N_ens`]"
    u::Array{DataContainer{FT}}
    "vector of the observed vector size [`N_obs`]"
    obs_mean::Vector{FT}
    "covariance matrix of the observational noise, of size [`N_obs × N_obs`]"
    obs_noise_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}}
    "ensemble size"
    N_ens::IT
    "Array of stores for forward model outputs, each of size  [`N_obs × N_ens`]"
    g::Array{DataContainer{FT}}
    "vector of errors"
    err::Vector{FT}
    "vector of timesteps used in each EK iteration"
    Δt::Vector{FT}
    "the particular EK process (`Inversion` or `Sampler` or `Unscented` or `SparseInversion`)"
    process::P
    "Random number generator object (algorithm + seed) used for sampling and noise, for reproducibility. Defaults to `Random.GLOBAL_RNG`."
    rng::AbstractRNG
    "struct storing failsafe update directives, implemented for (`Inversion`, `SparseInversion`, `Unscented`)"
    failure_handler::FailureHandler
    "Localization kernel, implemented for (`Inversion`, `SparseInversion`, `Unscented`)"
    localizer::Localizer
end

function EnsembleKalmanProcess(
    params::AbstractMatrix{FT},
    obs_mean,
    obs_noise_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}},
    process::P;
    Δt = FT(1),
    rng::AbstractRNG = Random.GLOBAL_RNG,
    failure_handler_method::FM = IgnoreFailures(),
    localization_method::LM = NoLocalization(),
) where {FT <: AbstractFloat, P <: Process, FM <: FailureHandlingMethod, LM <: LocalizationMethod}

    #initial parameters stored as columns
    init_params = DataContainer(params, data_are_columns = true)

    # dimensionality
    N_par, N_ens = size(init_params) #stored with data as columns
    N_obs = length(obs_mean)

    IT = typeof(N_ens)
    #store for model evaluations
    g = []
    # error store
    err = FT[]
    # timestep store
    Δt = Array([Δt])
    # failure handler
    fh = FailureHandler(process, failure_handler_method)
    # localizer
    loc = Localizer(localization_method, N_par, N_obs, FT)

    EnsembleKalmanProcess{FT, IT, P}([init_params], obs_mean, obs_noise_cov, N_ens, g, err, Δt, process, rng, fh, loc)
end

"""
    get_u(ekp::EnsembleKalmanProcess, iteration::IT; return_array=true) where {IT <: Integer}

Get for the EKI iteration. Returns a DataContainer object unless array is specified.
"""
function get_u(ekp::EnsembleKalmanProcess, iteration::IT; return_array = true) where {IT <: Integer}
    return return_array ? get_data(ekp.u[iteration]) : ekp.u[iteration]
end

"""
    get_g(ekp::EnsembleKalmanProcess, iteration::IT; return_array=true) where {IT <: Integer}

Get for the EKI iteration. Returns a DataContainer object unless array is specified.
"""
function get_g(ekp::EnsembleKalmanProcess, iteration::IT; return_array = true) where {IT <: Integer}
    return return_array ? get_data(ekp.g[iteration]) : ekp.g[iteration]
end

"""
    get_u(ekp::EnsembleKalmanProcess; return_array=true)

Get for the EKI iteration. Returns a `DataContainer` object unless array is specified.
"""
function get_u(ekp::EnsembleKalmanProcess; return_array = true) where {IT <: Integer}
    N_stored_u = get_N_iterations(ekp) + 1
    return [get_u(ekp, it, return_array = return_array) for it in 1:N_stored_u]
end

"""
    get_g(ekp::EnsembleKalmanProcess; return_array=true)

Get for the EKI iteration. Returns a `DataContainer` object unless array is specified.
"""
function get_g(ekp::EnsembleKalmanProcess; return_array = true) where {IT <: Integer}
    N_stored_g = get_N_iterations(ekp)
    return [get_g(ekp, it, return_array = return_array) for it in 1:N_stored_g]
end

"""
    get_u_final(ekp::EnsembleKalmanProcess, return_array=true)

Get the final or prior iteration of parameters or model ouputs, returns a `DataContainer` Object if `return_array` is false.
"""
function get_u_final(ekp::EnsembleKalmanProcess; return_array = true)
    return return_array ? get_u(ekp, size(ekp.u)[1]) : ekp.u[end]
end

"""
    get_u_prior(ekp::EnsembleKalmanProcess, return_array=true)

Get the final or prior iteration of parameters or model ouputs, returns a DataContainer Object if return_array is false.
"""

function get_u_prior(ekp::EnsembleKalmanProcess; return_array = true)
    return return_array ? get_u(ekp, 1) : ekp.u[1]
end

"""
    get_g_final(ekp::EnsembleKalmanProcess, return_array=true)

Get the final or prior iteration of parameters or model ouputs, returns a DataContainer Object if `return_array` is false.
"""

function get_g_final(ekp::EnsembleKalmanProcess; return_array = true)
    return return_array ? get_g(ekp, size(ekp.g)[1]) : ekp.g[end]
end

"""
    get_N_iterations(ekp::EnsembleKalmanProcess

Get number of times update has been called (equals `size(g)`, or `size(u)-1`).
"""
function get_N_iterations(ekp::EnsembleKalmanProcess)
    return size(ekp.u)[1] - 1
end

"""
    construct_initial_ensemble(
        rng::AbstractRNG,
        prior::ParameterDistribution,
        N_ens::IT;
        rng_seed::Union{IT, Nothing} = nothing,
    ) where {IT <: Int}
    construct_initial_ensemble(prior::ParameterDistribution, N_ens::IT; kwargs...) where {IT <: Int}

Construct the initial parameters, by sampling `N_ens` samples from specified
prior distribution. Returned with parameters as columns.
"""
function construct_initial_ensemble(
    rng::AbstractRNG,
    prior::ParameterDistribution,
    N_ens::IT;
    rng_seed::Union{IT, Nothing} = nothing,
) where {IT <: Int}
    # Ensuring reproducibility of the sampled parameter values:
    # re-seed the rng *only* if we're given a seed
    if rng_seed !== nothing
        rng = Random.seed!(rng, rng_seed)
    end
    return sample(rng, prior, N_ens) #of size [dim(param space) N_ens]
end
# first arg optional; defaults to GLOBAL_RNG (as in Random, StatsBase)
construct_initial_ensemble(prior::ParameterDistribution, N_ens::IT; kwargs...) where {IT <: Int} =
    construct_initial_ensemble(Random.GLOBAL_RNG, prior, N_ens; kwargs...)

"""
    compute_error!(ekp::EnsembleKalmanProcess)

Computes the covariance-weighted error of the mean forward model output, `(ḡ - y)'Γ_inv(ḡ - y)`.
The error is stored within the `EnsembleKalmanProcess`.
"""
function compute_error!(ekp::EnsembleKalmanProcess)
    mean_g = dropdims(mean(get_g_final(ekp), dims = 2), dims = 2)
    diff = ekp.obs_mean - mean_g
    X = ekp.obs_noise_cov \ diff # diff: column vector
    newerr = dot(diff, X)
    push!(ekp.err, newerr)
end

"""
    get_error(ekp::EnsembleKalmanProcess)

Returns the mean forward model output error as a function of algorithmic time.
"""
get_error(ekp::EnsembleKalmanProcess) = ekp.err

function set_Δt!(ekp::EnsembleKalmanProcess, Δt_new::T) where {T}
    if !isnothing(Δt_new)
        push!(ekp.Δt, Δt_new)
    elseif isnothing(Δt_new) && isempty(ekp.Δt)
        push!(ekp.Δt, FT(1))
    else
        push!(ekp.Δt, ekp.Δt[end])
    end
end

"""
    sample_empirical_gaussian(
        u::AbstractMatrix{FT},
        n::IT;
        inflation::Union{FT, Nothing} = nothing,
    ) where {FT <: Real, IT <: Int}

Returns `n` samples from an empirical Gaussian based on point estimates `u`, adding inflation
if the covariance is singular.
"""
function sample_empirical_gaussian(
    u::AbstractMatrix{FT},
    n::IT;
    inflation::Union{FT, Nothing} = nothing,
) where {FT <: Real, IT <: Int}
    cov_u_new = cov(u, u, dims = 2)
    if !isposdef(cov_u_new)
        @warn string("Sample covariance matrix over ensemble is singular.", "\n Appplying variance inflation.")
        if isnothing(inflation)
            # Reduce condition number to 1/sqrt(eps(FT))
            inflation = eigmax(cov_u_new) * sqrt(eps(FT))
        end
        cov_u_new = cov_u_new + inflation * I
    end
    mean_u_new = mean(u, dims = 2)
    return rand(MvNormal(mean_u_new[:], cov_u_new), n)
end

"""
     split_indices_by_success(g::AbstractMatrix{FT}) where {FT <: Real}

Returns the successful/failed particle indices given a matrix with output vectors stored as columns.
Failures are defined for particles containing at least one NaN output element.
"""
function split_indices_by_success(g::AbstractMatrix{FT}) where {FT <: Real}
    failed_ens = [i for i = 1:size(g, 2) if any(isnan.(g[:, i]))]
    successful_ens = filter(x -> !(x in failed_ens), collect(1:size(g, 2)))
    if length(failed_ens) > length(successful_ens)
        @warn string(
            "More than 50% of runs produced NaNs ($(length(failed_ens))/$(size(g, 2))).",
            "\nIterating... but consider increasing model stability.",
            "\nThis will affect optimization result.",
        )
    end
    return successful_ens, failed_ens
end

## include the different types of Processes and their exports:

# struct Inversion
export Inversion
export find_ekp_stepsize
include("EnsembleKalmanInversion.jl")

# struct SparseInversion
export SparseInversion
include("SparseEnsembleKalmanInversion.jl")

# struct Sampler
export Sampler
include("EnsembleKalmanSampler.jl")

# struct Unscented
export Unscented
export Gaussian_2d
export construct_initial_ensemble, construct_mean, construct_cov
include("UnscentedKalmanInversion.jl")


text:
using LCMCore
using CaesarLCMTypes

function typed_callback(channel::String, msg::apriltag_t)
    @show channel
    @show msg
end
lcm=LCM()
subscribe(lcm, "TEST_CHANNEL", typed_callback, apriltag_t)

while true
    handle(lcm)
end


text:
module TestGetIndex_VariableTimestep

using Mimi
using Test

my_model = Model()

#Testing that you cannot add two components of the same name
@defcomp testcomp1 begin
    var1 = Variable(index=[time])
    par1 = Parameter(index=[time])

    function run_timestep(p, v, d, t)
        v.var1[t] = p.par1[t]
    end
end

@defcomp testcomp2 begin
    var2 = Variable(index=[time])
    par2 = Parameter(index=[time])

    function run_timestep(p, v, d, t)
        v.var2[t] = p.par2[t]
    end
end

@defcomp testcomp3 begin
    var3 = Variable(index=[time])
    par3 = Parameter(index=[time])

    function run_timestep(p, v, d, t)
        v.var3[t] = p.par3[t]
    end
end

set_dimension!(my_model, :time, [2000:1:2014; 2015:5:2110])
add_comp!(my_model, testcomp1)

par = collect([2000:1:2014; 2015:5:2110])

update_param!(my_model, :testcomp1, :par1, par)
run(my_model)

# Regular get index
@test my_model[:testcomp1, :var1] == par

# Calling get index on nonexistent variable (with existing component)
@test_throws ErrorException my_model[:testcomp1, :var2]

# Calling index on component that does not exist
@test_throws ErrorException my_model[:testcomp2, :var2]

#Possibly more tests after adding another component

end #module