text:
using WordleSolver, ProgressMeter

function multiplay(word::String, scorer::Function, reps::Int)
    return [auto_play(word, scorer = scorer, verbose = false) for _ = 1:reps]
end

"""
    scorer_step_table(scorer::Function=min_max_score, reps::Int=10)
Use the `scorer` to play all possible wordle words `reps` times each
and create a dictionary of the average number of steps to solve.
"""
function scorer_step_table(scorer::Function = min_max_score, reps::Int = 5)
    d = Dict{String,Float64}()
    PM = Progress(WordleSolver.NA)
    for w ∈ WordleSolver.ANS_LIST
        s = sum(multiplay(w, scorer, reps))
        d[w] = s / reps
        next!(PM)
    end
    return d
end

function worst(d::Dict{T,S}) where {T,S<:Real}
    m = maximum(values(d))
    return [w for w in keys(d) if d[w] == m]
end

"""
    worst_words(scorer::Function=min_max_score, reps::Int=5)
Find the words that scores the worst using scoring method `scorer`.
"""
function worst_words(scorer::Function = min_max_score, reps::Int = 5)
    d = scorer_step_table(scorer, reps)
    return worst(d)
end

"""
    make_all_tables(reps::Int = 10)
Play all words `reps` times using all scoring methods. Return as a table of tables.
"""
function make_all_tables(reps::Int = 10)
    fns = [min_max_score, entropy_score, L2_score, equi_score]
    tables = Dict{Function, Dict{String, Float64}}()

    for f in fns
        @info "Playing all words $reps times using $f"
        tables[f] = scorer_step_table(f, reps)
    end

    @info "Done"
    return tables
end

function full_report(T::Dict{Function, Dict{String, Float64}})
    @info "Most difficult words by method"

    for f in keys(T)
        d = T[f]
        ww = worst(d)
        print(f,":\t")
        for w ∈ ww
            print(w," ")
        end
        println()
        println("Take about $(d[first(ww)]) steps")
        println()
    end
end

function run(reps::Int=10)
    T = make_all_tables(reps)
    full_report(T)
    return T
end


text:
"""

```
cutoffdistances!(cutoff::Float64,
                 x_solute::AbstractVector{T},
                 x_solvent::Vector{T},
                 lc_solvent::LinkedCells,
                 box::Box,
                 dc::CutoffDistances) where T
```

This routine that returns a list of the distances
between atoms that are smaller than a specified cutoff,
for a given set of coordinates.

Returns `nd`, the number of distances smaller than the cutoff, and modifies `dc`.

"""
function cutoffdistances!(cutoff::Float64,
                          x_solute::AbstractVector{T},
                          x_solvent::Vector{T},
                          lc_solvent::LinkedCells,
                          box::Box,
                          dc::CutoffDistances) where T

  # Reset the dc structure
  reset!(dc)

  # Loop over solute atoms
  for iat in eachindex(x_solute)
    xat = x_solute[iat]
    # Check the cell of this atom
    i, j, k = icell3D(xat,box)
    # Loop over vicinal cells to compute distances to solvent atoms, and
    # add data to dc structure (includes current cell)
    for ic in i-box.lcell:i+box.lcell
      for jc in j-box.lcell:j+box.lcell
        for kc in k-box.lcell:k+box.lcell
          cutoffdcell!(cutoff,iat,xat,x_solvent,lc_solvent,box,ic,jc,kc,dc)
        end
      end
    end
  end

  return nothing
end

ext: jl


text:
struct MinimalVec{V<:AbstractVector}
    vec::V
end

Base.getindex(v::MinimalVec) = v.vec # for convience, should not interfere

# minimal interface according to docs
Base.:*(a::Number, v::MinimalVec) = MinimalVec(a*v[])

Base.similar(v::MinimalVec) = MinimalVec(similar(v[]))

LinearAlgebra.axpy!(α, v::MinimalVec, w::MinimalVec) = (axpy!(α, v[], w[]); return w)
LinearAlgebra.axpby!(α, v::MinimalVec, β, w::MinimalVec) =
    (axpby!(α, v[], β, w[]); return w)
LinearAlgebra.rmul!(v::MinimalVec, α) = (rmul!(v[], α); return v)

LinearAlgebra.mul!(w::MinimalVec, v::MinimalVec, α) = (mul!(w[], v[], α); return w)
LinearAlgebra.dot(v::MinimalVec, w::MinimalVec) = dot(v[], w[])
LinearAlgebra.norm(v::MinimalVec) = norm(v[])


text:
# ------------------------------------------------------------------
# Licensed under the ISC License. See LICENSE in the project root.
# ------------------------------------------------------------------

"""
    readgeotable(args; coordnames=[:x,:y,:z], kwargs)

Read data from disk using `CSV.read`, optionally specifying
the columns `coordnames` with spatial coordinates.

The arguments `args` and keyword arguments `kwargs` are
forwarded to the `CSV.read` function, please check their
documentation for more details.

This function returns a [`GeoDataFrame`](@ref) object.
"""
readgeotable(args...; coordnames=[:x,:y,:z], kwargs...) =
  GeoDataFrame(read(args...; kwargs...), coordnames)

"""
    split(object, fraction, [normal])

Split spatial `object` into two parts where the first
part has a `fraction` of the total volume. The split
is performed along a `normal` direction. The default
direction is aligned with the first spatial dimension
of the object.
"""
function Base.split(object::AbstractSpatialObject,
                    fraction::Real, normal=nothing)
  if isnothing(normal)
    partition(object, FractionPartitioner(fraction))
  else
    partition(object, BisectFractionPartitioner(normal, fraction))
  end
end

"""
    groupby(sdata, var)

Partition spatial data `sdata` into groups of constant value
for spatial variable `var`.

### Notes

Missing values are grouped into a separate group.
"""
groupby(sdata::AbstractData, var::Symbol) =
  partition(sdata, VariablePartitioner(var))

"""
    boundbox(object)

Return the minimum axis-aligned bounding rectangle of the spatial `object`.

### Notes

Equivalent to `cover(object, RectangleCoverer())`
"""
boundbox(object::AbstractSpatialObject) = cover(object, RectangleCoverer())

"""
    sample(object, nsamples, [weights], replace=false)

Generate `nsamples` samples from spatial `object`
uniformly or using `weights`, with or without
replacement depending on `replace` option.
"""
function sample(object::AbstractSpatialObject, nsamples::Int,
                weights::AbstractVector=[]; replace=false)
  if isempty(weights)
    sample(object, UniformSampler(nsamples, replace))
  else
    sample(object, WeightedSampler(nsamples, weights, replace))
  end
end

"""
    join(sdata₁, sdata₂)

Join variables in spatial data `sdata₁` and `sdata₂`.
"""
Base.join(sdata₁::AbstractData, sdata₂::AbstractData) =
  join(sdata₁, sdata₂, VariableJoiner())

"""
    uniquecoords(sdata; aggreg=Dict())

Filter spatial data `sdata` to produce a new data
set with unique coordinates.

See [`UniqueCoordsFilter`](@ref) for more details.
"""
uniquecoords(sdata::AbstractData; aggreg=Dict()) =
  filter(sdata, UniqueCoordsFilter(aggreg))


text:
check_int(i::Int, key_path) = i
function check_int(o, key_path)
    try
        convert(Int, o)
    catch
        throw(AssertionError("$o must be an integer [$key_path]"))
    end
end

check_float(f::Float64, key_path) = f
function check_float(o, key_path)
    try
        convert(Float64, o)
    catch
        throw(AssertionError("$o must be a float [$key_path]"))
    end
end

check_boolean(b::Bool, key_path) = b
function check_boolean(o, key_path)
    try
        convert(Bool, o)
    catch
        throw(AssertionError("$o must be a boolean [$key_path]"))
    end
end

function check_array(val::AbstractArray, values::AbstractArray, key_path)
    if !issubset(val, values)
        throw(AssertionError("$val not in $values [$key_path]"))
    end
end
check_array(val, values::AbstractArray, key_path) =
    throw(AssertionError("$val must be an array (subset of $values) [$key_path]"))

function get_values(def, classes)
    if haskey(def, "values")
        return def["values"]
    elseif haskey(def, "class")
        return keys_or_vals(classes[def["class"]])
    else
        throw(AssertionError("array type definitions must specify `values` or `class`"))
    end
end

"""
Checks if an item meets the requirements of the field's definition.
"""
function check_item(val, def, key_path; keys = nothing, classes = nothing)
    # check type of value
    dtype = def["type"]
    if dtype == "int"
        val = check_int(val, key_path)
    elseif dtype == "float"
        val = check_float(val, key_path)
    elseif dtype == "boolean"
        val = check_boolean(val, key_path)
    elseif dtype == "enum"
        values = get_values(def, classes)
        @assert val in values "$val not in $values [$key_path]"
    elseif dtype == "array"
        values = get_values(def, classes)
        check_array(val, values, key_path)
    elseif dtype == "keys"
        check_array(val, keys, key_path)
    end

    # check range
    if haskey(def, "min")
        @assert val >= def["min"] "$val must be greater than $(def["min"]) [$key_path]"
    end
    if haskey(def, "max")
        @assert val <= def["max"] "$val must be less than $(def["max"]) [$key_path]"
    end

    return val
end


text:
export Point
export History, best_point, n_solved, n_params, load_history, save_history;
export f_values;

mutable struct Point{F}
    guessV :: Vector{F}
    solnV :: Vector{F}
    fStart :: F
    fVal :: F
    exitFlag :: Symbol
    runTime :: Second
    nEvals :: Int
end

empty_point(F) =
    Point{F}(zeros(F, 1), zeros(F, 1),
        zero(F), zero(F),
        :notRun, Second(0), 0);

n_params(p :: Point{F}) where F = length(p.guessV);
solved(p :: Point{F}) where F = (p.exitFlag != :notRun);
f_value(p :: Point{F}) where F = p.fVal;
fstart_value(p :: Point{F}) where F = p.fStart;

"""
	$(SIGNATURES)

Stores all points computed during the multistart.
Guesses are filled at construction. Solutions are filled in in any order.
"""
mutable struct History{F}
    fPath :: String
    # Initial population
    guessesV :: Vector{Vector{Float64}}
    pointV :: Vector{Point{F}}
end

Base.show(io :: IO, h :: History{F}) where F =
    print(io, "MultiStartLH history with $(n_solved(h)) points.");

get_points(h :: History{F}) where F = get_points(h, 1 : n_solved(h));
get_points(h :: History{F}, j) where F = h.pointV[j];
n_params(h :: History{F}) where F = length(first(h.guessesV));
n_solved(h :: History{F}) where F = sum(solved.(h.pointV));
f_values(h :: History{F}) where F = f_value.(get_points(h));
fstart_values(h :: History{F}) where F = fstart_value.(get_points(h));

function init_history(fPath :: String,
            guessesV :: AbstractVector{Vector{F1}}) where F1
    # nParams = length(first(guessesV));
    nPoints = length(guessesV);
    return History(fPath, guessesV, fill(empty_point(Float64), nPoints));
        # guessesV,
        # [Vector{F1}()  for j = 1 : nPoints],
        # [Vector{F1}()  for j = 1 : nPoints],
        # fill(1e8, nPoints),
        # fill(1e8, nPoints),
        # fill(:notRun, nPoints),
        # fill(Second(0), nPoints))
end

function add_to_history!(h :: History, j :: Integer, p :: Point{F}) where F
    @assert n_params(p) == n_params(h);
    # @assert length(startV) == n_params(h);
    # h.nPoints += 1;
    @assert !solved(h.pointV[j]);
    h.pointV[j] = p;
end

"""
	$(SIGNATURES)

Return best point of a history: parameters, fVal, exitFlag
"""
function best_point(h :: History)
    @assert n_solved(h) > 0;
    fVal, idx = findmin(f_values(h));
    return h.pointV[idx]
end

function save_history(h :: History)
    jldsave(h.fPath; h);
end

"""
	$(SIGNATURES)

Load a history. Returns a `History` object.
"""
function load_history(histFn)
    @assert isfile(histFn);
    h = load(histFn, "h");
    @assert h isa History;
    return h
end

function make_test_history(nPoints, nParams;
            fPath = test_history_path())
    lb = 1.0;
    ub = 2.0;
    rng = MersenneTwister(23);
    guessesV = [lb .+ (ub - lb) .* rand(rng, nParams)  for j = 1 : nPoints];
    h = init_history(fPath, guessesV);
    return h
end

test_history_path() = joinpath(test_dir(), "history_test.jld2");

# ----------


text:
"""
    upgrade(src, dest)
    upgrade(src, dest, project, manifest)
    upgrade(src::Vector, dest::Vector)

Handle upgrading existing JLSO files. This can be useful to:

1. Avoid repeatedly reading legacy jlso files (requiring structure updates at runtime)
2. Allow overriding project and manifest data rather than building it as part of the
upgrade step. Won't be needed once the v2 file format is dropped.
3. Allows for efficient batch updating of files created with the same environment.

!!! warning
    The manifest generated by this process are best guesses.
    JLSO v2 files do not include information on indirect dependencies, so the upgrade process relies on the package manager to work those out.
    The project generated is accurate, except for duplicated names (no UUID) and custom branchs.
"""
function upgrade(src, dest)
    info(LOGGER, "Upgrading $src -> $dest")
    jlso = open(io -> read(io, JLSOFile), src)
    write(dest, jlso)
end

function upgrade(src, dest, project, manifest)
    info(LOGGER, "Upgrading $src -> $dest")
    parsed = BSON.load(string(src))
    haskey(parsed["metadata"], "pkgs") && empty!(parsed["metadata"]["pkgs"])
    d = upgrade_jlso(parsed)

    if VersionNumber(d["metadata"]["version"]) < v"4"
        jlso = JLSOFile(
            VersionNumber(d["metadata"]["version"]),
            VersionNumber(d["metadata"]["julia"]),
            Symbol(d["metadata"]["format"]),
            Symbol(d["metadata"]["compression"]),
            d["metadata"]["image"],
            project,
            manifest,
            Dict{Symbol, Vector{UInt8}}(Symbol(k) => v for (k, v) in d["objects"]),
            ReentrantLock()
        )
    else
        # We don't need to upgrade v4 files
        jlso = read(src, JLSOFile)
    end

    write(dest, jlso)
end

function upgrade(sources::Vector, destinations::Vector)
    iter = zip(sources, destinations)

    if !isempty(iter)
        (src, dest), s = iterate(iter)
        mkpath(dirname(dest))
        jlso = open(io -> read(io, JLSOFile), src)
        write(dest, jlso)

        for (src, dest) in Iterators.rest(iter, s)
            mkpath(dirname(dest))
            upgrade(src, dest, jlso.project, jlso.manifest)
        end
    end
end

function upgrade_jlso(raw_dict::AbstractDict)
    result = copy(raw_dict)
    version = version_number(result["metadata"]["version"])

    while version < v"3"
        result = upgrade_jlso(result, Val(Int(version.major)))
        version = version_number(result["metadata"]["version"])
    end

    return result
end

# v1 and v2 stored version numbers but v3 stores strings to be bson parser agnostic
version_number(x::VersionNumber) = x
version_number(x::String) = VersionNumber(x)

# Metadata changes to upgrade file from v1 to v2
function upgrade_jlso(raw_dict::AbstractDict, ::Val{1})
    info(LOGGER, "Upgrading JLSO format from v1 to v2")
    metadata = copy(raw_dict["metadata"])
    version = version_number(metadata["version"])
    @assert version ∈ semver_spec("1")

    if metadata["format"] == :serialize
        metadata["format"] = :julia_serialize
    end

    metadata["compression"] = :none
    metadata["version"] = v"2"

    return Dict("metadata" => metadata, "objects" => get(raw_dict, "objects", Dict()))
end

# Metadata changes to upgrade from v2 to v3
function upgrade_jlso(raw_dict::AbstractDict, ::Val{2})
    info(LOGGER, "Upgrading JLSO format from v2 to v3")
    metadata = copy(raw_dict["metadata"])
    version = version_number(metadata["version"])
    @assert version ∈ semver_spec("2")

    # JLSO v3 expects to be loading a compressed manifest
    compression = get(metadata, "compression", :none)
    project, manifest = _upgrade_env(get(metadata, "pkgs", nothing))

    return Dict{String, Dict}(
        "metadata" => Dict{String, Union{String, Vector{UInt8}}}(
            "version" => "3",
            "julia" => string(get(metadata, "julia", VERSION)),
            "format" => string(get(metadata, "format", :julia_serialize)),
            "compression" => string(compression),
            "image" => get(metadata, "image", ""),
            "project" => project,
            "manifest" => read(compress(compression, IOBuffer(manifest))),
        ),
        "objects" => Dict{String, Vector{UInt8}}(
            k => v for (k, v) in get(raw_dict, "objects", Dict())
        ),
    )
end

#=
NOTE:

JLSO file v4 just moved the object data from the BSON doc to raw bytes at the
end of the file, avoiding BSON size limitations.
Extracting the object data is handled separately in `_extract_objects` in file_io.jl.
=#

# Used to convert the old "PkgName" => VersionNumber metadata to a
# Project.toml and Manifest.toml file.
function _upgrade_env(pkgs::Dict{String, VersionNumber})
    src_env = Base.active_project()

    try
        mktempdir() do tmp
            Pkg.activate(tmp)

            # In julia 1.5 a `require_not_empty` restriction was added for adding packages,
            # we use this early exit to avoid that.
            # JuliaLang/Pkg.jl@87aab4a#diff-a43f827629f8ae12bbcbb218a293d325afcdda3a7f3c7016e2292840bf29ed1cR58
            isempty(pkgs) && return _env()

            # We construct an array of PackageSpecs to avoid ordering problems with
            # adding each package individually
            try
                Pkg.add([
                    Pkg.PackageSpec(; name=key, version=value) for (key, value) in pkgs
                ])
            catch e
                # Warn about failure and fallback to simply trying to install the pacakges
                # without version constraints.
                warn(LOGGER) do
                    "Failed to construct an environment with the provide package version " *
                    "($pkgs): $e.\n Falling back to simply adding the packages."
                end
                Pkg.add([Pkg.PackageSpec(; name=key) for (key, value) in pkgs])
            end

            return _env()
        end
    finally
        Pkg.activate(src_env)
    end
end

_upgrade_env(::Nothing) = ("", "")

max_stars_repo_name: invenia/JLSO

max_issues_repo_name: invenia/JLSO

max_forks_repo_name: invenia/JLSO


text:
@testset "LinAlg" begin
    d = @data eye(3, 3)
    d[1, 1] = missing

    @test_nowarn svd(d)
end


text:
using Test, Distances
@testset "TE wrapper" begin
    d1 = rand(50)
    d2 = rand(50)
    r = rand(50)
    tol = 1e-8

    # Directly from time series
    te(d1, r, estimator = :tetogrid)
    @test te(d1, r, estimator = :tetogrid) >= 0 - tol
    @test te(d2, r, estimator = :tetogrid) >= 0 - tol
    @test te(d1, r, estimator = :tefreq) >= 0 - tol
    @test te(d2, r, estimator = :tefreq) >= 0 - tol
    @test te(d1, r, estimator = :tekNN) >= 0 - tol
    @test te(d2, r, estimator = :tekraskov) >= 0 - tol

    #########################
    # Surrogates
    #########################
    @test te(d1, r, which_is_surr = :driver) >= 0 - tol
    @test te(d1, r, which_is_surr = :response) >= 0 - tol
    @test te(d1, r, which_is_surr = :none) >= 0 - tol

    @test te(d1, r, which_is_surr = :both, surr_func = aaft) >= 0 - tol
    @test te(d1, r, which_is_surr = :both, surr_func = iaaft) >= 0 - tol
    @test te(d1, r, which_is_surr = :both, surr_func = randomphases) >= 0 - tol
    @test te(d1, r, which_is_surr = :both, surr_func = randomamplitudes) >= 0 - tol
    @test te(d1, r, which_is_surr = :both, surr_func = randomshuffle) >= 0 - tol

    #########################
    # Forward prediction lags
    #########################
    @test te(d1, r, ν = 1) >= 0 - tol
    @test te(d1, r, ν = 5) >= 0 - tol

    #########################
    # Embedding lags and dimensions
    #########################
    d1 = rand(100)
    d2 = rand(100)
    r = rand(100)
    @test te(d1, r, ν = 1, τ = 2, dim = 4) >= 0 - tol
    @test te(d1, r, ν = 3, τ = 3, dim = 5) >= 0 - tol

    #########################
    # Tuning the bins
    #########################
    @test te(d1, r, n_ϵ = 10) >= 0 - tol
    @test te(d1, r, n_ϵ = 10, max_numbins = 5, min_numbins = 2) >= 0 - tol

end
# maxlag = 2
# res1 = te_lags(d1, r, -maxlag:maxlag, method = :transferoperator_grid,
#     max_numbins = 7, min_numbins = 3)
# res2 = te_lags(d1, r, -maxlag:maxlag, method = :visitfreq,
#     max_numbins = 10, min_numbins = 3)
# res3 = te_lags(d1, r, -maxlag:maxlag, method = :transferoperator_grid)
# res4 = te_lags(d1, r, -maxlag:maxlag, method = :visitfreq)
# res5 = te_lags(d1, r, -maxlag:maxlag, method = :kraskov)
#
# @test all(res1 .>= 0)
# @test all(res2 .>= 0)
# @test all(res3 .>= 0)
# @test all(res4 .>= 0)
# @test all(res5 .>= 0)
#
# # From embedding with associated TEVars instance
# E = embed([d1, d2, r], [1, 2, 3, 3], [0, 0, 0, 1])
# v = TEVars([4], [3], [1, 2])
# res1 = te_lags(E, v, -maxlag:maxlag, method = :transferoperator_grid,
#     max_numbins = 7, min_numbins = 3)
# res2 = te_lags(E, v, -maxlag:maxlag, method = :visitfreq,
#     max_numbins = 10, min_numbins = 3)
# res3 = te_lags(E, v, -maxlag:maxlag, method = :transferoperator_grid)
# res4 = te_lags(E, v, -maxlag:maxlag, method = :visitfreq)
# res5 = te_lags(E, v, -maxlag:maxlag, method = :kraskov)
#
# @test all(res1 .>= 0)
# @test all(res2 .>= 0)
# @test all(res3 .>= 0)
# @test all(res4 .>= 0)
# @test all(res5 .>= 0)


text:
const TEST_NAME = "PCG_XSH_RS_64_32"

include("common.jl")

using RNG.PCG

r = PCGStateSetseq(UInt64, PCG_XSH_RS, (0x018cd83e277674ac, 0x436cd6f2434be066))

test_all(r, 100_000_000)


text:
"""
    F = JopNim(spc, c)

where `F` is the 'normalized integral method' operator F = int_0^t [ x(t) dt ] / int_0_T [ x(t) dt ]
"""
JopNim(spc::JetSpace{T}) where {T} = JopNl(dom = spc, rng = spc, f! = JopNim_f!, df! = JopNim_df!, df′! = JopNim_df′!)
export JopNim

function JopNim_f!(d::AbstractArray, m::AbstractArray)
    d .= cumsum(m,dims=1) ./ sum(m,dims=1)
end

# (-I_t[dm] I[m] + I_t[m] I[dm]) / (I_t[m])^2
function JopNim_df!(δd::AbstractArray, δm::AbstractArray; mₒ)
    δd .= (-sum(δm,dims=1) .* cumsum(mₒ,dims=1).+ sum(mₒ,dims=1) .* cumsum(δm,dims=1)) ./ sum(mₒ,dims=1).^2
end

# (-I_tᵀ[mᵀIᵀ] + I_t[m] Iᵀ) δd / (I_t[m])^2
function JopNim_df′!(δm::AbstractArray{T}, δd::AbstractArray{T}; mₒ) where {T}
    A = JopLeakyIntegration(JetSpace(T,size(δm)), α=1.0)
    δm .= conj.(( -sum( mₒ .* (A'*conj.(δd)) ,dims=1) .+ ( sum(mₒ,dims=1) .* (A'*conj.(δd)) ) ) ./ sum(mₒ,dims=1).^2)
end


text:
module DeepCorrections

using LinearAlgebra
using POMDPs
using POMDPModelTools
using POMDPPolicies
using Flux
using RLInterface
using DeepQLearning
using DeepQLearning: isrecurrent, create_dueling_network, sample, globalnorm
using DeepQLearning: initialize_replay_buffer, update_priorities!
using DeepQLearning: batch_train!, dqn_train!
using Parameters

export
    DeepCorrectionSolver,
    DeepCorrectionPolicy,
    CorrectionTrainGraph,
    lowfi_values,
    zeros_values,
    correction,
    additive_correction,
    multiplicative_correction

"""
Deep Correction algorithm to train an additive correction term to an existing value function

# Fields
    - dqn::DeepQLearningSolver
See the documentation in DeepQLearning.jl for the field of the solver, the qnetwork corresponds to the correction network
    - lowfi_values::Any
Function or object for the low fidelity value function to correct, returns a vector of action values, ordered according to action_index
If this is a function `f`, `f(mdp, s)` will be called to estimate the value.
If this is an object `o`, `lowfi_values(o, mdp, s, depth)` will be called.
default: zero_values::Function
    - correction::Any Function or object for the correction method used
If this is a function `f`, `f(mdp, q_lo, q_corr, weight)` will be called to estimate the value.
If this is an object `o`, `correction(o, mdp, q_lo, q_corr, weight)` will be called.
default: additive_correction::Function
    - correction_weight::Float64 a constant weight that provides more flexibility in tuning the correction method
"""
@with_kw mutable struct DeepCorrectionSolver <: Solver
    dqn::DeepQLearningSolver = DeepQLearningSolver()
    lowfi_values::Any = zeros_values
    correction::Any = additive_correction
    correction_weight::Float64 = 1.0
end

"""
    lowfi_values
return a low fidelity estimate of the action values, the returned value should be a vector of size `n_actions(mdp)``
the actions are ordered according to `action_index`
"""
function lowfi_values end
lowfi_values(f::Function, problem::Union{POMDP,MDP}, s) = f(problem, s)
lowfi_values(p::Policy, problem::Union{POMDP, MDP}, s) = actionvalues(p, s)

# default implementation
# should return a vector of size n_actions
function zeros_values(problem, s)
    na = n_actions(problem)
    return zeros(na)
end

"""
    correction
combine the output from the corrective network and the low fidelity values
"""
function correction end
correction(f::Function, problem::Union{POMDP, MDP}, q_lo::AbstractArray, q_corr::AbstractArray, weight::Float64) = f(problem, q_lo, q_corr, weight)

# default implementations

function additive_correction(problem::Union{POMDP, MDP}, q_lo::AbstractArray, q_corr::AbstractArray, weight::Float64)
    return q_lo .+ weight.*q_corr
end

function multiplicative_correction(problem::Union{POMDP, MDP}, q_lo::AbstractArray, q_corr::AbstractArray, weight::Float64)
    return q_lo.*q_corr
end

include("policy.jl")
include("solver.jl")

end # module


text:
"""
    accumulated_edges(adj[, num_V])

Return a vector which acts as a mapping table. The index is the vertex index,
value is accumulated numbers of edge (current vertex not included).
"""
function accumulated_edges(adj::AbstractVector{<:AbstractVector{<:Integer}},
                           num_V=size(adj,1))

    return [0, cumsum(map(length, adj))...]
end

Zygote.@nograd accumulated_edges

Zygote.@nograd function generate_cluster(M::AbstractArray{T,N}, accu_edge, V, E) where {T,N}
    cluster = similar(M, Int, E)
    @inbounds for i = 1:V
        j = accu_edge[i]
        k = accu_edge[i+1]
        cluster[j+1:k] .= i
    end
    cluster
end

"""
    vertex_pair_table(adj[, num_E])

Generate a mapping from edge index to vertex pair (i, j). The edge indecies are determined by
the sorted vertex indecies.
"""
function vertex_pair_table(adj::AbstractVector{<:AbstractVector{<:Integer}},
                           num_E=sum(map(length, adj)))
    table = similar(adj[1], Tuple{UInt32,UInt32}, num_E)
    e = one(UInt64)
    for (i, js) = enumerate(adj)
        js = sort(js)
        for j = js
            table[e] = (i, j)
            e += one(UInt64)
        end
    end
    table
end

function vertex_pair_table(eidx::Dict)
    table = Array{Tuple{UInt32,UInt32}}(undef, num_E)
    for (k, v) = eidx
        table[v] = k
    end
    table
end

Zygote.@nograd vertex_pair_table

"""
    edge_index_table(adj[, num_E])

Generate a mapping from vertex pair (i, j) to edge index. The edge indecies are determined by
the sorted vertex indecies.
"""
function edge_index_table(adj::AbstractVector{<:AbstractVector{<:Integer}},
                          num_E=sum(map(length, adj)))
    table = Dict{Tuple{UInt32,UInt32},UInt64}()
    e = one(UInt64)
    for (i, js) = enumerate(adj)
        js = sort(js)
        for j = js
            table[(i, j)] = e
            e += one(UInt64)
        end
    end
    table
end

function edge_index_table(vpair::AbstractVector{<:Tuple})
    table = Dict{Tuple{UInt32,UInt32},UInt64}()
    for (i, p) = enumerate(vpair)
        table[p] = i
    end
    table
end

Zygote.@nograd edge_index_table

function transform(X::AbstractArray, vpair::AbstractVector{<:Tuple}, num_V)
    dims = size(X)[1:end-1]..., num_V, num_V
    Y = similar(X, dims)
    for (i, p) in enumerate(vpair)
        view(Y, :, p[1], p[2]) .= view(X, :, i)
    end
    Y
end

function transform(X::AbstractArray, eidx::Dict)
    dims = size(X)[1:end-2]..., length(eidx)
    Y = similar(X, dims)
    for (k, v) in eidx
        view(Y, :, v) .= view(X, :, k[1], k[2])
    end
    Y
end


text:
#=
naive linear system solver
solves 6x6 system without reductions
=#

abstract type NaiveSystemSolver{T <: Real} <: SystemSolver{T} end

#=
direct sparse
=#

mutable struct NaiveSparseSystemSolver{T <: Real} <: NaiveSystemSolver{T}
    lhs::SparseMatrixCSC{T}
    fact_cache::SparseNonSymCache{T}
    hess_idxs::Vector
    mtt_idx::Int
    function NaiveSparseSystemSolver{T}(; fact_cache::SparseNonSymCache{T} =
        SparseNonSymCache{T}()) where {T <: Real}
        s = new{T}()
        s.fact_cache = fact_cache
        return s
    end
end

function load(
    syssolver::NaiveSparseSystemSolver{T},
    solver::Solver{T},
    ) where {T <: Real}
    syssolver.fact_cache.analyzed = false
    model = solver.model
    (n, p, q) = (model.n, model.p, model.q)
    tau_row = n + p + q + 1
    cones = model.cones
    cone_idxs = model.cone_idxs

    # form sparse LHS without Hessians in s row
    # TODO check for inefficiency, may need to implement more manually
    spz(a, b) = spzeros(T, a, b)
    spi(a, b) = sparse(one(T) * I, a, b)
    lhs = hvcat((5, 4, 4, 5, 3, 4),
        spz(n, n), model.A', model.G', model.c, spz(n, q + 1),
        -model.A, spz(p, p + q), model.b, spz(p, q + 1),
        -model.G, spz(q, p + q), model.h, -spi(q, q + 1),
        -model.c', -model.b', -model.h', spz(1, 1 + q), -1,
        spz(q, n + p), spi(q, q + 1), spi(q, q + 1),
        spz(1, n + p + q), 1, spz(1, q), 1)
    @assert lhs isa SparseMatrixCSC{T}
    dropzeros!(lhs)
    (Is, Js, Vs) = findnz(lhs)

    # add I, J, V for Hessians
    hess_nz_total = (isempty(cones) ? 0 : sum(
        Cones.hess_nz_count(cone_k) for cone_k in cones))
    H_Is = Vector{Int}(undef, hess_nz_total)
    H_Js = Vector{Int}(undef, hess_nz_total)
    offset = 1
    for (cone_k, idxs_k) in zip(cones, cone_idxs)
        z_start_k = n + p + first(idxs_k) - 1
        s_start_k = tau_row + first(idxs_k) - 1
        H_start_k = Cones.use_dual_barrier(cone_k) ? z_start_k : s_start_k
        for j in 1:Cones.dimension(cone_k)
            nz_rows_kj = s_start_k .+ Cones.hess_nz_idxs_col(cone_k, j)
            len_kj = length(nz_rows_kj)
            IJV_idxs = offset:(offset + len_kj - 1)
            offset += len_kj
            @. H_Is[IJV_idxs] = nz_rows_kj
            @. H_Js[IJV_idxs] = H_start_k + j
        end
    end
    append!(Is, H_Is)
    append!(Js, H_Js)
    append!(Vs, ones(T, hess_nz_total))

    # convert integer types here rather than inside external wrappers
    dim = size(lhs, 1)
    # integer type supported by the sparse system solver library to be used
    Ti = int_type(syssolver.fact_cache)
    Is = convert(Vector{Ti}, Is)
    Js = convert(Vector{Ti}, Js)
    lhs = syssolver.lhs = sparse(Is, Js, Vs, dim, dim)

    # cache indices of nonzeros of Hessians in sparse LHS nonzeros vector
    syssolver.hess_idxs = [Vector{Union{UnitRange, Vector{Int}}}(
        undef, Cones.dimension(cone_k)) for cone_k in cones]
    for (k, cone_k) in enumerate(cones)
        idxs_k = cone_idxs[k]
        z_start_k = n + p + first(idxs_k) - 1
        s_start_k = tau_row + first(idxs_k) - 1
        H_start_k = Cones.use_dual_barrier(cone_k) ? z_start_k : s_start_k
        for j in 1:Cones.dimension(cone_k)
            col = H_start_k + j
            # get nonzero rows in the current column of the LHS
            col_idx_start = lhs.colptr[col]
            nz_rows = lhs.rowval[col_idx_start:(lhs.colptr[col + 1] - 1)]
            # get nonzero rows in column j of the Hessian
            nz_hess_indices = Cones.hess_nz_idxs_col(cone_k, j)
            # get index corresponding to first nonzero element of the current col
            first_H = findfirst(isequal(
                s_start_k + first(nz_hess_indices)), nz_rows)
            # indices of nonzero values for cone k column j
            syssolver.hess_idxs[k][j] = (col_idx_start + first_H - 2) .+
                (1:length(nz_hess_indices))
        end
    end

    # get mu/tau/tau index
    syssolver.mtt_idx = lhs.colptr[tau_row + 1] - 1

    return syssolver
end

function update_lhs(syssolver::NaiveSparseSystemSolver, solver::Solver)
    for (k, cone_k) in enumerate(solver.model.cones)
        H_k = Cones.hess(cone_k)
        for j in 1:Cones.dimension(cone_k)
            nz_rows = Cones.hess_nz_idxs_col(cone_k, j)
            @views copyto!(syssolver.lhs.nzval[
                syssolver.hess_idxs[k][j]], H_k[nz_rows, j])
        end
    end
    tau = solver.point.tau[]
    syssolver.lhs.nzval[syssolver.mtt_idx] = solver.mu / tau / tau

    solver.time_upfact += @elapsed update_fact(syssolver.fact_cache,
        syssolver.lhs)

    return syssolver
end

function solve_system(
    syssolver::NaiveSparseSystemSolver,
    solver::Solver{T},
    sol::Point{T},
    rhs::Point{T},
    ) where {T <: Real}
    inv_prod(syssolver.fact_cache, sol.vec, syssolver.lhs, rhs.vec)
    return sol
end

#=
direct dense
=#

mutable struct NaiveDenseSystemSolver{T <: Real} <: NaiveSystemSolver{T}
    tau_row::Int
    lhs::Matrix{T}
    lhs_fact::Matrix{T}
    fact::LU{T, Matrix{T}}
    lhs_H_k::Vector
    function NaiveDenseSystemSolver{T}() where {T <: Real}
        syssolver = new{T}()
        return syssolver
    end
end

function load(
    syssolver::NaiveDenseSystemSolver{T},
    solver::Solver{T},
    ) where {T <: Real}
    model = solver.model
    (n, p, q) = (model.n, model.p, model.q)
    cones = model.cones
    cone_idxs = model.cone_idxs
    syssolver.tau_row = n + p + q + 1

    # TODO check for inefficiency, may need to implement more manually
    dz(a, b) = zeros(T, a, b)
    di(a, b) = Matrix(one(T) * I, a, b)
    lhs = hvcat((5, 4, 4, 5, 3, 4),
        dz(n, n), model.A', model.G', model.c, dz(n, q + 1),
        -model.A, dz(p, p + q), model.b, dz(p, q + 1),
        -model.G, dz(q, p + q), model.h, -di(q, q + 1),
        -model.c', -model.b', -model.h', dz(1, 1 + q), -1,
        dz(q, n + p), di(q, q + 1), di(q, q + 1),
        dz(1, n + p + q), 1, dz(1, q), 1)
    @assert lhs isa Matrix{T}
    syssolver.lhs = lhs
    syssolver.lhs_fact = zero(lhs)

    function view_H_k(cone_k, idxs_k)
        rows = syssolver.tau_row .+ idxs_k
        cols = Cones.use_dual_barrier(cone_k) ? (n + p) .+ idxs_k : rows
        return view(syssolver.lhs, rows, cols)
    end
    syssolver.lhs_H_k = [view_H_k(cone_k, idxs_k) for
        (cone_k, idxs_k) in zip(cones, cone_idxs)]

    return syssolver
end

function update_lhs(syssolver::NaiveDenseSystemSolver, solver::Solver)
    for (cone_k, lhs_H_k) in zip(solver.model.cones, syssolver.lhs_H_k)
        copyto!(lhs_H_k, Cones.hess(cone_k))
    end
    tau = solver.point.tau[]
    syssolver.lhs[end, syssolver.tau_row] = solver.mu / tau / tau

    solver.time_upfact += @elapsed syssolver.fact =
        nonsymm_fact_copy!(syssolver.lhs_fact, syssolver.lhs)

    if !issuccess(syssolver.fact)
        println("nonsymmetric linear system factorization failed")
    end

    return syssolver
end

function solve_system(
    syssolver::NaiveDenseSystemSolver,
    solver::Solver,
    sol::Point{T},
    rhs::Point{T},
    ) where {T <: Real}
    ldiv!(sol.vec, syssolver.fact, rhs.vec)
    return sol
end


text:
"""
    DerivativeCoefficients

The coefficients of a derivative operator on a nonperiodic grid.
"""
struct DerivativeCoefficients{T,LeftBoundary,RightBoundary,LeftBoundaryDerivatives,RightBoundaryDerivatives,
                              LowerOffset,UpperOffset,LeftWidth,RightWidth,Parallel,SourceOfCoefficients} <: AbstractDerivativeCoefficients{T}
    # coefficients defining the operator and its action
    left_boundary::LeftBoundary
    right_boundary::RightBoundary
    left_boundary_derivatives::LeftBoundaryDerivatives
    right_boundary_derivatives::RightBoundaryDerivatives
    lower_coef::SVector{LowerOffset, T}
    central_coef::T
    upper_coef::SVector{UpperOffset, T}
    left_weights::SVector{LeftWidth, T}
    right_weights::SVector{RightWidth, T}
    parallel::Parallel
    # corresponding orders etc.
    derivative_order::Int
    accuracy_order  ::Int
    symmetric       ::Bool
    source_of_coeffcients::SourceOfCoefficients

    function DerivativeCoefficients(left_boundary::LeftBoundary, right_boundary::RightBoundary,
                                    left_boundary_derivatives::LeftBoundaryDerivatives,
                                    right_boundary_derivatives::RightBoundaryDerivatives,
                                    lower_coef::SVector{LowerOffset, T}, central_coef::T, upper_coef::SVector{UpperOffset, T},
                                    left_weights::SVector{LeftWidth, T}, right_weights::SVector{RightWidth, T},
                                    parallel::Parallel, derivative_order::Int, accuracy_order::Int,
                                    source_of_coeffcients::SourceOfCoefficients) where {T,LeftBoundary,RightBoundary,LeftBoundaryDerivatives,RightBoundaryDerivatives,
                                                                                LowerOffset,UpperOffset,LeftWidth,RightWidth,Parallel,SourceOfCoefficients}
        @argcheck length(left_boundary) == LeftWidth
        @argcheck length(right_boundary) == RightWidth
        @argcheck length(left_boundary_derivatives) == length(right_boundary_derivatives)
        symmetric = LowerOffset == UpperOffset
        if symmetric
            @inbounds for i in Base.OneTo(LowerOffset)
                symmetric = symmetric && lower_coef[i] == upper_coef[i]
            end
        end
        symmetric = symmetric && length(left_boundary[end]) == length(left_boundary)
        #TODO: check boundary coefficients
        if derivative_order-1 != length(left_boundary_derivatives)
            warn("Derivative coefficients of degree $derivative_order should provide $(derivative_order-1) boundary derivatives.")
        end
        new{T,LeftBoundary,RightBoundary,LeftBoundaryDerivatives,RightBoundaryDerivatives,LowerOffset,UpperOffset,LeftWidth,RightWidth,Parallel,SourceOfCoefficients}(
            left_boundary, right_boundary, left_boundary_derivatives, right_boundary_derivatives,
            lower_coef, central_coef, upper_coef, left_weights, right_weights,
            parallel, derivative_order, accuracy_order, symmetric, source_of_coeffcients)
    end
end

@inline source_of_coeffcients(coefficients::AbstractDerivativeCoefficients) = coefficients.source_of_coeffcients

function Base.show(io::IO, coefficients::DerivativeCoefficients)
    if derivative_order(coefficients) == 1
        print(io, "Coefficients of the SBP 1st derivative operator of order ",
                    accuracy_order(coefficients), " given in \n")
    elseif  derivative_order(coefficients) == 2
        print(io, "Coefficients of the SBP 2nd derivative operator of order ",
                    accuracy_order(coefficients), " given in \n")
    elseif  derivative_order(coefficients) == 3
        print(io, "Coefficients of the SBP 3rd derivative operator of order ",
                    accuracy_order(coefficients), " given in \n")
    else
        print(io, "Coefficients of the SBP ", derivative_order(coefficients),
                    "th derivative operator of order ", accuracy_order(coefficients),
                    " given in \n")
    end
    print(io, source_of_coeffcients(coefficients))
end

"""
    mul!(dest::AbstractVector, coefficients::DerivativeCoefficients, u::AbstractVector, α, β)

Compute `α*D*u + β*dest` and store the result in `dest`.
"""
function mul!(dest::AbstractVector, coefficients::DerivativeCoefficients, u::AbstractVector, α, β)
    @unpack left_boundary, right_boundary, lower_coef, central_coef, upper_coef, parallel = coefficients

    @boundscheck begin
        @argcheck length(dest) == length(u) DimensionMismatch
        @argcheck length(u) > length(lower_coef)+length(upper_coef) DimensionMismatch
        @argcheck length(u) > length(left_boundary) + length(right_boundary) DimensionMismatch
    end

    convolve_boundary_coefficients!(dest, left_boundary, right_boundary, u, α, β)
    convolve_interior_coefficients!(dest, lower_coef, central_coef, upper_coef, u, α, β, length(left_boundary), length(right_boundary), parallel)
end

"""
    mul!(dest::AbstractVector, coefficients::DerivativeCoefficients, u::AbstractVector, α)

Compute `α*D*u` and store the result in `dest`.
"""
function mul!(dest::AbstractVector, coefficients::DerivativeCoefficients, u::AbstractVector, α)
    @unpack left_boundary, right_boundary, lower_coef, central_coef, upper_coef, parallel = coefficients

    @boundscheck begin
        @argcheck length(dest) == length(u) DimensionMismatch
        @argcheck length(u) > length(lower_coef)+length(upper_coef) DimensionMismatch
        @argcheck length(u) > length(coefficients.left_boundary) + length(coefficients.right_boundary) DimensionMismatch
    end

    convolve_boundary_coefficients!(dest, left_boundary, right_boundary, u, α)
    convolve_interior_coefficients!(dest, lower_coef, central_coef, upper_coef, u, α, length(left_boundary), length(right_boundary), parallel)
end

"""
    DerivativeCoefficientRow{T,Start,Length}

A struct representing a row in the boundary block of an SBP derivative operator
with scalar type `T`.
"""
struct DerivativeCoefficientRow{T,Start,Length}
    coef::SVector{Length, T}
end

function Base.length(::DerivativeCoefficientRow{T,Start,Length}) where {T,Start,Length}
    Length
end

function offset(::DerivativeCoefficientRow{T,Start,Length}) where {T,Start,Length}
    Start
end

function -(coef_row::DerivativeCoefficientRow{T,Start,Length}) where {T,Start,Length}
    DerivativeCoefficientRow{T,Start,Length}(-coef_row.coef)
end

@inline function convolve_left_row(coef_row::DerivativeCoefficientRow{T,Start,Length}, u) where {T,Start,Length}
    @inbounds tmp = coef_row.coef[1]*u[Start]
    @inbounds for i in 2:Length
        tmp += coef_row.coef[i]*u[Start+i-1]
    end
    tmp
end

@inline function convolve_right_row(coef_row::DerivativeCoefficientRow{T,Start,Length}, u) where {T,Start,Length}
    @inbounds tmp = coef_row.coef[1]*u[end-Start+1]
    @inbounds for i in 2:Length
        tmp += coef_row.coef[i]*u[end-(Start+i-2)]
    end
    tmp
end

@unroll function convolve_boundary_coefficients!(dest::AbstractVector, left_boundary, right_boundary, u::AbstractVector, α, β)
    @unroll for i in 1:length(left_boundary)
        tmp = convolve_left_row(left_boundary[i], u)
        @inbounds dest[i] = β*dest[i] + α*tmp
    end
    @unroll for i in 1:length(right_boundary)
        tmp = convolve_right_row(right_boundary[i], u)
        @inbounds dest[end-i+1] = β*dest[end-i+1] + α*tmp
    end
end

@unroll function convolve_boundary_coefficients!(dest::AbstractVector, left_boundary, right_boundary, u::AbstractVector, α)
    @unroll for i in 1:length(left_boundary)
        tmp = convolve_left_row(left_boundary[i], u)
        @inbounds dest[i] = α*tmp
    end
    @unroll for i in 1:length(right_boundary)
        tmp = convolve_right_row(right_boundary[i], u)
        @inbounds dest[end-i+1] = α*tmp
    end
end

@generated function convolve_interior_coefficients!(dest::AbstractVector, lower_coef::SVector{LowerOffset}, central_coef, upper_coef::SVector{UpperOffset}, u::AbstractVector, α, β, left_boundary_width, right_boundary_width, parallel) where {LowerOffset, UpperOffset}
    ex = :( lower_coef[$LowerOffset]*u[i-$LowerOffset] )
    for j in LowerOffset-1:-1:1
        ex = :( $ex + lower_coef[$j]*u[i-$j] )
    end
    ex = :( $ex + central_coef*u[i] )
    for j in 1:UpperOffset
        ex = :( $ex + upper_coef[$j]*u[i+$j] )
    end

    quote
        @inbounds for i in (left_boundary_width+1):(length(dest)-right_boundary_width)
            dest[i] = β*dest[i] + α*$ex
        end
    end
end

function convolve_interior_coefficients!(dest::AbstractVector{T}, lower_coef::SVector{LowerOffset}, central_coef, upper_coef::SVector{UpperOffset}, u::AbstractVector, α, β, left_boundary_width, right_boundary_width, parallel::Val{:threads}) where {T, LowerOffset, UpperOffset}
    Threads.@threads for i in (left_boundary_width+1):(length(dest)-right_boundary_width) @inbounds begin
        tmp = zero(T)
        for j in Base.OneTo(LowerOffset)
            tmp += lower_coef[j]*u[i-j]
        end
        tmp += central_coef*u[i]
        for j in Base.OneTo(UpperOffset)
            tmp += upper_coef[j]*u[i+j]
        end
        dest[i] = β*dest[i] + α*tmp
    end end
end

@generated function convolve_interior_coefficients!(dest::AbstractVector, lower_coef::SVector{LowerOffset}, central_coef, upper_coef::SVector{UpperOffset}, u::AbstractVector, α, left_boundary_width, right_boundary_width, parallel) where {LowerOffset, UpperOffset}
    ex = :( lower_coef[$LowerOffset]*u[i-$LowerOffset] )
    for j in LowerOffset-1:-1:1
        ex = :( $ex + lower_coef[$j]*u[i-$j] )
    end
    ex = :( $ex + central_coef*u[i] )
    for j in 1:UpperOffset
        ex = :( $ex + upper_coef[$j]*u[i+$j] )
    end

    quote
        @inbounds for i in (left_boundary_width+1):(length(dest)-right_boundary_width)
            dest[i] = α*$ex
        end
    end
end

function convolve_interior_coefficients!(dest::AbstractVector{T}, lower_coef::SVector{LowerOffset}, central_coef, upper_coef::SVector{UpperOffset}, u::AbstractVector, α, left_boundary_width, right_boundary_width, parallel::Val{:threads}) where {T, LowerOffset, UpperOffset}
    Threads.@threads for i in (left_boundary_width+1):(length(dest)-right_boundary_width) @inbounds begin
        tmp = zero(T)
        for j in Base.OneTo(LowerOffset)
            tmp += lower_coef[j]*u[i-j]
        end
        tmp += central_coef*u[i]
        for j in Base.OneTo(UpperOffset)
            tmp += upper_coef[j]*u[i+j]
        end
        dest[i] = α*tmp
    end end
end

"""
    DerivativeOperator

A derivative operator on a nonperiodic finite difference grid.
"""
struct DerivativeOperator{T,LeftBoundary,RightBoundary,LeftBoundaryDerivatives,RightBoundaryDerivatives,
                          LowerOffset,UpperOffset,LeftWidth,RightWidth,Parallel,SourceOfCoefficients,Grid} <: AbstractDerivativeOperator{T}
    coefficients::DerivativeCoefficients{T,LeftBoundary,RightBoundary,LeftBoundaryDerivatives,RightBoundaryDerivatives,
                                         LowerOffset,UpperOffset,LeftWidth,RightWidth,Parallel,SourceOfCoefficients}
    grid::Grid
    Δx::T
    factor::T

    function DerivativeOperator(coefficients::DerivativeCoefficients{T,LeftBoundary,RightBoundary,LeftBoundaryDerivatives,RightBoundaryDerivatives,
                                                                     LowerOffset,UpperOffset,LeftWidth,RightWidth,Parallel,SourceOfCoefficients},
                                grid::Grid) where {T,LeftBoundary,RightBoundary,LeftBoundaryDerivatives,RightBoundaryDerivatives,
                                                   LowerOffset,UpperOffset,LeftWidth,RightWidth,Parallel,SourceOfCoefficients,Grid}
        @argcheck length(grid) > LowerOffset+UpperOffset DimensionMismatch
        @argcheck length(grid) > length(coefficients.left_boundary) + length(coefficients.right_boundary) DimensionMismatch

        Δx = step(grid)
        factor = inv(Δx^derivative_order(coefficients))
        new{T,LeftBoundary,RightBoundary,LeftBoundaryDerivatives,RightBoundaryDerivatives,LowerOffset,UpperOffset,LeftWidth,RightWidth,Parallel,SourceOfCoefficients,Grid}(
            coefficients, grid, Δx, factor)
    end
end

@inline source_of_coeffcients(D::DerivativeOperator) = source_of_coeffcients(D.coefficients)

function Base.show(io::IO, D::DerivativeOperator{T}) where {T}
    if derivative_order(D) == 1
        print(io, "SBP 1st derivative operator of order ")
    elseif  derivative_order(D) == 2
        print(io, "SBP 2nd derivative operator of order ")
    elseif  derivative_order(D) == 3
        print(io, "SBP 3rd derivative operator of order ")
    else
        print(io, "SBP ", derivative_order(D), "th derivative operator of order ")
    end
    print(io, accuracy_order(D), " {T=", T, ", Parallel=", typeof(D.coefficients.parallel), "} \n")
    print(io, "on a grid in [", first(grid(D)), ", ", last(grid(D)),
                "] using ", length(grid(D)), " nodes \n")
    print(io, "and coefficients given in \n")
    print(io, source_of_coeffcients(D))
end

"""
    derivative_left(D::DerivativeOperator, u, der_order::Val{N})

Compute the `N`-th derivative of the function given by the coefficients `u` at
the left boundary of the grid.
"""
@inline function derivative_left(D::DerivativeOperator, u, der_order::Val{N}) where {N}
    convolve_left_row(D.coefficients.left_boundary_derivatives[N], u) / D.Δx^N
end

"""
    derivative_right(D::DerivativeOperator, u, der_order::Val{N})

Compute the `N`-th derivative of the function given by the coefficients `u` at
the right boundary of the grid.
"""
@inline function derivative_right(D::DerivativeOperator, u, der_order::Val{N}) where {N}
    convolve_right_row(D.coefficients.right_boundary_derivatives[N], u) / D.Δx^N
end

"""
  add_transpose_derivative_left!(u, D::DerivativeOperator, der_order::Val{N})

Compute the transposed `N`-th derivative to the function given by the coefficients `u` at
the left boundary of the grid.
"""
@inline function add_transpose_derivative_left!(u::AbstractVector, D::DerivativeOperator, der_order::Val{N}, α) where {N}
    factor = α / D.Δx^N
    coef = D.coefficients.left_boundary_derivatives[N].coef
    for i in 1:length(coef)
        u[i] += factor * coef[i]
    end
end

"""
  add_transpose_derivative_right!(u, D::DerivativeOperator, der_order::Val{N})

Compute the transposed `N`-th derivative to the function given by the coefficients `u` at
the right boundary of the grid.
"""
@inline function add_transpose_derivative_right!(u::AbstractVector, D::DerivativeOperator, der_order::Val{N}, α) where {N}
    factor = α / D.Δx^N
    coef = D.coefficients.right_boundary_derivatives[N].coef
    for i in 1:length(coef)
        u[end-i+1] += factor * coef[i]
    end
end

"""
    mul!(dest::AbstractVector, D::DerivativeOperator, u::AbstractVector, α, β)

Compute `α*D*u + β*dest` and store the result in `dest`.
"""
Base.@propagate_inbounds function mul!(dest::AbstractVector, D::DerivativeOperator, u::AbstractVector, α, β)
    @boundscheck begin
        @argcheck size(D, 2) == length(u) DimensionMismatch
        @argcheck size(D, 1) == length(dest) DimensionMismatch
    end
    @inbounds mul!(dest, D.coefficients, u, α*D.factor, β)
end

"""
    mul!(dest::AbstractVector, D::DerivativeOperator, u::AbstractVector, α)

Compute `α*D*u` and store the result in `dest`.
"""
Base.@propagate_inbounds function mul!(dest::AbstractVector, D::DerivativeOperator, u::AbstractVector, α)
    @boundscheck begin
        @argcheck size(D, 2) == length(u) DimensionMismatch
        @argcheck size(D, 1) == length(dest) DimensionMismatch
    end
    @inbounds mul!(dest, D.coefficients, u, α*D.factor)
end

"""
    integrate(func, u, D::DerivativeOperator)

Map the function `func` to the coefficients `u` and integrate with respect to
the quadrature rule associated with the SBP derivative operator `D`.
"""
function integrate(func, u::AbstractVector, D::DerivativeOperator)
    @boundscheck begin
        length(u) == length(grid(D))
    end
    @unpack Δx = D
    @unpack left_weights, right_weights = D.coefficients

    @inbounds res = sum(func, view(u,1+length(left_weights):length(u)-length(right_weights)))
    @inbounds for i in Base.OneTo(length(left_weights))
        res += left_weights[i]*func(u[i])
    end
    @inbounds for i in Base.OneTo(length(right_weights))
        res += right_weights[i]*func(u[end-i+1])
    end

    Δx * res
end

"""
    derivative_operator(source_of_coefficients, derivative_order, accuracy_order,
                        xmin, xmax, N, parallel=Val{:serial}())

Create a `DerivativeOperator` approximating the `derivative_order`-th derivative
on a grid between `xmin` and `xmax` with `N` grid points up to order of accuracy
`accuracy_order`. with coefficients given by `source_of_coefficients`.
The evaluation of the derivative can be parallised using threads by chosing
`parallel=Val{:threads}())`.
"""
function derivative_operator(source_of_coefficients, derivative_order, accuracy_order, xmin, xmax, N, parallel=Val{:serial}())
    grid = construct_grid(source_of_coefficients, accuracy_order, xmin, xmax, N)
    if derivative_order == 1
        coefficients = first_derivative_coefficients(source_of_coefficients, accuracy_order, eltype(grid), parallel)
    elseif derivative_order == 2
        coefficients = second_derivative_coefficients(source_of_coefficients, accuracy_order, eltype(grid), parallel)
    elseif derivative_order == 3
        coefficients = third_derivative_coefficients(source_of_coefficients, accuracy_order, eltype(grid), parallel)
    elseif derivative_order == 4
        coefficients = fourth_derivative_coefficients(source_of_coefficients, accuracy_order, eltype(grid), parallel)
    else
        throw(ArgumentError("Derivative order $derivative_order not implemented."))
    end
    DerivativeOperator(coefficients, grid)
end

@inline construct_grid(source_of_coefficients, accuracy_order, xmin, xmax, N) = range(xmin, stop=xmax, length=N)

function lower_bandwidth(D::DerivativeOperator)
    @unpack left_boundary, right_boundary, lower_coef = D.coefficients

    l = length(lower_coef)
    for (i,coef_row) in enumerate(left_boundary)
        l = max(l, i-offset(coef_row))
    end
    for (i,coef_row) in enumerate(right_boundary)
        l = max(l, i-offset(coef_row))
    end
    l
end

function upper_bandwidth(D::DerivativeOperator)
    @unpack left_boundary, right_boundary, upper_coef = D.coefficients

    u = length(upper_coef)
    for (i,coef_row) in enumerate(left_boundary)
        u = max(u, 1+length(coef_row)-offset(coef_row)-i)
    end
    for (i,coef_row) in enumerate(right_boundary)
        u = max(u, 1+length(coef_row)-offset(coef_row)-i)
    end
    u
end

function left_boundary_weight(D::DerivativeOperator)
    @inbounds retval = D.Δx * D.coefficients.left_weights[1]
    retval
end

function right_boundary_weight(D::DerivativeOperator)
    @inbounds retval = D.Δx * D.coefficients.right_weights[1]
    retval
end


text:
######################################################################
#
#     IJulia Integration
#
#     - consist in defining a writemime(::IO, m::MIME"text/html", v::VegaLiteVis)
#       including scripts that load the required libraries (D3, vega, vega-lite)
#     - Will it work with other html backends ?  probably not.
#
######################################################################

import Base.writemime

# FIXME : Apparently, loading local js files is not allowed by the browser
#   => libraries are loaded externally in the `require.config`

# function jslibpath(url...)
#   libpath = joinpath(dirname(@__FILE__), "..", "assets", "bower_components", url...)
#   replace(libpath, "\\", "/")  # for windows...
# end
# // d3: "$(jslibpath("d3","d3.min.js"))",
# // vega: "$(jslibpath("vega", "vega.js"))",
# // vegalite: "$(jslibpath("vega-lite", "vega-lite.js"))",
# // vegaembed: "$(jslibpath("vega-embed", "vega-embed.js"))"

function writemime(io::IO, m::MIME"text/html", v::VegaLiteVis)
  divid = "vl" * randstring(10) # generated id for this plot

  fh = """

    <meta charset="utf-8">

  """
  # FIXME : understand why vega-embed can't be loaded
  #   SVG and SAVE_BUTTONS flags are ignored
  # vg.embed("#$divid", embedSpec, function(error, result) {

  write(io, fh)
end

export writemime


text:
function two_scales_correlation(prices::DataFrame, times::Vector{<:Real}, asseti::Symbol, assetj::Symbol, gamma::Real, num_grids::Real)
    vol_plus_version  = two_scales_volatility(prices[:,asseti] * gamma + (1-gamma) * prices[:,assetj] , times, Symbol("CompoundAsset__",asseti,"__+__",assetj,"__",gamma), num_grids)[1]
    vol_minus_version = two_scales_volatility(prices[:,asseti] * gamma - (1-gamma) * prices[:,assetj] , times, Symbol("CompoundAsset__",asseti,"__-__",assetj,"__",gamma), num_grids)[1]
    covv = (1/(4*gamma*(1-gamma))) * ( vol_plus_version^2 -  vol_minus_version^2 )
    # Note that we calculate these again rather than feeding in the ones calculated earlier because it is improtant the vols are calculated with the same data as the sums/differences.
    asseti_vol = two_scales_volatility(prices[:,asseti], times, Symbol("CompoundAsset__",asseti,"__ONLY"), num_grids)[1]
    assetj_vol = two_scales_volatility(prices[:,assetj], times, Symbol("CompoundAsset__",assetj,"__ONLY"), num_grids)[1]
    correl = covv/(asseti_vol * assetj_vol)
    return correl
end

function get_refresh_times_and_prices(ts::SortedDataFrame, asset1::Symbol, asset2::Symbol)
    assets = [asset1, asset2]
    at_times    = get_all_refresh_times(ts, assets)
    prices = random_value_in_interval(ts, at_times; assets = assets)
    times = prices[:,:Time]
    return prices, times
end

"""
    two_scales_covariance(ts::SortedDataFrame, assets::Vector{Symbol} = get_assets(ts);
                          regularisation::Union{Missing,Symbol} = :correlation_default,
                          regularisation_params::Dict = Dict(), only_regulise_if_not_PSD::Bool = false,
                          equalweight::Bool = false, num_grids::Real = default_num_grids(ts))

Estimation of a CovarianceMatrix using the two scale covariance method.
### Inputs
* `ts` - The tick data.
* `assets` - The assets you want to estimate volatilities for.
* `regularisation` - A symbol representing what regularisation technique should be used. If missing no regularisation is performed.
* `regularisation_params` - keyword arguments to be consumed in the regularisation algorithm.
* `only_regulise_if_not_PSD` - Should regularisation only be attempted if the matrix is not psd already.
* `equalweight` - Should we use equal weight for the two different linear combinations of assets. If false then an optimal weight is calculated (from volatilities).
* `num_grids` - Number of grids used in order in two scales estimation.
### Returns
* A `CovarianceMatrix`.
"""
function two_scales_covariance(ts::SortedDataFrame, assets::Vector{Symbol} = get_assets(ts);
                               regularisation::Union{Missing,Symbol} = :correlation_default,
                               regularisation_params::Dict = Dict(), only_regulise_if_not_PSD::Bool = false,
                               equalweight::Bool = false, num_grids::Real = default_num_grids(ts))

    two_scales_vol, micro_noise = two_scales_volatility(ts, assets; num_grids = num_grids)

    N = length(assets)
    mat = zeros(N,N)
    for i in 1:N
        asseti = assets[i]
        for j in i:N
            assetj = assets[j]
            gamma = min(max(0.1, equalweight ? 0.5 : two_scales_vol[assetj] / (two_scales_vol[asseti] + two_scales_vol[assetj]))   , 0.9)
            if (i == j)
                mat[i,j] = 1
            else
                zero_vol = (two_scales_vol[asseti] < 2*eps()) |  (two_scales_vol[assetj] < 2*eps())
                if zero_vol
                    mat[i,j] = 0.0
                else
                    prices, times = get_refresh_times_and_prices(ts, asseti, assetj)
                    mat[i,j] = two_scales_correlation(prices, times, asseti, assetj, gamma, num_grids)
                end
            end
        end
    end
    mat = Hermitian(mat)

    # Regularisation - It is done on the correlation matrix for this algo rather than the covariance matrix.
    dont_regulise = ismissing(regularisation) || (only_regulise_if_not_PSD && is_psd_matrix(mat))
    covmat = dont_regulise ? mat : regularise(mat, ts, assets, regularisation; regularisation_params... )

    vols = map(a -> two_scales_vol[a], assets)
    covmat = CovarianceMatrix(covmat, vols, assets, ts.time_period_per_unit)

    return covmat
end


text:
module Deriv_Utils
    using Statistics
    using Plots
    using BenchmarkTools
    using Profile
    using FFTW
    using Statistics
    using Optim
    using Images, FileIO, ImageIO
    using Printf

    #using Profile
    using LinearAlgebra

    # put the cwd on the path so Julia can find the module
    push!(LOAD_PATH, pwd()*"/main")
    using DHC_2DUtils

    function wst_S1_deriv_fast(image::Array{Float64,2}, filter_hash::Dict)

        # Use 2 threads for FFT
        FFTW.set_num_threads(2)

        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        # allocate output array
        dSdp  = zeros(Float64, Nx, Nx, Nf)

        # allocate image arrays for internal use
        #im_rdc_0_1 = Array{ComplexF64, 3}(undef, Nx, Ny, Nf)  # real domain, complex

        # Not sure what to do here -- for gradients I don't think we want these
        ## 0th Order
        #S0[1]   = mean(image)
        #norm_im = image.-S0[1]
        #S0[2]   = sum(norm_im .* norm_im)/(Nx*Ny)
        #norm_im ./= sqrt(Nx*Ny*S0[2])

        ## 1st Order
        im_fd_0 = fft(image)  # total power=1.0

        # unpack filter_hash
        f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_hash["filt_value"]

        zarr = zeros(ComplexF64, Nx, Nx)  # temporary array to fill with zvals

        # make a FFTW "plan" for an array of the given size and type
        P = plan_ifft(im_fd_0)  # P is an operator, P*im is ifft(im)

        # Loop over filters
        for f = 1:Nf
            f_i = f_ind[f]  # CartesianIndex list for filter
            f_v = f_val[f]  # Values for f_i

            zarr[f_i] = f_v .* im_fd_0[f_i]
            #Old: I_λ = P*zarr  # complex valued ifft of zarr, Should have been z_lambda in the latex conv
            #im_rdc_0_1[:,:,f] = I_lambda

            #xfac = 2 .* real(I_λ)
            #yfac = 2 .* imag(I_λ)
            # it is clearly stupid to transform back and forth, but we will need these steps for the S20 part
            #Old: ψ_λ  = realspace_filter(Nx, f_i, f_v)
            # convolution requires flipping wavelet direction, equivalent to flipping sign of imaginary part.
            # dSdp[:,:,f] = real.(conv(xfac,real(ψ_λ))) - real.(conv(yfac,imag(ψ_λ)))
            zarr[f_i] = f_v .* zarr[f_i] #Reusing zarr for F[conv(z_λ,ψ_λ)] = F[z_λ].*F[ψ_λ] = fz_fψ
            dSdp[:,:,f] = 2 .* real.(P*(zarr)) #real.(conv(I_λ, ψ_λ))
            zarr[f_i] .= 0   # reset zarr for next loop

        end
        return dSdp

    end

    function wst_S1S2_derivfast(image::Array{Float64,2}, filter_hash::Dict)
        #Works now.
        #Possible Bug: You assumed zero freq in wrong place and N-k is the wrong transformation.
        # Use 2 threads for FFT
        FFTW.set_num_threads(2)

        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        # allocate output array
        dS1dp  = zeros(Float64, Nf, Nx, Nx) #Mod zeros(Float64, Nx, Nx, Nf)
        dS2dp  = zeros(Float64, Nf, Nf, Nx, Nx) #Mod

        # allocate image arrays for internal use
        im_rdc_0_1 = Array{ComplexF64, 3}(undef, Nx, Ny, Nf)  # real domain, complex
        im_rd_0_1  = Array{Float64, 3}(undef, Nx, Ny, Nf)  # real domain, complex

        # Not sure what to do here -- for gradients I don't think we want these
        ## 0th Order
        #S0[1]   = mean(image)
        #norm_im = image.-S0[1]
        #S0[2]   = sum(norm_im .* norm_im)/(Nx*Ny)
        #norm_im ./= sqrt(Nx*Ny*S0[2])

        ## 1st Order
        im_fd_0 = fft(image)  # total power=1.0

        # unpack filter_hash
        f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_hash["filt_value"]

        f_ind_rev = [[CartesianIndex(mod1(Nx+2 - ci[1],Nx), mod1(Nx+2 - ci[2],Nx)) for ci in f_Nf] for f_Nf in f_ind]

        #CartesianIndex(17,17) .- f_ind
        # tmp arrays for all tmp[f_i] = arr[f_i] .* f_v opns
        zarr1 = zeros(ComplexF64, Nx, Nx)
        fz_fψ1 = zeros(ComplexF64, Nx, Nx)
        fterm_a = zeros(ComplexF64, Nx, Nx)
        fterm_ct1 = zeros(ComplexF64, Nx, Nx)
        fterm_ct2 = zeros(ComplexF64, Nx, Nx)
        #rterm_bt2 = zeros(ComplexF64, Nx, Nx)
        # make a FFTW "plan" for an array of the given size and type
        P = plan_ifft(im_fd_0)  # P is an operator, P*im is ifft(im)

        # Loop over filters
        for f1 = 1:Nf
            f_i1 = f_ind[f1]  # CartesianIndex list for filter1
            f_v1 = f_val[f1]  # Values for f_i1
            f_i1rev = f_ind_rev[f1]

            zarr1[f_i1] = f_v1 .* im_fd_0[f_i1] #F[z] #BUG: here????
            zarr1_rd = P*zarr1 #for frzi
            fz_fψ1[f_i1] = f_v1 .* zarr1[f_i1]
            fz_fψ1_rd = P*fz_fψ1
            dS1dp[f1, :,:] = 2 .* real.(fz_fψ1_rd) #real.(conv(I_λ, ψ_λ)) #Mod
            #CHECK; that this equals derivS1fast code

            #dS2 loop prep
            #ψ_λ1  = realspace_filter(Nx, f_i1, f_v1) #Slow
            #fcψ_λ1 = fft(conj.(ψ_λ1)) #Slow
            I_λ1 = sqrt.(abs2.(zarr1_rd))
            fI_λ1 = fft(I_λ1)
            rterm_bt1 = zarr1_rd./I_λ1 #Z_λ1/I_λ1_bar.
            rterm_bt2 = conj.(zarr1_rd)./I_λ1
            for f2 = 1:Nf
                f_i2 = f_ind[f2]  # CartesianIndex list for filter2
                f_v2 = f_val[f2]  # Values for f_i2

                fterm_a[f_i2] = fI_λ1[f_i2] .* (f_v2).^2 #fterm_a = F[I_λ1].F[ψ_λ]^2
                rterm_a = P*fterm_a                #Finv[fterm_a]
                fterm_bt1 = fft(rterm_a .* rterm_bt1) #F[Finv[fterm_a].*Z_λ1/I_λ1_bar] for T1
                fterm_bt2 = fft(rterm_a .* rterm_bt2) #F[Finv[fterm_a].*Z_λ1_bar/I_λ1_bar] for T2
                fterm_ct1[f_i1] = fterm_bt1[f_i1] .* f_v1    #fterm_b*F[ψ_λ]
                #fterm_ct2 = fterm_bt2 .* fcψ_λ1             #Slow
                #println(size(fterm_ct2[f_i1rev]),size(fterm_bt2[f_i1rev]),size(f_v1),size(conj.(f_v1)))
                fterm_ct2[f_i1rev] = fterm_bt2[f_i1rev] .* f_v1 #f_v1 is real
                #fterm_ct2slow = fterm_bt2 .* fcψ_λ1
                dS2dp[f1, f2, :, :] = real.(P*(fterm_ct1 + fterm_ct2)) #Mod THIS SHOULD be Nx*f2-1 + f1 or do the reshaping in loss func.

                #Reset f2 specific tmp arrays to 0
                fterm_a[f_i2] .= 0
                fterm_ct1[f_i1] .=0
                fterm_ct2[f_i1rev] .=0

            end
            # reset all reused variables to 0
            zarr1[f_i1] .= 0
            fz_fψ1[f_i1] .= 0

        end
        return vcat(dS1dp, reshape(dS2dp, (Nf^2, Nx, Nx))) #Mod
    end

    function wst_S20_deriv(image::Array{Float64,2}, filter_hash::Dict, nthread::Int=1)

        # Use nthread threads for FFT -- but for Nx<512 nthread=1 is fastest.  Overhead?
        FFTW.set_num_threads(nthread)

        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        # allocate output array
        dS20dα  = zeros(Float64, Nx, Nx, Nf, Nf)

        # allocate image arrays for internal use
        im_rdc = Array{ComplexF64, 3}(undef, Nx, Ny, Nf)  # real domain, complex
        im_rd  = Array{Float64, 3}(undef, Nx, Ny, Nf)  # real domain, complex
        im_fd  = fft(image)

        # unpack filter_hash
        f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_hash["filt_value"]
        zarr = zeros(ComplexF64, Nx, Nx)  # temporary array to fill with zvals

        # make a FFTW "plan" a complex array, both forward and inverse transform
        P_fft  = plan_fft(im_fd)   # P_fft is an operator,  P_fft*im is fft(im)
        P_ifft = plan_ifft(im_fd)  # P_ifft is an operator, P_ifft*im is ifft(im)

        # Loop over filters
        for f = 1:Nf
            f_i = f_ind[f]  # CartesianIndex list for filter
            f_v = f_val[f]  # Values for f_i

            zarr[f_i] = f_v .* im_fd[f_i]
            Z_λ = P_ifft*zarr  # complex valued ifft of zarr
            zarr[f_i] .= 0   # reset zarr for next loop
            im_rdc[:,:,f] = Z_λ
            im_rd[:,:,f]  = abs.(Z_λ)
        end

        zarr = zeros(ComplexF64, Nx, Nx)  # temporary array to fill with zvals
        for f2 = 1:Nf
            f_i = f_ind[f2]  # CartesianIndex list for filter
            f_v = f_val[f2]  # Values for f_i
            uvec = im_rdc[:,:,f2] ./ im_rd[:,:,f2]
            for f1 = 1:Nf
                temp = P_fft*(im_rd[:,:,f1].*uvec)
                zarr[f_i] = f_v .* temp[f_i]

                Z1dZ2 = real.(P_ifft*zarr)
                #  It is possible to do this with rifft, but it is not much faster...
                #   Z1dZ2 = myrealifft(zarr)
                dS20dα[:,:,f1,f2] += Z1dZ2
                dS20dα[:,:,f2,f1] += Z1dZ2
                zarr[f_i] .= 0   # reset zarr for next loop
            end
        end
        return dS20dα
    end

    function wst_S12_deriv(image::Array{Float64,2}, filter_hash::Dict)
        FFTW.set_num_threads(2)

        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        # unpack filter_hash
        f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_hash["filt_value"]

        zarr = zeros(ComplexF64, Nf, Nx, Nx)  # temporary array to fill with zvals
        ## 1st Order
        im_fd_0 = fft(image)  # total power=1.0
        # make a FFTW "plan" for an array of the given size and type
        P = plan_ifft(im_fd_0)  # P is an operator, P*im is ifft(im)
        for f1 = 1:Nf
            f_i1 = f_ind[f1]  # CartesianIndex list for filter1
            f_v1 = f_val[f1]  # Values for f_i1
            zarr[f1, f_i1] = im_fd_0[f_i1] .* f_v1 #check if this notation is fine
        end
        zarrc =conj.(zarr)
        absz = abs.(zarr)

        #Implemented in the more memory intensive / parallelized ops way
        #fhash_indsarr = [hcat((x->[x[1], x[2]]).(fh)...)' for fh in fhash["filt_index"]] #Each elem gives you fhash["filt_index"] in non-CartesianIndices
        #dabsz = #Mem Req: Nf*Nx^2*sp*Nx^2 and sp~2%

        #Less Mem Route
        dabsz_absz = zeros(Float64, Nf, Nf, Nx, Nx)  # Lookup arrays for dS12 terms
        xigrid = reshape(CartesianIndices((1:Nx, 1:Nx)), (1, Nx^2))
        xigrid = hcat((x->[x[1], x[2]]).(xigrid)...) #2*Nx^2

        for f1 = 1:Nf
            f_i1 = f_ind[f1]  # CartesianIndex list for filter1
            f_v1 = f_val[f1]  # Values for f_i1

            for f2 = 1:Nf
                f_i2 = f_ind[f2]  # CartesianIndex list for filter1
                f_v2 = f_val[f2]  # Values for f_i1
                pnz = intersect(f_i1, f_i2) #Fd indices where both filters are nonzero
                if length(pnz)!=0
                    pnz_arr = vcat((x->[x[1], x[2]]').(pnz)...) #pnz*2
                    Φmat =  exp.((-2π*im)/Nx .* ((pnz_arr .- 1) * (xigrid .- 1))) #pnz*Nx^2
                    f_v1pnz = f_v1[findall(in(pnz), f_i1)]
                    t2 = real.(zarrc[f1, pnz] .* Φmat) #Check
                    #println(size(f_v1pnz), size(t2), size(absz[f2, pnz]))
                    term = sum(((absz[f2, pnz] ./ absz[f1, pnz]) .* f_v1pnz) .* t2, dims=1) #p*Nx^2 -> 1*Nx^2
                    dabsz_absz[f1, f2, :, :] = reshape(term, (Nx, Nx))
                end
            end
        end
        dS12 = dabsz_absz + permutedims(dabsz_absz, [2, 1, 3, 4])
        return dS12
    end

    function wst_S2_deriv_sum(image::Array{Float64,2}, filter_hash::Dict, wt::Array{Float64})
        FFTW.set_num_threads(2)

        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        if size(wt)!=(Nf, Nf) error("Wt has wrong shape") end

        f_ind   = filter_hash["filt_index"]  # (J, L) array of filters represented as index value pairs
        f_val   = filter_hash["filt_value"]
        f_ind_rev = [[CartesianIndex(mod1(Nx+2 - ci[1],Nx), mod1(Nx+2 - ci[2],Nx)) for ci in f_Nf] for f_Nf in f_ind]

        im_fd_0 = fft(image)
        P = plan_fft(im_fd_0)
        Pi = plan_ifft(im_fd_0)
        zarr1 = zeros(ComplexF64, Nx, Nx)
        Ilam1 = zeros(ComplexF64, Nx, Nx)
        fterm_bt1 = zeros(ComplexF64, Nx, Nx)
        fterm_bt2 = zeros(ComplexF64, Nx, Nx)
        ifterm = zeros(ComplexF64, Nx, Nx)
        dLossdα = zeros(Float64, Nx, Nx)
        for f1=1:Nf
            vlam1 = zeros(ComplexF64, Nx, Nx) #defined for each lambda1, computed by summing over lambda2
            f_i1 = f_ind[f1]
            f_v1 = f_val[f1]
            f_i1rev = f_ind_rev[f1]

            zarr1[f_i1] = im_fd_0[f_i1] .* f_v1
            zarr1_rd = Pi*(zarr1)
            Ilam1 = abs.(zarr1_rd)
            fIlam1 = P*Ilam1

            for f2=1:Nf
                f_i2 = f_ind[f2]
                f_v2 = f_val[f2]
                vlam1[f_i2] += wt[f1, f2] .* f_v2.^2
            end
            #vlam1= sum_λ2 w.ψl1,l2^2
            vlam1 = (Pi*(fIlam1 .* vlam1))./Ilam1 #sum over f2 -> (Nx, Nx) -> real space Vλ1/Iλ1
            fterm_bt1 = P*(vlam1 .* zarr1_rd)
            fterm_bt2 = P*(vlam1 .* conj.(zarr1_rd))
            ifterm[f_i1] = fterm_bt1[f_i1] .* f_v1
            ifterm[f_i1rev] += fterm_bt2[f_i1rev] .* f_v1
            dLossdα += real.(Pi*ifterm)

            #zero out alloc
            zarr1[f_i1] .= 0
            ifterm[f_i1] .=0
            ifterm[f_i1rev] .=0

        end
        return dLossdα

    end

    function wst_S1S2_derivsum_comb(image::Array{Float64,2}, filter_hash::Dict, wt::Array{Float64})
        # array sizes
        (Nx, Ny)  = size(image)
        (Nf, )    = size(filter_hash["filt_index"])

        if size(wt)!=(Nf + (Nf*Nf),) error("Wt has wrong shape: the argument should contain both S1 and S2 weights") end
        wts1 = reshape(wt[1:Nf], (Nf, 1))
        wts2 = reshape(wt[Nf+1:end], (Nf, Nf))

        dwS1 = reshape(DHC_2DUtils.wst_S1_deriv(image, filter_hash), (Nx^2, Nf)) * wts1 #NOTE: Replace this with DHC_2DUtils.wst_S1_deriv
        dwS2 = reshape(wst_S2_deriv_sum(image, filter_hash, wts2), (Nx^2, 1))
        #println(size(dwS1), size(dwS2))
        return reshape(dwS1 + dwS2, Nx^2)

    end

    #######Testing Utilities Begin Here##################
    function derivtestfast(Nx)
       eps = 1e-4
       fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
       im = rand(Float64, Nx, Nx).*0.1
       im[6,6]=1.0
       im1 = copy(im)
       im0 = copy(im)
       im1[2,3] += eps/2
       im0[2,3] -= eps/2
       blarg = wst_S1_deriv(im, fhash)
       blargfast = wst_S1_deriv_fast(im, fhash)

       der0=DHC(im0,fhash,doS2=false)
       der1=DHC(im1,fhash,doS2=false)
       dS = (der1-der0) ./ eps

       diff = dS[3:end]-blarg[2,3,:]
       println(dS[3:end])
       println("and")
       println(blarg[2,3,:])
       println("stdev: ",std(diff))
       println("New fewer conv ds1")
       println(dS[3:end])
       println("and")
       println(blargfast[2,3,:])
       println("stdev: ",std(diff))
       #=#Delete: blarg shouldnt be 0 for all other pixels since the coefficient isnt constant wrt other pixels!!
       #Check: that the derivative is 0 for all pixels?? THink about this... need to check against autodiff?
       Nf = length(fhash["filt_index"])
       mask = trues((Nx,Nx, Nf))
       println("Nf",Nf,"Check",(size(mask)==size(blargfast)))
       mask[2, 3, :] .= false
       println("Mean deriv wrt to other pixels (should be 0 for all)",mean(blargfast[mask]), "Maximum",maximum(blargfast[mask]))
       #For Nx=128, Mean is actually 0.002, the max is 0.10. Seems concerning since eps is 1e-4!?
       #For Nx=8, Mean is 0.01, max is 0.90 -- some kind of normalization issue?
       #plot(dS[3:end])
       #plot!(blarg[2,3,:])
       #plot(diff)=#
       return
   end

   function derivtestS1S2(Nx; mode="slow")
       eps = 1e-4
       fhash = fink_filter_hash(1, 8, nx=Nx, pc=1, wd=1)
       im = rand(Float64, Nx, Nx).*0.1
       im[6,6]=1.0
       im1 = copy(im)
       im0 = copy(im)
       im1[2,3] += eps/2
       im0[2,3] -= eps/2
       dS1dp_existing = wst_S1_deriv(im, fhash)
       if mode=="slow"
           dS1dp, dS2dp = wst_S1S2_deriv(im, fhash)
       else
           dS1dp, dS2dp = wst_S1S2_derivfast(im, fhash)
       end

       der0=DHC(im0,fhash,doS2=true,doS20=false)
       der1=DHC(im1,fhash,doS2=true,doS20=false)
       dSlim = (der1-der0) ./ eps
       Nf = length(fhash["filt_index"])
       lasts1ind = 2+Nf

       dS1lim23 = dSlim[3:lasts1ind]
       dS2lim23 = dSlim[lasts1ind+1:end]
       dS1dp_23 = dS1dp[2, 3, :]
       dS2dp_23 = dS2dp[2, 3, :, :]
       #=
       println("Checking dS1dp using existing deriv")
       Nf = length(fhash["filt_index"])
       lasts1ind = 2+Nf
       diff = dS1dp_existing[2, 3, :]-dSlim[3:lasts1ind]
       println(dSlim[3:lasts1ind])
       println("and")
       println(dS1dp_existing[2, 3, :])
       println("stdev: ",std(diff))
       println("--------------------------")=#
       println("Checking dS1dp using dS1S2")
       derdiff = dS1dp_23 - dS1lim23
       #println(dS1lim23)
       #println("and")
       #println(dS1dp_23)
       txt= @sprintf("Range of S1 der0, der1: Max= %.3E, %.3E Min=%.3E, %.3E ",maximum(der0[3:lasts1ind]),maximum(der1[3:lasts1ind]),minimum(der0[3:lasts1ind]),minimum(der1[3:lasts1ind]))
       print(txt)
       txt = @sprintf("Checking dS1dp using S1S2 deriv, mean(abs(diff/dS1lim)): %.3E", mean(abs.(derdiff./dS1lim23)))
       print(txt)
       println("stdev: ",std(derdiff))
       txt= @sprintf("Range of dS1lim: Max= %.3E Min= %.3E",maximum(abs.(dS1lim23)),minimum(abs.(dS1lim23)))
       print(txt)
       println("--------------------------")
       println("Checking dS2dp using S1S2 deriv")
       Nf = length(fhash["filt_index"])
       println("Shape check",size(dS2dp_23),size(dS2lim23))
       derdiff = reshape(dS2dp_23, Nf*Nf) - dS2lim23 #Column major issue here?
       println("Difference",derdiff)
       txt = @sprintf("Range of S2 der0, der1: Max= %.3E, %.3E Min=%.3E, %.3E \n",maximum(der0[lasts1ind+1:end]),maximum(der1[lasts1ind+1:end]),minimum(der0[lasts1ind+1:end]),minimum(der1[lasts1ind+1:end]))
       print(txt)
       txt = @sprintf("Range of dS2lim: Max=%.3E  Min=%.3E \n",maximum(abs.(dS2lim23)),minimum(abs.(dS2lim23)))
       print(txt)
       txt = @sprintf("Difference between dS2dp and dSlim Mean: %.3E stdev: %.3E mean(abs(derdiff/dS2lim)): %.3E \n",mean(derdiff),std(derdiff), mean(abs.(derdiff./dS2lim23)))
       print(txt)
       println("Examining only those S2 coeffs which are greater than eps=1e-3 for der0")
       eps_mask = findall(der0[lasts1ind+1:end] .> 1E-3)
       der0_s2good = der0[lasts1ind+1:end][eps_mask]
       der1_s2good = der1[lasts1ind+1:end][eps_mask]
       dSlim_s2good = dS2lim23[eps_mask]
       derdiff_good = derdiff[eps_mask]
       print(derdiff_good)
       txt = @sprintf("\nRange of S2 der0, der1: Max= %.3E, %.3E Min=%.3E, %.3E \n",maximum(der0_s2good),maximum(der1_s2good),minimum(der0_s2good),minimum(der1_s2good))
       print(txt)
       txt = @sprintf("Range of dS2lim: Max=%.3E  Min=%.3E \n",maximum(abs.(dSlim_s2good)),minimum(abs.(dSlim_s2good)))
       print(txt)
       txt = @sprintf("Mean: %.3E stdev: %.3E mean(abs(derdiff/dS2lim)): %.3E \n",mean(derdiff_good),std(derdiff_good), mean(abs.(derdiff_good ./dSlim_s2good)))
       print(txt)

       #plot(dS[3:end])
       #plot!(blarg[2,3,:])
       #plot(diff)
       return
   end
end

max_stars_repo_name: andrew-saydjari/DHC

max_issues_repo_name: andrew-saydjari/DHC

max_forks_repo_name: andrew-saydjari/DHC


text:
using Turing, Random, Test
using Turing.Core: ResampleWithESSThreshold
using Turing.Inference: getspace, resample_systematic, resample_multinomial

using Random

dir = splitdir(splitdir(pathof(Turing))[1])[1]
include(dir*"/test/test_utils/AllUtils.jl")

@testset "SMC" begin
    @turing_testset "constructor" begin
        s = SMC()
        @test s.resampler == ResampleWithESSThreshold()
        @test getspace(s) === ()

        s = SMC(:x)
        @test s.resampler == ResampleWithESSThreshold()
        @test getspace(s) === (:x,)

        s = SMC((:x,))
        @test s.resampler == ResampleWithESSThreshold()
        @test getspace(s) === (:x,)

        s = SMC(:x, :y)
        @test s.resampler == ResampleWithESSThreshold()
        @test getspace(s) === (:x, :y)

        s = SMC((:x, :y))
        @test s.resampler == ResampleWithESSThreshold()
        @test getspace(s) === (:x, :y)

        s = SMC(0.6)
        @test s.resampler === ResampleWithESSThreshold(resample_systematic, 0.6)
        @test getspace(s) === ()

        s = SMC(0.6, (:x,))
        @test s.resampler === ResampleWithESSThreshold(resample_systematic, 0.6)
        @test getspace(s) === (:x,)

        s = SMC(resample_multinomial, 0.6)
        @test s.resampler === ResampleWithESSThreshold(resample_multinomial, 0.6)
        @test getspace(s) === ()

        s = SMC(resample_multinomial, 0.6, (:x,))
        @test s.resampler === ResampleWithESSThreshold(resample_multinomial, 0.6)
        @test getspace(s) === (:x,)

        s = SMC(resample_systematic)
        @test s.resampler === resample_systematic
        @test getspace(s) === ()

        s = SMC(resample_systematic, (:x,))
        @test s.resampler === resample_systematic
        @test getspace(s) === (:x,)
    end

    @turing_testset "models" begin
        @model function normal()
            a ~ Normal(4,5)
            3 ~ Normal(a,2)
            b ~ Normal(a,1)
            1.5 ~ Normal(b,2)
            a, b
        end

        tested = sample(normal(), SMC(), 100);

        # failing test
        @model function fail_smc()
            a ~ Normal(4,5)
            3 ~ Normal(a,2)
            b ~ Normal(a,1)
            if a >= 4.0
                1.5 ~ Normal(b,2)
            end
            a, b
        end

        @test_throws ErrorException sample(fail_smc(), SMC(), 100)
    end

    @turing_testset "logevidence" begin
        Random.seed!(100)

        @model function test()
            a ~ Normal(0, 1)
            x ~ Bernoulli(1)
            b ~ Gamma(2, 3)
            1 ~ Bernoulli(x / 2)
            c ~ Beta()
            0 ~ Bernoulli(x / 2)
            x
        end

        chains_smc = sample(test(), SMC(), 100)

        @test all(isone, chains_smc[:x])
        @test chains_smc.logevidence ≈ -2 * log(2)
    end
end

@testset "PG" begin
    @turing_testset "constructor" begin
        s = PG(10)
        @test s.nparticles == 10
        @test s.resampler == ResampleWithESSThreshold()
        @test getspace(s) === ()

        s = PG(20, :x)
        @test s.nparticles == 20
        @test s.resampler == ResampleWithESSThreshold()
        @test getspace(s) === (:x,)

        s = PG(30, (:x,))
        @test s.nparticles == 30
        @test s.resampler == ResampleWithESSThreshold()
        @test getspace(s) === (:x,)

        s = PG(40, :x, :y)
        @test s.nparticles == 40
        @test s.resampler == ResampleWithESSThreshold()
        @test getspace(s) === (:x, :y)

        s = PG(50, (:x, :y))
        @test s.nparticles == 50
        @test s.resampler == ResampleWithESSThreshold()
        @test getspace(s) === (:x, :y)

        s = PG(60, 0.6)
        @test s.nparticles == 60
        @test s.resampler === ResampleWithESSThreshold(resample_systematic, 0.6)
        @test getspace(s) === ()

        s = PG(70, 0.6, (:x,))
        @test s.nparticles == 70
        @test s.resampler === ResampleWithESSThreshold(resample_systematic, 0.6)
        @test getspace(s) === (:x,)

        s = PG(80, resample_multinomial, 0.6)
        @test s.nparticles == 80
        @test s.resampler === ResampleWithESSThreshold(resample_multinomial, 0.6)
        @test getspace(s) === ()

        s = PG(90, resample_multinomial, 0.6, (:x,))
        @test s.nparticles == 90
        @test s.resampler === ResampleWithESSThreshold(resample_multinomial, 0.6)
        @test getspace(s) === (:x,)

        s = PG(100, resample_systematic)
        @test s.nparticles == 100
        @test s.resampler === resample_systematic
        @test getspace(s) === ()

        s = PG(110, resample_systematic, (:x,))
        @test s.nparticles == 110
        @test s.resampler === resample_systematic
        @test getspace(s) === (:x,)
    end

    @turing_testset "logevidence" begin
        Random.seed!(100)

        @model function test()
            a ~ Normal(0, 1)
            x ~ Bernoulli(1)
            b ~ Gamma(2, 3)
            1 ~ Bernoulli(x / 2)
            c ~ Beta()
            0 ~ Bernoulli(x / 2)
            x
        end

        chains_pg = sample(test(), PG(10), 100)

        @test all(isone, chains_pg[:x])
        @test chains_pg.logevidence ≈ -2 * log(2) atol = 0.01
    end
end

# @testset "pmmh.jl" begin
#     @turing_testset "pmmh constructor" begin
#         N = 2000
#         s1 = PMMH(N, SMC(10, :s), MH(1,(:m, s -> Normal(s, sqrt(1)))))
#         s2 = PMMH(N, SMC(10, :s), MH(1, :m))
#         s3 = PIMH(N, SMC())
#
#         c1 = sample(gdemo_default, s1)
#         c2 = sample(gdemo_default, s2)
#         c3 = sample(gdemo_default, s3)
#     end
#     @numerical_testset "pmmh inference" begin
#         alg = PMMH(2000, SMC(20, :m), MH(1, (:s, GKernel(1))))
#         chain = sample(gdemo_default, alg)
#         check_gdemo(chain, atol = 0.1)
#
#         # PMMH with prior as proposal
#         alg = PMMH(2000, SMC(20, :m), MH(1, :s))
#         chain = sample(gdemo_default, alg)
#         check_gdemo(chain, atol = 0.1)
#
#         # PIMH
#         alg = PIMH(2000, SMC())
#         chain = sample(gdemo_default, alg)
#         check_gdemo(chain)
#
#         # MoGtest
#         pmmh = PMMH(2000,
#             SMC(10, :z1, :z2, :z3, :z4),
#             MH(1, :mu1, :mu2))
#         chain = sample(MoGtest_default, pmmh)
#
#         check_MoGtest_default(chain, atol = 0.1)
#     end
# end

# @testset "ipmcmc.jl" begin
#     @turing_testset "ipmcmc constructor" begin
#         Random.seed!(125)
#
#         N = 50
#         s1 = IPMCMC(10, N, 4, 2)
#         s2 = IPMCMC(10, N, 4)
#
#         c1 = sample(gdemo_default, s1)
#         c2 = sample(gdemo_default, s2)
#     end
#     @numerical_testset "ipmcmc inference" begin
#         alg = IPMCMC(30, 500, 4)
#         chain = sample(gdemo_default, alg)
#         check_gdemo(chain)
#
#         alg2 = IPMCMC(15, 100, 10)
#         chain2 = sample(MoGtest_default, alg2)
#         check_MoGtest_default(chain2, atol = 0.2)
#     end
# end

@turing_testset "resample.jl" begin
    D = [0.3, 0.4, 0.3]
    num_samples = Int(1e6)
    resSystematic = Turing.Inference.resample_systematic(D, num_samples )
    resStratified = Turing.Inference.resample_stratified(D, num_samples )
    resMultinomial= Turing.Inference.resample_multinomial(D, num_samples )
    resResidual   = Turing.Inference.resample_residual(D, num_samples )
    Turing.Inference.resample(D)
    resSystematic2=Turing.Inference.resample(D, num_samples )

    @test sum(resSystematic .== 2) ≈ (num_samples * 0.4) atol=1e-3*num_samples
    @test sum(resSystematic2 .== 2) ≈ (num_samples * 0.4) atol=1e-3*num_samples
    @test sum(resStratified .== 2) ≈ (num_samples * 0.4) atol=1e-3*num_samples
    @test sum(resMultinomial .== 2) ≈ (num_samples * 0.4) atol=1e-2*num_samples
    @test sum(resResidual .== 2) ≈ (num_samples * 0.4) atol=1e-2*num_samples
end


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

using Markdown
using InteractiveUtils

# ╔═╡ 0831131a-077d-11eb-2412-fbd0e7cc3436
using DrWatson

# ╔═╡ 08314e98-077d-11eb-1d86-f9c9dd52e2c3
begin
	@quickactivate "SR2TuringPluto"
	using Turing
	using StatisticalRethinking
end

# ╔═╡ 041b4048-077d-11eb-20b8-6f8b8ee72626
md"## Clip-05-18-27t.jl"

# ╔═╡ f2973922-077b-11eb-3833-9f97491e0ae2
begin
	df = CSV.read(sr_datadir("WaffleDivorce.csv"), DataFrame)
	df.D = zscore(df.Divorce)
	df.M = zscore(df.Marriage)
	z_age, unz_age = zscore_transform(df.MedianAgeMarriage)
	df.A = z_age(df.MedianAgeMarriage)
end;

# ╔═╡ 3d995aa2-08e1-11eb-09f7-35b6aa59fe73
@model function m5_3_A(A, M, D)
    # A -> M
    aM ~ Normal(0, 0.2)
    bAM ~ Normal(0, 0.5)
    σ_M ~ Exponential(1)
    μ_M = lin(aM, A, bAM)
    M ~ MvNormal(μ_M, σ_M)
    # A -> D <- M
    a ~ Normal(0, 0.2)
    bM ~ Normal(0, 0.5)
    bA ~ Normal(0, 0.5)
    σ ~ Exponential(1)
    μ = lin(a, A, bA, M, bM)
    D ~ MvNormal(μ, σ)
end

# ╔═╡ 3d99976a-08e1-11eb-344e-f193f1387780
begin
	quap5_3_At = quap(m5_3_A(df.A, df.M, df.D))

	A_seq = range(-2, 2, length = 30)
	post5_3_At = DataFrame(rand(quap5_3_At.distr, 1000)', quap5_3_At.params)

	μM = lin(post5_3_At.aM', A_seq, post5_3_At.bAM')
	M = rand.(Normal.(μM, post5_3_At.σ_M'))
	μD = lin(post5_3_At.a', A_seq, post5_3_At.bA', M, post5_3_At.bM')
	D = rand.(Normal.(μD, post5_3_At.σ')) |> meanlowerupper

	plot(A_seq, D.mean, ribbon = (D.mean .- D.lower, D.upper .- D.mean))
	plot!(legend = false, xlabel = "manipulated A", ylabel = "counterfactual D")
end

# ╔═╡ 3dac58dc-08e1-11eb-33df-43cc24f8f0a5
begin
	M_seq = range(-2, 2, length = 30)
	A_seq1 = zeros(length(M_seq))
	M2 = rand.(Normal.(lin(post5_3_At.aM', A_seq, post5_3_At.bAM'), post5_3_At.σ_M'))
	μD2 = lin(post5_3_At.a', A_seq1, post5_3_At.bA', M_seq, post5_3_At.bM')
	D1 = rand.(Normal.(μD2, post5_3_At.σ')) |> meanlowerupper

	plot(M_seq, D1.mean, ribbon = (D1.mean .- D1.lower, D1.upper .- D1.mean))
	plot!(legend = false, xlabel = "manipulated M", ylabel = "counterfactual D")
end

# ╔═╡ 3dad5228-08e1-11eb-1b14-6d412246f714
begin
	post = DataFrame(rand(quap5_3_At.distr, 1000)', quap5_3_At.params)
	M_sim = rand.(Normal.(post.aM' .+ A_seq .* post.bAM', post.σ_M'))
	D_sim = rand.(Normal.(post.a' .+ A_seq .* post.bA', post.σ'))
end

# ╔═╡ 3d9a2388-08e1-11eb-28d1-131c48d85129
begin
	M1 = rand.(Normal.(lin(post.aM', (20.0, 30.0) |> z_age, post5_3_At.bAM'), post5_3_At.σ_M'))
	μD1 = lin(post5_3_At.a', (20.0, 30.0) |> z_age, post5_3_At.bA', M1, post5_3_At.bM') |> meanlowerupper
	μD1.mean[2] - μD1.mean[1]
end

# ╔═╡ 3db8aeb6-08e1-11eb-3947-8fb3cd806df2
md"## End of clip-05--19-27t.jl"

# ╔═╡ Cell order:
# ╟─041b4048-077d-11eb-20b8-6f8b8ee72626
# ╠═0831131a-077d-11eb-2412-fbd0e7cc3436
# ╠═08314e98-077d-11eb-1d86-f9c9dd52e2c3
# ╠═f2973922-077b-11eb-3833-9f97491e0ae2
# ╠═3d995aa2-08e1-11eb-09f7-35b6aa59fe73
# ╠═3d99976a-08e1-11eb-344e-f193f1387780
# ╠═3d9a2388-08e1-11eb-28d1-131c48d85129
# ╠═3dac58dc-08e1-11eb-33df-43cc24f8f0a5
# ╠═3dad5228-08e1-11eb-1b14-6d412246f714
# ╟─3db8aeb6-08e1-11eb-3947-8fb3cd806df2


text:
# mostly a proof-of-concept (more explicitly: a terrible hack) for the moment
# to activate, include() this file at the normal (LineEdit) REPL
# Then, at any time, press Ctrl-O to edit in normal mode.
# To be able to execute code and search history etc, going back to insert mode is neccesary.
# For the moment, the first keypress after going to insert mode disappears sometimes.
import REPL
import REPL.LineEdit

include("Neovim.jl")

mutable struct NvimReplState
    active::Bool
    nv::Neovim.NvimClient
    s::LineEdit.MIState
    rbuf::IOBuffer
    function NvimReplState()
        state = new(false)
        state.nv = Neovim.nvim_spawn(state)
        state
    end
end

function Neovim.on_notify(s::NvimReplState, nv, name, args)
    @async begin
        update_screen()
        if name == "update"
            #done
        elseif name == "insert"
            s.active = false
        end
    end
end

const rstate = NvimReplState()
const channel = rstate.nv.channel_id

Neovim.command(rstate.nv, "au CursorMoved,TextChanged * call rpcnotify($channel, 'update')")
Neovim.command(rstate.nv, "au InsertEnter * call rpcnotify($channel, 'insert')")
Neovim.command(rstate.nv, "set ft=julia")

const nvim_keymap = Dict(
    "^O" => function (s,repl,c)
        ps = LineEdit.state(s, LineEdit.mode(s))
        nvim_normal(ps.terminal,s,repl)
    end
)

function update_screen()
    if !rstate.active
        return
    end
    nv = rstate.nv
    #FIXME: this is a terrible hack, use abstract_ui later
    code = current_buffer(nv)[:]
    curline, curpos = get_cursor(current_window(nv))
    truncate(rstate.rbuf, 0)
    write(rstate.rbuf, join(code, "\n"))
    cpos = curpos
    for i in 1:(curline-1)
        cpos += length(code[i])+1
    end
    seek(rstate.rbuf, cpos)
    LineEdit.refresh_line(rstate.s)
end

function nvim_normal(term, s, repl)
    rstate.s = s
    rstate.rbuf = rbuf = LineEdit.buffer(s)
    nv = rstate.nv

    buf_pos = position(rbuf)
    seek(rbuf, 0)
    cur_col = buf_pos
    lines = ByteString[]
    more = true
    input(nv, "\033")
    cursor = [0,0]
    while more
        line = readline(rbuf)
        if length(line) > 0 && line[end] == '\n'
            line = line[1:end-1]
        else
            more = false
        end
        push!(lines, line)
        if 0 <= cur_col < length(line)+1
            cursor = [length(lines), cur_col]
        end
        cur_col -= length(line)+1
    end
    current_buffer(nv)[:] = lines
    set_cursor(current_window(nv), cursor)

    rstate.active = true
    while rstate.active
        char = read(term, Uint8)
        if !rstate.active
            # FIXME: handle this in a better way
            break
        end

        input(nv, bytestring([char]))
        #LineEdit.edit_insert
        if char == 1; rstate.active = false; end
        update_screen()
    end
end

repl = Base.active_repl
main_mode = repl.interface.modes[1]
main_mode.keymap_dict = LineEdit.keymap_merge(main_mode.keymap_dict, LineEdit.normalize_keys(nvim_keymap))

nothing