text:
using DataDrivenDiffEq
using ModelingToolkit
using OrdinaryDiffEq

using LinearAlgebra

# Create a test problem
function lorenz(u,p,t)
    x, y, z = u

    ẋ = 10.0*(y - x)
    ẏ = x*(28.0-z) - y
    ż = x*y - (8/3)*z
    return [ẋ, ẏ, ż]
end

u0 = [1.0;0.0;0.0]
tspan = (0.0,100.0)
dt = 0.1
prob = ODEProblem(lorenz,u0,tspan)
sol = solve(prob, Tsit5(), saveat = dt, progress = true)

## Start the automatic discovery
ddprob = ContinuousDataDrivenProblem(sol)

@variables t x(t) y(t) z(t)
u = [x;y;z]
basis = Basis(polynomial_basis(u, 5), u, iv = t)
opt = STLSQ(exp10.(-5:0.1:-1))
ddsol = solve(ddprob, basis, opt, normalize = true)
print(ddsol, Val{true})


text:
function xorshift64(a::UInt64)::UInt64
	a ⊻= (a << 13)
	a ⊻= (a >> 7)
	a ⊻= (a << 17)
	a
end

function unique_sequence(n, seed)
	a = UInt64(seed)
	m = zeros(UInt64, n)
	for i in 1:n
		a = xorshift64(a)
		m[i] = a
	end
	m
end

function unique_sequence(::Type{T}, n, seed) where {T <: Int64}
	m = unique_sequence(n, seed)
	for i in 1:n
		m[i] >>= 1
	end
	Int64.(m)
end

function unique_sequence(::Type{T}, n, seed) where {T <: Float64}
	m = unique_sequence(n, seed)
	Float64.(m)
end

ext: jl


text:
import Base: +, -, *, /

immutable Site2D
    # NB: origin is at (x,y) = (0,0)
    x::Int
    y::Int

    Lx::Int
    Ly::Int
    index::Int

    function Site2D(x::Int, y::Int, Lx::Int, Ly::Int)
        # we don't want these asserts so as to have PBC capabilities
        # @assert 0 <= x < Lx
        # @assert 0 <= y < Ly
        index = mod(x, Lx) * Ly + mod(y, Ly) + 1
        return new(div(index-1, Ly), mod(index-1, Ly), Lx, Ly, index)
    end

    function Site2D(index::Int, Lx::Int, Ly::Int)
        @assert 1 <= index <= Lx * Ly
        return new(div(index-1, Ly), mod(index-1, Ly), Lx, Ly, index)
    end
end

function +(r1::Site2D, r2::Site2D)
    @assert r1.Lx == r2.Lx
    @assert r1.Ly == r2.Ly
    return Site2D(r1.x + r2.x, r1.y + r2.y, r1.Lx, r1.Ly)
end

function -(r1::Site2D, r2::Site2D)
    @assert r1.Lx == r2.Lx
    @assert r1.Ly == r2.Ly
    return Site2D(r1.x - r2.x, r1.y - r2.y, r1.Lx, r1.Ly)
end

*(c::Number, r::Site2D) = Site2D(convert(Int, c * r.x), convert(Int, c * r.y), r.Lx, r.Ly)
*(r::Site2D, c::Number) = Site2D(convert(Int, c * r.x), convert(Int, c * r.y), r.Lx, r.Ly)
/(r::Site2D, c::Number) = Site2D(convert(Int, r.x / c), convert(Int, r.y / c), r.Lx, r.Ly)

export +, -, *, /

unitx(Lx::Int, Ly::Int) = Site2D(1, 0, Lx, Ly)
unity(Lx::Int, Ly::Int) = Site2D(0, 1, Lx, Ly)

# rectangular, contiguous subsystem LAx/LAy sites long in x/y direction
# whose bottom-/left-most site is located at (x0, y0)
function make_subsystem(LAx::Int, LAy::Int, Lx::Int, Ly::Int, x0::Int = 0, y0::Int = 0)
    @assert 0 < LAx <= Lx
    @assert 0 < LAy <= Ly
    @assert 0 <= x0 < Lx
    @assert 0 <= y0 < Ly

    origin = Site2D(x0, y0, Lx, Ly)

    subsystem = Int[]
    for i in 1:LAx
        bottom_site = origin + (i-1)*unitx(Lx, Ly)
        for j in 1:LAy
            push!(subsystem, (bottom_site + (j-1)*unity(Lx, Ly)).index)
        end
    end

    return subsystem
end

#########################################################
#################### Neighbor tables ####################
#########################################################

function zigzag_strip_periodic_nn(Lx::Int)
    xhat = unitx(Lx, 1)
    Nsites = Lx
    table = Vector{Int}[]
    for i in 1:Nsites
        r = Site2D(i, Lx, 1) # actually a site on a 1D chain..
        push!(table, [(r + 2xhat).index,
                      (r + xhat).index,
                      (r - xhat).index,
                      (r - 2xhat).index])
    end
    return table
end

function kagome_strip_periodic_nn(Lx::Int)
    xhat = unitx(Lx, 3)
    yhat = unity(Lx, 3)
    Nsites = 3Lx
    table = Vector{Int}[]
    for i in 1:Nsites
        r = Site2D(i, Lx, 3)
        if r.y == 0
            push!(table, [(r + xhat).index,
                          (r + xhat + yhat).index,
                          (r + yhat).index,
                          (r - xhat).index])
        elseif r.y == 1
            push!(table, [(r + yhat).index,
                          (r - xhat + yhat).index,
                          (r - xhat - yhat).index,
                          (r - yhat).index])
        elseif r.y == 2
            push!(table, [(r + xhat).index,
                          (r - xhat).index,
                          (r - yhat).index,
                          (r + xhat - yhat).index])
        else
            @assert false
        end
    end
    return table
end

function kagome_strip_open_nn(Lx::Int)
    xhat = unitx(Lx, 3)
    yhat = unity(Lx, 3)
    Nsites = 3Lx-2
    table = Vector{Int}[]
    for i in 1:Nsites-1
        r = Site2D(i, Lx, 3)
        if r.x > 0 && r.x < Lx-2
            if r.y == 0
                push!(table, [(r + xhat).index,
                              (r + xhat + yhat).index,
                              (r + yhat).index,
                              (r - xhat).index])
            elseif r.y == 1
                push!(table, [(r + yhat).index,
                              (r - xhat + yhat).index,
                              (r - xhat - yhat).index,
                              (r - yhat).index])
            elseif r.y == 2
                push!(table, [(r + xhat).index,
                              (r - xhat).index,
                              (r - yhat).index,
                              (r + xhat - yhat).index])
            else
                @assert false
            end
        elseif r.x == 0
            if r.y == 0
                push!(table, [(r + xhat).index,
                              (r + xhat + yhat).index,
                              (r + yhat).index])
            elseif r.y == 1
                push!(table, [(r + yhat).index,
                              (r - yhat).index])
            elseif r.y == 2
                push!(table, [(r + xhat).index,
                              (r - yhat).index,
                              (r + xhat - yhat).index])
            else
                @assert false
            end
        elseif r.x == Lx-2
            if r.y == 0
                push!(table, [Nsites,
                              (r + yhat).index,
                              (r - xhat).index])
            elseif r.y == 1
                push!(table, [(r + yhat).index,
                              (r - xhat + yhat).index,
                              (r - xhat - yhat).index,
                              (r - yhat).index])
            elseif r.y == 2
                push!(table, [(r - xhat).index,
                              (r - yhat).index,
                              Nsites])
            else
                @assert false
            end
        else
            @assert false
        end
    end
    push!(table, [Nsites-1, Nsites-3]) # deal with the last site separately
    return table
end

function triangular_ladder_PBCx_PBCy_nn(Lx::Int, Ly::Int)
    xhat = unitx(Lx, Ly)
    yhat = unity(Lx, Ly)
    Nsites = Lx * Ly
    table = Vector{Int}[]
    for i in 1:Nsites
        r = Site2D(i, Lx, Ly)
        push!(table, [(r + xhat).index,
                      (r + xhat + yhat).index,
                      (r + yhat).index,
                      (r - xhat).index,
                      (r - xhat - yhat).index,
                      (r - yhat).index])
    end
    return table
end

function triangular_ladder_subsystem_wrap_nn(Lx::Int, Ly::Int, LAx::Int, LAy::Int)
    @assert LAy == Ly # require cylindrical subsystem for now
    xhat = unitx(Lx, Ly)
    yhat = unity(Lx, Ly)
    Nsites = Lx * Ly

    subsystemA = make_subsystem(LAx, LAy, Lx, Ly)
    tableA = Vector{Int}[]
    for i in subsystemA
        r = Site2D(i, Lx, Ly)
        push!(tableA, [in((r + xhat).index, subsystemA) ? (r + xhat).index : (r - (LAx-1)*xhat).index,
                       in((r + xhat + yhat).index, subsystemA) ? (r + xhat + yhat).index : (r - (LAx-1)*xhat + yhat).index,
                       (r + yhat).index, # this guy's always still in subsystemA given the assert
                       in((r - xhat).index, subsystemA) ? (r - xhat).index : (r + (LAx-1)*xhat).index,
                       in((r - xhat - yhat).index, subsystemA) ? (r - xhat - yhat).index : (r + (LAx-1)*xhat - yhat).index,
                       (r - yhat).index]) # this guy's always still in subsystemA given the assert
        @assert issubset(tableA[end], subsystemA)
    end

    subsystemB = setdiff(collect(1:Nsites), subsystemA)
    tableB = Vector{Int}[]
    for i in subsystemB
        r = Site2D(i, Lx, Ly)
        push!(tableB, [in((r + xhat).index, subsystemB) ? (r + xhat).index : (r - (LAx-1)*xhat).index,
                       in((r + xhat + yhat).index, subsystemB) ? (r + xhat + yhat).index : (r - (LAx-1)*xhat + yhat).index,
                       (r + yhat).index, # this guy's always still in subsystemB given the assert
                       in((r - xhat).index, subsystemB) ? (r - xhat).index : (r + (LAx-1)*xhat).index,
                       in((r - xhat - yhat).index, subsystemB) ? (r - xhat - yhat).index : (r + (LAx-1)*xhat - yhat).index,
                       (r - yhat).index]) # this guy's always still in subsystemB given the assert
        @assert issubset(tableB[end], subsystemB)
    end

    return tableA, tableB
end

function square_lattice_PBCx_PBCy_nn(Lx::Int, Ly::Int)
    xhat = unitx(Lx, Ly)
    yhat = unity(Lx, Ly)
    Nsites = Lx * Ly
    table = Vector{Int}[]
    for i in 1:Nsites
        r = Site2D(i, Lx, Ly)
        push!(table, [(r + xhat).index,
                      (r + yhat).index,
                      (r - xhat).index,
                      (r - yhat).index])
    end
    return table
end


text:
using MultipleViewGeometry, Test
using MultipleViewGeometry.ModuleTypes

ℳ = [Point2DH(rand(1,3)) for i = 1:10000]

moments(FundamentalMatrix(), (ℳ,ℳ))

# TODO


text:
__precompile__(true)

"""
Main module for `RNG.jl` -- a random number generator package for Julia Language.

This module exports two types and four submodules:

- [`AbstractRNG`](@ref).
- [`WrappedRNG`](@ref).
- [`PCG`](@ref).
- [`MersenneTwisters`](@ref).
- [`Random123`](@ref).
- [`Xorshifts`](@ref).
"""
module RNG

    export AbstractRNG
    export WrappedRNG
    export PCG, MersenneTwisters, Random123, Xorshifts

    include("common.jl")
    include("utils.jl")

    include("wrapped_rng.jl")

    include("./PCG/PCG.jl")
    include("./MersenneTwisters/MersenneTwisters.jl")
    include("./Random123/Random123.jl")
    include("./Xorshifts/Xorshifts.jl")

end


text:
using NeuralOperators
using Flux
using CUDA

if has_cuda()
    @info "CUDA is on"
    device = gpu
    CUDA.allowscalar(false)
else
    device = cpu
end

m = FourierNeuralOperator() |> device
loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]

n_train = 1800
n_test = 200
batchsize = 100
𝐱, 𝐲 = get_burgers_data(n=2048)

𝐱_train, 𝐲_train = 𝐱[:, :, 1:n_train], 𝐲[:, 1:n_train]
loader_train = Flux.DataLoader((𝐱_train, 𝐲_train), batchsize=batchsize, shuffle=true)

𝐱_test, 𝐲_test = 𝐱[:, :, end-n_test+1:end], 𝐲[:, end-n_test+1:end]
loader_test = Flux.DataLoader((𝐱_test, 𝐲_test), batchsize=batchsize, shuffle=false)

function validate()
    validation_losses = [loss(device(𝐱), device(𝐲)) for (𝐱, 𝐲) in loader_test]
    @info "loss: $(sum(validation_losses)/length(loader_test))"
end

data = [(𝐱, 𝐲) for (𝐱, 𝐲) in loader_train] |> device
opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
call_back = Flux.throttle(validate, 5, leading=false, trailing=true)
Flux.@epochs 500 @time(Flux.train!(loss, params(m), data, opt, cb=call_back))


text:
#  Copyright 2016, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
#  This Source Code Form is subject to the terms of the Mozilla Public
#  License, v. 2.0. If a copy of the MPL was not distributed with this
#  file, You can obtain one at http://mozilla.org/MPL/2.0/.

# generates code which converts an expression into a NodeData array (tape)
# parent is the index of the parent expression
# values is the name of the list of constants which appear in the expression
function parseNLExpr(m, x, tapevar, parent, values)

    if isexpr(x,:call) && length(x.args) >= 2 && isexpr(x.args[2],:generator)
        header = x.args[1]
        if issum(header)
            operatorid = operator_to_id[:+]
        elseif isprod(header)
            operatorid = operator_to_id[:*]
        else
            error("Unrecognized expression $header(...)")
        end
        codeblock = :(let; end)
        block = codeblock.args[1]
        push!(block.args, :(push!($tapevar, NodeData(CALL, $operatorid, $parent))))
        parentvar = gensym()
        push!(block.args, :($parentvar = length($tapevar)))

        # we have a filter condition
        if isexpr(x.args[2].args[2],:filter)
            cond = x.args[2].args[2].args[1]
            # generate inner loop code first and then wrap in for loops
            innercode = parseNLExpr(m, x.args[2].args[1], tapevar, parentvar, values)
            code = quote
                if $(esc(cond))
                    $innercode
                end
            end
            for level in length(x.args[2].args[2]):-1:2
                _idxvar, idxset = parseIdxSet(x.args[2].args[2].args[level]::Expr)
                idxvar = esc(_idxvar)
                code = :(let
                    $(localvar(idxvar))
                    for $idxvar in $(esc(idxset))
                        $code
                    end
                end)
            end
            push!(block.args, code)
        else # no condition
            innercode = parseNLExpr(m, x.args[2].args[1], tapevar, parentvar, values)
            code = quote
                $innercode
            end
            for level in length(x.args[2].args):-1:2
                _idxvar, idxset = parseIdxSet(x.args[2].args[level]::Expr)
                idxvar = esc(_idxvar)
                code = :(let
                    $(localvar(idxvar))
                    for $idxvar in $(esc(idxset))
                        $code
                    end
                end)
            end
            push!(block.args, code)
        end
        return codeblock
    end
    if isexpr(x,:call) && length(x.args) >= 2 && isexpr(x.args[2],:flatten)
        flatten_error(x.args[2])
    end

    if isexpr(x, :call)
        if length(x.args) == 2 # univariate
            code = :(let; end)
            block = code.args[1]
            @assert isexpr(block, :block)
            if haskey(univariate_operator_to_id,x.args[1])
                operatorid = univariate_operator_to_id[x.args[1]]
                push!(block.args, :(push!($tapevar, NodeData(CALLUNIVAR, $operatorid, $parent))))
            else
                opname = quot(x.args[1])
                errorstring = "Unrecognized function $opname used in nonlinear expression."
                errorstring2 = "Incorrect number of arguments for $opname in nonlinear expression."
                lookupcode = quote
                    if !haskey(univariate_operator_to_id,$opname)
                        if haskey(operator_to_id,$opname)
                            error($errorstring2)
                        else
                            error($errorstring)
                        end
                    end
                    operatorid = univariate_operator_to_id[$opname]
                end
                push!(block.args, :($lookupcode; push!($tapevar, NodeData(CALLUNIVAR, operatorid, $parent))))
            end
            parentvar = gensym()
            push!(block.args, :($parentvar = length($tapevar)))
            push!(block.args, parseNLExpr(m, x.args[2], tapevar, parentvar, values))
            return code
        else
            code = :(let; end)
            block = code.args[1]
            @assert isexpr(block, :block)
            if haskey(operator_to_id,x.args[1]) # fast compile-time lookup
                operatorid = operator_to_id[x.args[1]]
                push!(block.args, :(push!($tapevar, NodeData(CALL, $operatorid, $parent))))
            elseif haskey(comparison_operator_to_id,x.args[1])
                operatorid = comparison_operator_to_id[x.args[1]]
                push!(block.args, :(push!($tapevar, NodeData(COMPARISON, $operatorid, $parent))))
            else # could be user defined
                opname = quot(x.args[1])
                errorstring = "Unrecognized function $opname used in nonlinear expression."
                errorstring2 = "Incorrect number of arguments for $opname in nonlinear expression."
                lookupcode = quote
                    if !haskey(operator_to_id,$opname)
                        if haskey(univariate_operator_to_id,$opname)
                            error($errorstring2)
                        else
                            error($errorstring)
                        end
                    end
                    operatorid = operator_to_id[$opname]
                end
                push!(block.args, :($lookupcode; push!($tapevar, NodeData(CALL, operatorid, $parent))))
            end
            parentvar = gensym()
            push!(block.args, :($parentvar = length($tapevar)))
            for i in 1:length(x.args)-1
                push!(block.args, parseNLExpr(m, x.args[i+1], tapevar, parentvar, values))
            end
            return code
        end
    end
    if isexpr(x, :comparison)
        code = :(let; end)
        block = code.args[1]
        op = x.args[2]
        operatorid = comparison_operator_to_id[op]
        for k in 2:2:length(x.args)-1
            @assert x.args[k] == op # don't handle a <= b >= c
        end
        parentvar = gensym()
        push!(block.args, :(push!($tapevar, NodeData(COMPARISON, $operatorid, $parent))))
        push!(block.args, :($parentvar = length($tapevar)))
        for k in 1:2:length(x.args)
            push!(block.args, parseNLExpr(m, x.args[k], tapevar, parentvar, values))
        end
        return code
    end
    if isexpr(x, :&&) || isexpr(x, :||)
        code = :(let; end)
        block = code.args[1]
        op = x.head
        operatorid = logic_operator_to_id[op]
        parentvar = gensym()
        push!(block.args, :(push!($tapevar, NodeData(LOGIC, $operatorid, $parent))))
        push!(block.args, :($parentvar = length($tapevar)))
        push!(block.args, parseNLExpr(m, x.args[1], tapevar, parentvar, values))
        push!(block.args, parseNLExpr(m, x.args[2], tapevar, parentvar, values))
        return code
    end
    if isexpr(x, :curly)
        header = x.args[1]
        if length(x.args) < 3
            error("Need at least two arguments for $header")
        end
        if issum(header)
            operatorid = operator_to_id[:+]
        elseif isprod(header)
            operatorid = operator_to_id[:*]
        else
            error("Unrecognized expression $header{...}")
        end
        codeblock = :(let; end)
        block = codeblock.args[1]
        push!(block.args, :(push!($tapevar, NodeData(CALL, $operatorid, $parent))))
        parentvar = gensym()
        push!(block.args, :($parentvar = length($tapevar)))

        # we have a filter condition
        if isexpr(x.args[2],:parameters)
            cond = x.args[2]
            if length(cond.args) != 1
                error("No commas after semicolon allowed in $header{} expression, use && for multiple conditions")
            end
            # generate inner loop code first and then wrap in for loops
            innercode = parseNLExpr(m, x.args[3], tapevar, parentvar, values)
            code = quote
                if $(esc(cond.args[1]))
                    $innercode
                end
            end
            for level in length(x.args):-1:4
                _idxvar, idxset = parseIdxSet(x.args[level]::Expr)
                idxvar = esc(_idxvar)
                code = :(let
                    $(localvar(idxvar))
                    for $idxvar in $(esc(idxset))
                        $code
                    end
                end)
            end
            push!(block.args, code)
        else # no condition
            innercode = parseNLExpr(m, x.args[2], tapevar, parentvar, values)
            code = quote
                $innercode
            end
            for level in length(x.args):-1:3
                _idxvar, idxset = parseIdxSet(x.args[level]::Expr)
                idxvar = esc(_idxvar)
                code = :(let
                    $(localvar(idxvar))
                    for $idxvar in $(esc(idxset))
                        $code
                    end
                end)
            end
            push!(block.args, code)
        end
        return codeblock
    end
    # at the lowest level?
    return :( parseNLExpr_runtime($(esc(m)),$(esc(x)), $tapevar, $parent, $values) )

end

function parseNLExpr_runtime(m::Model, x::Number, tape, parent, values)
    push!(values, x)
    push!(tape, NodeData(VALUE, length(values), parent))
    nothing
end

# Temporary hack for deprecation of @defNLExpr syntax
const __last_model = Array(Model,1)

function parseNLExpr_runtime(m::Model, x::Variable, tape, parent, values)
    __last_model[1] = x.m
    x.m === m || error("Variable in nonlinear expression does not belong to corresponding model")
    push!(tape, NodeData(VARIABLE, x.col, parent))
    nothing
end

function parseNLExpr_runtime(m::Model, x::NonlinearExpression, tape, parent, values)
    push!(tape, NodeData(SUBEXPRESSION, x.index, parent))
    nothing
end

function parseNLExpr_runtime(m::Model, x::NonlinearParameter, tape, parent, values)
    push!(tape, NodeData(PARAMETER, x.index, parent))
    nothing
end

function parseNLExpr_runtime(m::Model, x::Vector, tape, parent, values)
    error("Unexpected vector $x in nonlinear expression. Nonlinear expressions may contain only scalar expressions.")
end

macro processNLExpr(m, ex)
    parsed = parseNLExpr(m, ex, :tape, -1, :values)
    quote
        tape = NodeData[]
        values = Float64[]
        $parsed
        NonlinearExprData(tape, values)
    end
end

# Construct a NonlinearExprData from a Julia expression.
# Variable objects should be spliced into the expression.
function NonlinearExprData(m::Model, ex::Expr)
    ex = spliceref(m,ex)
    nd, values = ReverseDiffSparse.expr_to_nodedata(ex)
    return NonlinearExprData(nd, values)
end
NonlinearExprData(m::Model, ex) = NonlinearExprData(m, :($ex + 0))

# recursively replace Variable(m, i) with Expr(:ref,:x,i) in ex
spliceref(m::Model, ex::Expr) = Expr(ex.head,map(e -> spliceref(m,e), ex.args)...)
function spliceref(m::Model, v::Variable)
    v.m === m || error("Variable $v does not belong to this model")
    return Expr(:ref, :x, linearindex(v))
end
spliceref(m::Model, ex) = ex


text:
struct Δprod{p}
  w::Dict{SVector{p, Int},Float64}
end
function Δprod(::Type{Val{p}}) where p
  Δprod{p}(Dict{SVector{p, Int}, Float64}())
end
function Base.:+(x::Δprod{p}, y::Δprod{p}) where p
  Δprod(x.w + y.w)
end
function Base.:+(x::Δprod{p}...) where p
  Δprod(sum(x.w))
end
function setindex!(A::Δprod{p}, x::Float64, i::Int64...) where {p}
  A.w[SVector{p,Int64}(i)] = x
end
function getindex(A::Δprod{p}, i::SVector{p,Int64}) where p
  A.w[i]
end
Base.keys(A::Δprod) = keys(A.w)
Base.values(A::Δprod) = values(A.w)
Base.length(A::Δprod) = length(A.w)

struct NestedGrid{p, q <: NestedQuadratureRule}
  Qs::Dict{Int64, q}
  Δs::Dict{Int64, Δ{q}}
  Δ_prods::Dict{SVector{p, Int64}, Δprod{p}}
  node_value::Dict{Int64,Float64}
  seq::Array{Int64,1}
  grid::Δprod{p}
end
function NestedGrid(::Type{Val{p}}, ::Type{q}) where {p, q <: NestedQuadratureRule}
  NestedGrid{p, q}(Dict{Int64,q}(), Dict{Int64,Δ{q}}(), Dict{SVector{p, Int64}, Δprod}(), Dict{Int64,Float64}(), default(q), Δprod(Val{p}))
end
function NestedGrid(::Type{Val{p}}, ::Type{q}, seq::Vector{Int}) where {p, q <: NestedQuadratureRule}
  NestedGrid{p, q}(Dict{Int64,q}(), Dict{Int64,Δ{q}}(), Dict{SVector{p, Int64}, Δprod}(), Dict{Int64,Float64}(), seq, Δprod(Val{p}))
end

#Type unstable calls.
NestedGrid(p::Int, ::Type{q}) where q = NestedGrid(Val{p}, q)
NestedGrid(p::Int, ::Type{q}, seq::Array{Int64,1}) where q = NestedGrid(Val{p}, q, seq)

function get_node(Grid::NestedGrid, i::Int)
  Grid.node_value[i]
end
function calc_Q!(Grid::NestedGrid{p, q}, i::Int) where {p, q <: NestedQuadratureRule}
  new_q = q(Grid.seq[i])
  if length(new_q.n) > length(Grid.node_value)
    merge!(Grid.node_value, new_q.n)
  end
  new_q
end
function get_Q!(Grid::NestedGrid{p, q}, i::Int) where {p, q <: NestedQuadratureRule}
  get!(() -> calc_Q!(Grid, i), Grid.Qs, i)
end
function calc_Δ!(Grid::NestedGrid{p,q}, i::Int) where {p, q <: NestedQuadratureRule}
  if i > 1
    if Grid.seq[i] != Grid.seq[i-1]
      return get_Q!(Grid, i) - get_Q!(Grid, i - 1)
    else
      return Δ{q}(Dict{Int64,Float64}(), Dict{Int64,Float64}())
    end
  elseif i == 1
    return Δ{q}(get_Q!(Grid, i).n, get_Q!(Grid, i).w)
  end
end
function get_Δ!(Grid::NestedGrid{p,q}, i::Int) where {p, q <: NestedQuadratureRule}
  get!(() -> calc_Δ!(Grid, i), Grid.Δs, i)
end
function get_Δ_weight!(Grid::NestedGrid, i::Int, j::Int)
  get_Δ!(Grid, i).w[j]
end
function get_Δ_weight!(Grid::NestedGrid, sv::SVector{p,Int64}, tup::Tuple) where {p}
  out = 1.0
  for i ∈ 1:p
    out *= get_Δ!(Grid, sv[i]).w[tup[i]]
  end
  out
end
#get_Δprod takes a vector of rule-sizes
#Δprods themselves are dictionaries with keys that have ints for node location.
function get_Δ_prod!(Grid::NestedGrid{p,q}, i::SVector{p, Int64}) where {p,q<:NestedQuadratureRule}
  get!(() -> calc_Δ_prod!(Grid, i), Grid.Δ_prods, i)
end
function expand_grid!(Grid::NestedGrid{p,q}, Δ_prod::Δprod{p}) where {p,q}
  for (i, j) ∈ Δ_prod.w
    Grid.grid.w[i] = get(Grid.grid.w, i, 0.0) + j
  end
end
function expand_grid!(Grid::NestedGrid{p,q}, j::SVector{p, Int64}) where {p,q}
  for (i, k) ∈ get_Δ_prod!(Grid, j).w
    Grid.grid.w[i] = get(Grid.grid.w, i, 0.0) + k
  end
end

@generated function calc_Δ_prod!(Grid::NestedGrid{p,q}, arg_indices::SVector{p,Int64}) where {p,q<:NestedQuadratureRule}
  quote
    out = Δprod(Val{$p})
    @nloops $p i dim -> begin
      keys(get_Δ!(Grid, arg_indices[dim]).w)
    end begin
      (@nref $p out i) = get_Δ_weight!( Grid, arg_indices, (@ntuple $p i) )
    end
    expand_grid!(Grid, out)
    out
  end
end
@generated function calc_Δ_prod(Grid::NestedGrid{p,q}, arg_indices::SVector{p,Int64}) where {p,q<:NestedQuadratureRule}
  quote
    out = Δprod(Val{$p})
    @nloops $p i dim -> begin
      keys(get_Δ!(Grid, arg_indices[dim]).w)
    end begin
      (@nref $p out i) = get_Δ_weight!( Grid, arg_indices, (@ntuple $p i) )
    end
    out
  end
end

function Base.length(Grid::NestedGrid)
  length(Grid.grid)
end


text:
module SliderFloats

Base.Experimental.@optlevel 1

using Redux
using CImGui
import ..ReduxImGui: get_label, get_values, get_value

# actions
abstract type AbstractSliderFloatAction <: AbstractSyncAction end

"""
    Rename(label, new_label)
Change label to `new_label`.
Note that renaming may also change the identifier, please refer to `help?> CImGui.PushID`.
"""
struct Rename <: AbstractSliderFloatAction
    label::String
    new_label::String
end

"""
    ChangeSize(label, size)
Change widget size to `size`.
Note changing size also changes widget appearance. If the old size is bigger,
those unused states are left intact.
"""
struct ChangeSize <: AbstractSliderFloatAction
    label::String
    size::Int
end

"""
    ChangeRange(label, range)
Change widget range to `range`.
"""
struct ChangeRange <: AbstractSliderFloatAction
    label::String
    range::Tuple{Cfloat,Cfloat}
end

"""
    ChangeRangeMaxTo(label, max)
Change widget max range to `max`.
"""
struct ChangeRangeMaxTo <: AbstractSliderFloatAction
    label::String
    max::Cfloat
end

"""
    ChangeRangeMinTo(label, min)
Change widget min range to `min`.
"""
struct ChangeRangeMinTo <: AbstractSliderFloatAction
    label::String
    min::Cfloat
end

"""
    ChangeFormat(label, format)
Change value's format to `format`, e.g. "%d".
"""
struct ChangeFormat <: AbstractSliderFloatAction
    label::String
    format::String
end

"""
    ChangePower(label, pow)
Change power to `pow`, e.g. "%d".
"""
struct ChangePower <: AbstractSliderFloatAction
    label::String
    power::Cfloat
end

# state
struct State <: AbstractImmutableState
    label::String
    vals::Vector{Cfloat}
    size::Int
    range::Tuple{Cfloat,Cfloat}
    format::String
    power::Cfloat
end
State(
    label::AbstractString;
    vals::Vector{Cfloat} = Cfloat[0, 0, 0, 0],
    size = 1,
    range = (0.0f0,0.0f0),
    format = "%.3f",
    power = 1.0f0,
) = State(label, vals, size, range, format, power)

State(label::AbstractString, val::AbstractFloat; args...) =
    State(label; vals = Cfloat[val, 0, 0, 0], args...)

# reducers
slider_float(state::AbstractState, action::AbstractAction) = state
slider_float(state::Vector{<:AbstractState}, action::AbstractAction) = state
slider_float(state::Dict{String,<:AbstractState}, action::AbstractAction) = state

function slider_float(state::Dict{String,State}, action::AbstractSliderFloatAction)
    s = Dict{String,State}()
    for (k, v) in state
        s[k] = get_label(v) == action.label ? slider_float(v, action) : v
    end
    return s
end

slider_float(s::State, a::Rename) = State(a.new_label, s.vals, s.size, s.range, s.format, s.power)
slider_float(s::State, a::ChangeSize) = State(s.label, s.vals, a.size, s.range, s.format, s.power)
slider_float(s::State, a::ChangeRange) = State(s.label, s.vals, s.size, a.range, s.format, s.power)
slider_float(s::State, a::ChangeRangeMaxTo) = State(s.label, s.vals, s.size, (s.range[1], a.max), s.format, s.power)
slider_float(s::State, a::ChangeRangeMinTo) = State(s.label, s.vals, s.size, (a.min, s.range[2]), s.format, s.power)
slider_float(s::State, a::ChangeFormat) = State(s.label, s.vals, s.size, s.range, a.format, s.power)
slider_float(s::State, a::ChangePower) = State(s.label, s.vals, s.size, s.range, s.format, a.power)

# helper
"""
    SliderFloat(store::AbstractStore, get_state=Redux.get_state) -> Bool
Return `true` when triggered.
"""
function SliderFloat(store::AbstractStore, get_state = Redux.get_state)
    s = get_state(store)
    min, max = s.range
    if s.size == 1
        is_triggered = CImGui.SliderFloat(get_label(s), s.vals, min, max, s.format, s.power)
    elseif s.size == 2
        is_triggered = CImGui.SliderFloat2(get_label(s), s.vals, min, max, s.format, s.power)
    elseif s.size == 3
        is_triggered = CImGui.SliderFloat3(get_label(s), s.vals, min, max, s.format, s.power)
    elseif s.size == 4
        is_triggered = CImGui.SliderFloat4(get_label(s), s.vals, min, max, s.format, s.power)
    else
        is_triggered = false
    end
    return is_triggered
end

get_label(s::State) = s.label
get_values(s::State) = s.vals
get_value(s::State) = first(s.vals)

end # module


text:
"""

`draw_grid(video::Video, object::AbstractObject, frame::Int; direction::AbstractString = "TR", line_gap = 25)`

Draws an oriented grid on the given frame of a Video.

# Arguments
- `direction::AbstractString`: Where grid animation finishes. Default: `"TR"` Available Orientations:
  - `"TR"` - Animation finishes in the **T**op **R**ight corner of the frame.
  - `"TL"` - Animation finishes in the **T**op **L**eft corner of the frame.
  - `"BR"` - Animation finishes in the **B**ottom **R**ight corner of the frame.
  - `"BL"` - Animation finishes in the **B**ottom **L**eft corner of the frame.
- `line_gap`: How many pixels between each line. Default: `25`

# Example
Example call of this function within an `Object`.
```
...
 Object(1:100, :line, draw_grid(direction = "TL", line_gap = 25))
...
```

"""
function draw_grid(; direction::AbstractString = "TR", line_gap = 25)
    return (video, object, frame) ->
        _draw_grid(video, object, frame; direction = direction, line_gap = line_gap)
end

function _draw_grid(
    video::Video,
    object::AbstractObject,
    frame::Int;
    direction::AbstractString = "TR",
    line_gap = 25,
)

    min_width = div(video.width, -2, RoundDown)
    max_width = div(video.width, 2, RoundUp)

    min_height = div(video.height, -2, RoundDown)
    max_height = div(video.height, 2, RoundUp)

    # This determines how quickly the animation is drawn
    step = (frame - first(get_frames(object))) / (length(get_frames(object)) - 1)

    if direction[1] == 'T'
        # Bottom to top for vertical grid lines
        for x_point in min_width:line_gap:max_width
            start_point = Point(x_point, max_height)
            finish_point = Point(x_point, min_height)
            end_point = start_point + step * (finish_point - start_point)
            line(start_point, end_point, :stroke)
        end
    else
        # Top to bottom motion for vertical grid lines
        for x_point in min_width:line_gap:max_width
            start_point = Point(x_point, min_height)
            finish_point = Point(x_point, max_height)
            end_point = start_point + step * (finish_point - start_point)
            line(start_point, end_point, :stroke)
        end
    end

    if direction[2] == 'R'
        # Left to right motion for horizontal grid lines
        for y_point in min_height:line_gap:max_height
            start_point = Point(min_width, y_point)
            finish_point = Point(max_width, y_point)
            end_point = start_point + step * (finish_point - start_point)
            line(start_point, end_point, :stroke)
        end
    else
        # Right to left motion for horizontal grid lines
        for y_point in min_height:line_gap:max_height
            start_point = Point(max_width, y_point)
            finish_point = Point(min_width, y_point)
            end_point = start_point + step * (finish_point - start_point)
            line(start_point, end_point, :stroke)
        end
    end

end

"""

`zero_lines(video::Video, object::AbstractObject, frame::Int; direction::AbstractString = "TR",
            line_thickness = 10)`

Draws zero lines on the given frame of a Video.

# Arguments
- `direction::AbstractString`: Direction for how vertical and horizontal axes are drawn.
Default: `"TR"` Available Orientations:
  - `"TR"` - Vertical axis drawn towards the **T**op and horizontal axis drawn
    to the **R**ight of the frame.
  - `"TL"` - Vertical axis drawn towards the **T**op and horizontal axis drawn
    to the **L**eft of the frame.
  - `"BR"` - Vertical axis drawn towards the **B**ottom and horizontal axis drawn
    to the **R**ight of the frame.
  - `"BL"` - Vertical axis drawn towards the **B**ottom and horizontal axis drawn
    to the **L**eft of the frame.
- `line_thickness`: Defines the thickness of the zero lines. Default: `10`

# Example
This example will produce an animation with the vertical axis being drawn towards the top
and the horizontal axis being drawn towards the left.
One will need to define their own path for `tempdirectory` and `pathname`.

```
...
 Object(1:100, :line, zero_lines(direction = "TL", line_thickness = 10)),
...
```

"""
function zero_lines(; direction::AbstractString = "TR", line_thickness = 10)
    return (video, object, frame) -> _zero_lines(
        video,
        object,
        frame;
        direction = direction,
        line_thickness = line_thickness,
    )
end

function _zero_lines(
    video::Video,
    object::AbstractObject,
    frame::Int;
    direction::AbstractString,
    line_thickness,
)

    min_width = div(video.width, -2, RoundDown)
    max_width = div(video.width, 2, RoundUp)

    min_height = div(video.height, -2, RoundDown)
    max_height = div(video.height, 2, RoundUp)

    setline(line_thickness)

    # This determines how quickly the animation is drawn
    step = (frame - first(get_frames(object))) / (length(get_frames(object)) - 1)

    if direction[1] == 'B'
        # Top to bottom motion for vertical line
        start_point_y = Point(0, min_height)
        end_point_y = start_point_y + step * (Point(0, max_height) - start_point_y)
    else
        # Bottom to top motion for vertical line
        start_point_y = Point(0, max_height)
        end_point_y = start_point_y + step * (Point(0, min_height) - start_point_y)
    end

    if direction[2] == 'R'
        # Left to right motion for horizontal line
        start_point_x = Point(min_width, 0)
        end_point_x = start_point_x + step * (Point(max_width, 0) - start_point_x)
    else
        # Right to left motion for horizontal line
        start_point_x = Point(max_width, 0)
        end_point_x = start_point_x + step * (Point(min_width, 0) - start_point_x)
    end

    line(start_point_x, end_point_x, :stroke)
    line(start_point_y, end_point_y, :stroke)

end

export draw_grid, zero_lines


text:
# wrappers of low-level functionality

function cufftGetProperty(property::libraryPropertyType)
  value_ref = Ref{Cint}()
  cufftGetProperty(property, value_ref)
  value_ref[]
end

version() = VersionNumber(cufftGetProperty(CUDA.MAJOR_VERSION),
                          cufftGetProperty(CUDA.MINOR_VERSION),
                          cufftGetProperty(CUDA.PATCH_LEVEL))


text:
@doc Markdown.doc"""
    rand!(a::nf_elem, U::AbstractArray) -> nf_elem
Inplace, set the coefficients of $a$ to random elements in $U$.
$a$ is returned.
"""
function rand!(a::nf_elem, U::AbstractArray)
  for i=1:degree(parent(a))
    _num_setcoeff!(a, i-1, rand(U))
  end
  return a
end

@doc Markdown.doc"""
    rand(K::AnticNumberField, U::AbstractArray) -> nf_elem
Find an element in $K$ where the coefficients are selceted at random in $U$.
"""
function rand(K::AnticNumberField, U::AbstractArray)
  a = K()
  return rand!(a, U)
end

@doc Markdown.doc"""
    rand!(A::Generic.Mat{nf_elem}, U::AbstractArray) -> Generic.Mat{nf_elem}
Inplace, replace each element in $A$ by an element where the coefficients are
sected at random in $U$.
Returns $A$.
"""
function rand!(A::Generic.Mat{nf_elem}, U::AbstractArray)
  for i=1:nrows(A)
    for j=1:ncols(A)
      rand!(A[i,j], U)
    end
  end
  return A
end

@doc Markdown.doc"""
    rand(A::Generic.MatSpace{nf_elem}, U::AbstractArray) -> Generic.Mat{nf_elem}
Create a random matrix in $A$ where the coefficients are selected from $U$.
"""
function rand(A::Generic.MatSpace{nf_elem}, U::AbstractArray)
  return rand!(A(), U)
end

@doc Markdown.doc"""
    modular_lift(ap::Array{fq_nmod_mat, 1}, me::modular_env) -> Array
Given an array of matrices as computed by \code{modular_proj},
compute a global pre-image using some efficient CRT.
"""
function modular_lift(ap::Array{fq_nmod_mat, 1}, me::modular_env)
  A = zero_matrix(me.K, nrows(ap[1]), ncols(ap[1]))
  for i=1:nrows(A)
    for j=1:ncols(A)
      A[i,j] = modular_lift([ap[k][i,j] for k=1:length(ap)], me)
    end
  end
  return A
end

@doc Markdown.doc"""
    mod!(A::Generic.Mat{nf_elem}, m::fmpz)
Inplace: reduce all entries of $A$ modulo $m$, into the positive residue system.
"""
function mod!(A::Generic.Mat{nf_elem}, m::fmpz)
  for i=1:nrows(A)
    for j=1:ncols(A)
      mod!(A[i, j], m)
    end
  end
end

@doc Markdown.doc"""
    mod_sym!!(A::Generic.Mat{nf_elem}, m::fmpz)
Inplace: reduce all entries of $A$ modulo $m$, into the symmetric residue system.
"""
function mod_sym!(A::Generic.Mat{nf_elem}, m::fmpz)
  for i = 1:nrows(A)
    for j = 1:ncols(A)
      mod_sym!(A[i, j], m)
    end
  end
end

function small_coeff(a::nf_elem, B::fmpz, i::Int)
  z = fmpz()
  Nemo.num_coeff!(z, a, i)
  return cmpabs(z, B) <= 0
end

@doc Markdown.doc"""
    rational_reconstruction(A::Generic.Mat{nf_elem}, M::fmpz) -> Bool, Generic.Mat{nf_elem}
Apply \code{rational_reconstruction} to each entry of $M$.
"""
function rational_reconstruction2(A::Generic.Mat{nf_elem}, M::fmpz)
  B = similar(A)
  sM = root(M, 2)
  d = one(A[1,1])
  di = inv(d)
  for i=1:nrows(A)
    for j=1:ncols(A)
      a = A[i,j]*d
      mod_sym!(a, M)
      if all(i->small_coeff(a, sM, i), 1:a.elem_length)
        B[i,j] = a*di
      else
        n, dn = algebraic_reconstruction(a, M)
        d*=dn
        if any(i->!small_coeff(d, sM, i), 1:a.elem_length)
          println("early $i $j abort")
          return false, B
        end
        di*=inv(dn)
        B[i,j] = n*di
      end
    end
  end
  println("final den: $d")
  return true, B//d
end

function rational_reconstruction(A::Generic.Mat{nf_elem}, M::fmpz)
  B = similar(A)
  for i=1:nrows(A)
    for j=1:ncols(A)
      fl, B[i,j] = rational_reconstruction(A[i,j], M)
      if !fl
        return false, B
      end
    end
  end
  return true, B
end

function algebraic_reconstruction(a::nf_elem, M::fmpz)
  K = parent(a)
  n = degree(K)
  Znn = MatrixSpace(FlintZZ, n, n, false)
#  L = [ Znn(1) representation_matrix_q(a)[1] ; Znn(0) Znn(M)]
  L = vcat(hcat(Znn(1), representation_matrix_q(a)[1]), hcat(Znn(0),Znn(M)))
  lll!(L)
  d = Nemo.elem_from_mat_row(K, sub(L, 1:1, 1:n), 1, fmpz(1))
  n = Nemo.elem_from_mat_row(K, sub(L, 1:1, n+1:2*n), 1, fmpz(1))
  return n,d
  return true, n//d
end

function algebraic_reconstruction(a::nf_elem, M::NfAbsOrdIdl)
  K = parent(a)
  n = degree(K)
  Znn = MatrixSpace(FlintZZ, n, n, false)
  L = [ Znn(1) representation_matrix_q(a)[1] ; Znn(0) basis_matrix(M, copy = false)]
  lll!(L)
  d = Nemo.elem_from_mat_row(K, sub(L, 1:1, 1:n), 1, fmpz(1))
  n = Nemo.elem_from_mat_row(K, sub(L, 1:1, n+1:2*n), 1, fmpz(1))
  return n,d
  return true, n//d
end

@doc Markdown.doc"""
    algebraic_split(a::nf_elem) -> nf_elem, nf_elem

Writes the input as a quotient of two "small" algebraic integers.
"""
function algebraic_split(a::nf_elem)
  n = degree(parent(a))
  d = denominator(a)
  M, dd = representation_matrix_q(a)
  @assert d == dd
  M = [M identity_matrix(FlintZZ, n); -d*identity_matrix(FlintZZ, n) zero_matrix(FlintZZ, n, n)]
  L = lll(M)[1, n+1:2*n]
  ga = parent(a)(collect(L))
  be = a*ga
  @assert denominator(be) == 1
  return be, ga
end

#function denominator_ideal(M::Generic.MatSpaceElem{nf_elem}, den::nf_elem)
function denominator_ideal(M::Array{nf_elem, 1}, den::nf_elem)
  k = parent(M[1,1])
  zk = maximal_order(k)
  _, d = integral_split(M[1]//den * zk)
  g = simplify(den*zk * inv(d)).num ## should be small
  if isone(g)
    return d
  end

  for m = M
    i = simplify(ideal(zk, minimum(g)^2, zk(m)))
    j = simplify(i*inv(g))
    if denominator(j) == 1
      continue
    end
    _, dd = integral_split(m//den * zk)
    d = lcm(d, dd)
    g = simplify(simplify(den*zk * inv(d)).num)
    if isone(g)
      break
    end
  end
  return d
end

@doc Markdown.doc"""
    divexact!(A::Generic.Mat{nf_elem}, p::fmpz)
Inplace: divide each entry by $p$.
"""
function divexact!(A::Generic.Mat{nf_elem}, p::fmpz)
  for i=1:nrows(A)
    for j=1:ncols(A)
      A[i,j] = A[i,j]//p
    end
  end
end

#TODO/ To experiment:
# - vector reconstruction ala Storjohan
# - reconstruction with algebraic denominators
# - vector reconstruction with algebraic denominators
# - Talk to Bill: fq_nmod_mat is missing in Nemo, the inv is dreadfully slow
# - extend to non-unique solutions
# - make Aip*D mult faster, A*y as well?
#
function solve_dixon(A::Generic.Mat{nf_elem}, B::Generic.Mat{nf_elem})
  p = next_prime(p_start)
  K = base_ring(A)

  me = modular_init(K, p)
  ap = modular_proj(A, me)
  @time ap = [inv(x) for x= ap]
  Aip = modular_lift(ap, me)
  sol = 0*B
  D = B
  pp = fmpz(1)
  last_SOL = false
  nd = 0
  while true
    nd += 1
    y = Aip*D
    mod!(y, fmpz(p))
    sol += y*pp
    pp *= p
    fl, SOL = rational_reconstruction(sol, pp)
#    t = A*sol-B
#    mod!(t, pp)
#    @assert iszero(t)
    if fl
      if last_SOL == SOL && A*SOL == B
        #println("used $nd $p-adic digits")
        return SOL
      else
        last_SOL = SOL
      end
    end
    D = D - A*y
    divexact!(D, fmpz(p))
#    if nbits(pp) > 10000 # a safety device to avoid infinite loops
#      error("not work")
#    end
  end
end


text:
gw = GridWorld(sx=2, sy=2, rs=[GridWorldState(1,1)], rv=[10.0])

pvec = fill(GridWorldAction(:left), 5)

solver = VectorSolver(pvec)

p = solve(solver, gw)

for s in iterator(states(gw))
    @test action(p, s) == GridWorldAction(:left)
end

p2 = VectorPolicy(gw, pvec)
for s in iterator(states(gw))
    @test action(p2, s) == GridWorldAction(:left)
end


text:
# a blue/green/yellow theme of no specific origin
const lime_palette = reverse([colorant"#271924", # dark blue
                      colorant"#394256", # semi dark blue
                      colorant"#30727F", # green blue
                      colorant"#36A58F", # turqoise
                      colorant"#80D584", # green,
                      colorant"#EBFB73"]) # yellow

black = lime_palette[6]

_themes[:lime] = PlotTheme(
    bg = black,
    bginside = black,
    fg = lime_palette[1],
    fgtext = lime_palette[2],
    fgguide = lime_palette[2],
    fglegend = lime_palette[2],
    palette = expand_palette(black, lime_palette[1:4]),
    gradient = cgrad(:viridis).colors
)


text:
using BenchmarkTools
using Transducers

function manual_filter_map_reduce(xs)
    acc = zero(eltype(xs))
    for x in xs
        if -0.5 < x < 0.5
            acc += 2x
        end
    end
    return acc
end

let suite = BenchmarkGroup()

    xf = Filter(x -> -0.5 < x < 0.5) |> Map(x -> 2x)
    n = 10^5

    let xs = randn(10)
        @assert manual_filter_map_reduce(xs) == mapfoldl(xf, +, 0.0, xs)
    end

    suite["xf"] = @benchmarkable(
        mapfoldl($xf, +, 0.0, xs),
        setup=(xs = randn($n)))
    suite["man"] = @benchmarkable(
        manual_filter_map_reduce(xs),
        setup=(xs = randn($n)))

    suite
end


text:
cd(@__DIR__); include("setups/rgc100.jl")
gr(dpi = 200)

plt = plot(0:(N - 1), 𝛌, c = :black, lw = 1, ylims = [0, 5], legend = false,
        shape = :circle, frame = :box, xlab = latexstring("l"),
        ylab = latexstring("\\tilde{\\lambda}_l"), size = (500, 400))
savefig(plt, "../figs/RGC100_eigenvalues.png")


text:
const DO_NOT_PUSH_TO_THESE_DESTINATIONS = String[
    "MirrorUpdater",
    "MirrorUpdater.jl",
    "UnofficialJuliaMirror/MirrorUpdater",
    "UnofficialJuliaMirror/MirrorUpdater.jl",
    "UnofficialJuliaMirrorBot/MirrorUpdater",
    "UnofficialJuliaMirrorBot/MirrorUpdater.jl",
    ]


text:
module NeXLParticle

using DataFrames
using FileIO
using Reexport
using Requires
using StatsBase
using AxisArrays

@reexport using NeXLSpectrum

include("zeppelin.jl")
export Zeppelin # struct holding a hdz/pxz file pair
export classes # An ordered list of particle classes
export header # A alphabetically sorted list of header items
export data # Provides access to the DataFrame containing the particle data
export iszeppelin # Is this IOStream or filename a Zeppelin file?
export spectrum # read a spectrum associate with a Zeppelin file (or missing)
export eachparticle # The range of partice row indices
export rowsmax, rowsmin, rowsclass
export MORPH_COLS, COMP_COLS, CLASS_COLS
export beamenergy, probecurrent, magdata
export maxparticle
export writeZep

export ParticleClassifier

include("signature.jl")
export Signature
export signature # Compures the particle signature
export quantify # Quantifies the particle spectrum data and constructs a fresh Zeppelin with the results
export NSigmaCulling
export NoCulling

include("classrule.jl")
export classify
export OrderedRuleSet
export GSRRules
export BaseRules
export NullRules

include("multiternary.jl")
export multiternary
export TernPalette
export TernColorblind

include("blob.jl")
export Blob
export blob
export separate # Split agglomerated features in a Blob in many Blob(s)
export multiseparate # Multithreshold algorithm to separate agglomerated blobs
export curvature
export perimeter
export crosscorr
export area
export perimeterlength
export colorizedimage
export maskedimage
export perimeterlength
export ecd
export scorer
export interiorregions

include("rca.jl")
export rca
export metrics

include("watershed.jl")
export soille_watershed

include("particlesegregation.jl")
export Watershed

include("diluvian.jl")
export DiluvianCluster # Constructs the clusters from the data
export asmask # As a binary mask
export membercounts # Number of items in each cluster
export summarizeclusters # Statistics summarizing the clusters as a DataFrame
export clusterstats # Statistics summarizing one or more clusters
export clusters # The raw cluster assignments
export asimage
export defaultpalette

include("align.jl")
export rough_align
export AlignIntermediary
export offset, score, findbest, align

include("translate.jl")
export XYPosition
export translate

function __init__()
    @require Gadfly = "c91e804a-d5a3-530f-b6f0-dfbca275c004" include("gadflysupport.jl")
end

end


text:
"""
    coefficient_diff(β::AbstractVector, method=:concat)

Returns the coefficient difference statistic W[j] = |β[j]| - |β[j + p]|
from a univariate (single response) regression.

# Inputs
+ `β`: Vector of regression coefficients
+ `method`: Either `:concatenated` (default) if all knockoffs are concatenated
    at the end (e.g. [XX̃]) or `:interleaved` if each variable is immediately
    followed by its knockoff (e.g. [x₁x̃₁x₂x̃₂...])
"""
function coefficient_diff(β::AbstractVector, method::Symbol=:concatenated)
    iseven(length(β)) || error("length of β should be even but was odd.")
    p = length(β) >> 1
    W = Vector{eltype(β)}(undef, p)
    if method == :concatenated
        for j in 1:p
            W[j] = abs(β[j]) - abs(β[j + p])
        end
    elseif method == :interleaved
        for j in 1:p
            W[j] = abs(β[2j - 1]) - abs(β[2j])
        end
    else
        error("method should be :concatenated or :interleaved but got $method")
    end
    return W
end

"""
    coefficient_diff(β::AbstractVector, original::AbstractVector{Int}, knockoff::AbstractVector{Int})

Returns the coefficient difference statistic W[j] = |β[j]| - |β[j + p]|
from a univariate (single response) regression, where the `j`th variable is stored
in position `original[j]` of `β`, and its knockoff is stored in position `knockoff[j]`

# Inputs
+ `β`: Vector of regression coefficients
+ `original`: The index of original variables in `β`
+ `knockoff`: The index of knockoff variables in `β`
"""
function coefficient_diff(β::AbstractVector, original::AbstractVector{Int}, knockoff::AbstractVector{Int})
    return abs.(β[original]) - abs.(β[knockoff])
end

"""
    coefficient_diff(β::AbstractVector, original::AbstractVector{Int}, knockoff::AbstractVector{Int})

Returns the coefficient difference statistic for grouped variables
W[G] = sum_{j in G} |β[j]| - sum_{j in G} |β[j + p]|.

# Inputs
+ `β`: Vector of regression coefficients
+ `groups`: Vector storing group membership. `groups[i]` is the group of `β[i]`
+ `original`: The index of original variables in `β`
+ `knockoff`: The index of knockoff variables in `β`
"""
function coefficient_diff(β::AbstractVector, groups::AbstractVector{Int},
    original::AbstractVector{Int}, knockoff::AbstractVector{Int})
    unique_groups = unique(groups)
    β_groups = zeros(length(unique_groups))
    # find which variables are Knockoffs
    knockoff_idx = falses(length(β))
    knockoff_idx[knockoff] .= true
    # loop over each variable
    for i in 1:length(β)
        g = groups[i]
        if knockoff_idx[i]
            β_groups[g] -= abs(β[i])
        else
            β_groups[g] += abs(β[i])
        end
    end
    return β_groups
end

"""
    coefficient_diff(B::AbstractMatrix)

Returns the coefficient difference statistic
W[j] = |B[j, 1]| - |B[j + p, 1]| + ... +  |B[j, r]| - |B[j + p, r]|
from a multivariate (multiple response) regression.

# Inputs
+ `β`: Matrix of regression coefficients. Each column is a vector of β.
+ `method`: Either `:concatenated` if all knockoffs are concatenated at the end
    (e.g. [XX̃]) or `:interleaved` if each variable is immediately followed by
    its knockoff (e.g. [x₁x̃₁x₂x̃₂...])
"""
function coefficient_diff(B::AbstractMatrix, method::Symbol=:concatenated)
    iseven(size(B, 1)) || error("Number of covariates in B should be even but was odd.")
    p = size(B, 1) >> 1
    r = size(B, 2)
    T = eltype(B)
    W = Vector{T}(undef, p)
    if method == :concatenated
        for j in 1:p
            Wj = zero(T)
            for i in 1:r
                Wj += abs(B[j, i]) - abs(B[j + p, i])
            end
            W[j] = Wj
        end
    elseif method == :interleaved
        for j in 1:p
            Wj = zero(T)
            for i in 1:r
                Wj += abs(B[2j - 1, i]) - abs(B[2j, i])
            end
            W[j] = Wj
        end
    else
        error("method should be :concatenated or :interleaved but got $method")
    end
    return W
end

# """
#     coefficient_diff(B::AbstractMatrix, original::AbstractVector{Int}, knockoff::AbstractVector{Int})

# Returns the coefficient difference statistic W[j] = |β[j]| - |β[j + p]|
# from a multivariate (multiple response) regression, where the `j`th variable is stored
# in position `original[j]` of `β`, and its knockoff is stored in position `knockoff[j]`

# # Inputs
# + `β`: Vector of regression coefficients
# + `original`: The index of original variables in `β`
# + `knockoff`: The index of knockoff variables in `β`
# """
# function coefficient_diff(B::AbstractMatrix, original::AbstractVector{Int}, knockoff::AbstractVector{Int})
#     p = size(B, 1) >> 1
#     length(original) == length(knockoff) == p || error("Number of variables in " *
#         "B should be twice the length of original and knockoff.")
#     W = Vector{eltype(B)}(undef, p)
#     for j in 1:p
#         Wj = zero(T)
#         for i in 1:r
#             Wj += abs(B[original[j], i]) - abs(B[knockoff[j], i])
#         end
#         W[j] = Wj
#     end
#     return W
# end

"""
    extract_beta(β̂_knockoff::AbstractVector, fdr::Number, method::Symbol=:concatenated)

Given estimated β of original variables and their knockoffs, compute β for the
original design matrix that controls the FDR.
"""
function extract_beta(β̂_knockoff::AbstractVector{T}, fdr::Number,
    method::Symbol=:concatenated
    ) where T <: AbstractFloat
    # first handle errors
    0 ≤ fdr ≤ 1 || error("Target FDR should be between 0 and 1 but got $fdr")
    p = length(β̂_knockoff) >> 1
    # length(original) == length(knockoff) == p || error("Length of " *
    #     "β should be twice of original and knockoff.")
    # find set of selected predictors
    W = coefficient_diff(β̂_knockoff, method)
    τ = threshold(W, fdr)
    detected = findall(W .≥ τ)
    # construct original β
    β = zeros(T, p)
    if method == :concatenated
        for i in detected
            β[i] = β̂_knockoff[i]
        end
    elseif method == :interleaved
        for i in detected
            β[i] = β̂_knockoff[2i - 1]
        end
    else
        error("method should be :concatenated or :interleaved but got $method")
    end
    return β
end

function extract_beta(β̂_knockoff::AbstractVector{T}, fdr::Number,
    original::AbstractVector{Int}, knockoff::AbstractVector{Int},
    method::Symbol=:knockoff, combine::Bool=false) where T <: AbstractFloat
    # first handle errors
    p = length(β̂_knockoff) >> 1
    0 ≤ fdr ≤ 1 || error("Target FDR should be between 0 and 1 but got $fdr")
    # length(original) == length(knockoff) == p || error("Length of " *
    #     "β should be twice of original and knockoff.")
    # find set of selected predictors
    W = coefficient_diff(β̂_knockoff, original, knockoff)
    τ = threshold(W, fdr, method)
    detected = findall(W .≥ τ)
    # construct final β
    β = zeros(T, p)
    if combine
        for i in detected
            β[i] = β̂_knockoff[original[i]] + β̂_knockoff[knockoff[i]]
        end
    else
        for i in detected
            β[i] = β̂_knockoff[original[i]]
        end
    end
    return β
end

function extract_beta(β̂_knockoff::AbstractVector{T}, fdr::Number, groups::Vector{Int},
    original::AbstractVector{Int}, knockoff::AbstractVector{Int}, method=:knockoff
    ) where T <: AbstractFloat
    # first handle errors
    p = length(β̂_knockoff) >> 1
    0 ≤ fdr ≤ 1 || error("Target FDR should be between 0 and 1 but got $fdr")
    # length(β̂_knockoff) == length(groups) ||
    #     error("β̂_knockoff should have same length as groups")
    # length(original) == length(knockoff) == p ||
    #     error("Length of β should be twice of original and knockoff.")
    # find set of selected predictors
    W = coefficient_diff(β̂_knockoff, groups, original, knockoff)
    τ = threshold(W, fdr, method)
    detected_groups = findall(W .≥ τ)
    # construct the full β
    β = zeros(T, length(β̂_knockoff))
    for g in detected_groups
        group_idx = findall(x -> x == g, groups)
        β[group_idx] .= @view(β̂_knockoff[group_idx])
    end
    return β[original]
end


text:
export
    ishexagonal,
    nthhexagonal, nhexagonal,
    allhexagonal, somehexagonal, exacthexagonal

ishexagonal(n::Integer) = begin
    t = 8*n + 1
    r = isqrt(t)
    (r^2 == t) && (r % 4 == 3)
end

nthhexagonal(n::Integer) = n * (2*n - 1)
nhexagonal(n::Int, T::Type = Int) = collect(exacthexagonal(n, T))

allhexagonal(T::Type = Int) = HexagonalIterator{T}()
somehexagonal(xmax::T) where {T<:Integer} = @pipe allhexagonal(T) |> takewhile(x -> x <= xmax, _)
exacthexagonal(n::Int, T::Type = Int) = @pipe allhexagonal(T) |> take(_, n)

struct HexagonalIterator{T<:Integer}
end

Base.start(::HexagonalIterator{T}) where {T<:Integer} = zero(T), one(T)
Base.next(::HexagonalIterator{T}, state) where {T<:Integer} = begin
    s, n = state
    s += 4*n - 3
    s, (s, n + 1)
end
Base.done(::HexagonalIterator, _) = false

Base.eltype(it::HexagonalIterator) = Base.eltype(typeof(it))
Base.eltype(::Type{HexagonalIterator{T}}) where {T<:Integer} = T

Base.iteratorsize(::HexagonalIterator) = Base.IsInfinite()


text:
function check(x::Matrix{T} where T <: Real)
    r, c = size(x)
    if c > r
        @warn "the column size is greater than row size, you may need to transform your data"
    end
end


text:
# Structs for stress and velocity fields (used for multiple dispatch)

abstract type Fields end

"""
Return a list of fields which contain `s`.
Returns all the elastic fields defined here by default i.e., for `ndims=3`.
"""
function Fields(s = ""; ndims = _fd.ndims)
    @assert(ndims ∈ [2, 3])
    f = filter(x -> contains(string(x), string(s)), nameof.(subtypes(Fields)))
    if (ndims == 2)
        f = filter(x -> !contains(string(x), "y"), f) # filter out y when 2D
    end
    return f
end
function Fields(::T, s = ""; ndims = _fd.ndims) where {T<:Elastic}
    f = Fields(s, ndims = ndims) # get all
    f = filter(x -> !contains(string(x), "p"), f) # filter out pressure
end
function Fields(::T, s = ""; ndims = _fd.ndims) where {T<:Acoustic}
    f = Fields(s, ndims = ndims) # get all
    f = filter(x -> !contains(string(x), "tau"), f) # filter out tau
    f = filter(x -> !in(x, [:dvxdy, :dvxdz, :dvydx, :dvydz, :dvzdx, :dvzdy]), f) # filter out dvxdy
end

# define pressure for acoustic
struct p <: Fields end
# define stress and velocity (tauxx, tauxy, tauxz,)
for f in vcat(dim_names(3, "v"), dim_names(3, "tau"; order = 2))
    @eval struct $f <: Fields end
end
# define spatial derivatives of stress, velocity, pressure
dnames = dim_names(3, "d") # derivatives
vnames = dim_names(3, "v") # velocity
taunames = dim_names(3, "tau"; order = 2)
for (i, x) in enumerate(dim_names(3))
    tn = filter(y -> contains(string(y), string(x)), taunames)
    for t in vcat(tn, vnames, [:p])
        @eval struct $(Symbol("d", string(t), "d", string(x))) <: Fields end
    end
end

#-------------------------------------------------------
# grid indices and linear-interpolation weights
#-------------------------------------------------------

for dimnames in [zip([:1, :2, :3], [:z, :y, :x]), zip([:1, :2], [:z, :x])]
    grids = broadcast(x -> Symbol(string("m", x)), getindex.(collect(dimnames), 2))
    sizes1 = broadcast(x -> Symbol(string("n1", x)), getindex.(collect(dimnames), 2))
    sizes = broadcast(x -> Symbol(string("n", x)), getindex.(collect(dimnames), 2))
    # initialize fields by simply getting grids
    for f in Fields(ndims = length(collect(dimnames)))
        n1pml = Symbol("n1", string(last(string(f))))
        @eval function Base.zeros(
            ::$f,
            attrib_mod,
            $(sizes...);
            pml::Bool = false,
        ) where {T<:Fields}
            # get new sizes
            $(
                (
                    quote
                        $n1 = length(
                            get_mgrid(
                                $f(),
                                attrib_mod,
                                $((
                                        quote
                                            range(0, length = $n, stop = 1)
                                        end for n in sizes
                                    )...),
                            )[$i],
                        )
                    end for (i, n1) in enumerate(sizes1)
                )...
            )
            # if pml flag size on corresponding dimension is just 2*npml
            if (pml)
                $n1pml = 2 * _fd.npml
            end
            return Data.Array(zeros($(sizes1...)))
        end
    end
end

#-------------------------------------------------------
# Acoustic
#-------------------------------------------------------
get_mgrid(::p, ::FdtdAcoustic, mz, mx) = [mz, mx]
get_mgrid(::p, ::FdtdAcoustic, mz, my, mx) = [mz, my, mx]
get_mgrid(::dpdx, ::FdtdAcoustic, mz, mx) = [
    range(
        mz[1] + step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - 2(_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]
get_mgrid(::dpdx, ::FdtdAcoustic, mz, my, mx) = [
    range(
        mz[1] + step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - 2(_fd.order - 1),
    ),
    range(
        my[1] + step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - 2(_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]
get_mgrid(::dpdz, ::FdtdAcoustic, mz, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        mx[1] + step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - 2(_fd.order - 1),
    ),
]
get_mgrid(::dpdz, ::FdtdAcoustic, mz, my, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        my[1] + step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - 2(_fd.order - 1),
    ),
    range(
        mx[1] + step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - 2(_fd.order - 1),
    ),
]
get_mgrid(::dpdy, ::FdtdAcoustic, mz, my, mx) = [
    range(
        mz[1] + step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - 2(_fd.order - 1),
    ),
    range(
        my[1] + 0.5 * step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - (_fd.order - 1),
    ),
    range(
        mx[1] + step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - 2(_fd.order - 1),
    ),
]
get_mgrid(::vx, ::FdtdAcoustic, mz, mx) = [
    mz,
    range(
        mx[1] - 0.5 *  step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) + (_fd.order - 1),
    ),
]
get_mgrid(::vz, ::FdtdAcoustic, mz, mx) = [
    range(
        mz[1] - 0.5 * step(mz) *(_fd.order - 1),
        step = step(mz),
        length = length(mz) + (_fd.order - 1),
    ),
    mx,
]
get_mgrid(::vx, ::FdtdAcoustic, mz, my, mx) = [
    mz,
    my,
    range(
        mx[1] - 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) + (_fd.order - 1),
    ),
]
get_mgrid(::vz, ::FdtdAcoustic, mz, my, mx) = [
    range(
        mz[1] - 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) + (_fd.order - 1),
    ),
    my,
    mx,
]
get_mgrid(::vy, ::FdtdAcoustic, mz, my, mx) = [
    mz,
    range(
        my[1] - 0.5 * step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) + (_fd.order - 1),
    ),
    mx,
]
get_mgrid(::dvxdx, ::FdtdAcoustic, mz, mx) = [mz, mx]
get_mgrid(::dvzdz, ::FdtdAcoustic, mz, mx) = [mz, mx]
get_mgrid(::dvxdx, ::FdtdAcoustic, mz, my, mx) = [mz, my, mx]
get_mgrid(::dvzdz, ::FdtdAcoustic, mz, my, mx) = [mz, my, mx]
get_mgrid(::dvydy, ::FdtdAcoustic, mz, my, mx) = [mz, my, mx]

#-------------------------------------------------------
# Elastic
#-------------------------------------------------------
get_mgrid(::tauxx, ::FdtdElastic, mz, my, mx) = [mz, my, mx]
get_mgrid(::tauyy, ::FdtdElastic, mz, my, mx) = [mz, my, mx]
get_mgrid(::tauzz, ::FdtdElastic, mz, my, mx) = [mz, my, mx]
get_mgrid(::tauxy, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - 2(_fd.order - 1),
    ),
    range(
        my[1] + 0.5 * step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - (_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]
get_mgrid(::tauxz, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        my[1] + step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - 2(_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]
get_mgrid(::tauyz, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        my[1] + 0.5 * step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - (_fd.order - 1),
    ),
    range(
        mx[1] + step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - 2(_fd.order - 1),
    ),
]

get_mgrid(::tauxx, ::FdtdElastic, mz, mx) = [mz, mx]
get_mgrid(::tauzz, ::FdtdElastic, mz, mx) = [mz, mx]
get_mgrid(::tauxz, ::FdtdElastic, mz, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]

get_mgrid(::vx, ::FdtdElastic, mz, my, mx) = [
    mz,
    my,
    range(
        mx[1] - 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) + (_fd.order - 1),
    ),
]
get_mgrid(::vy, ::FdtdElastic, mz, my, mx) = [
    mz,
    range(
        my[1] - 0.5 * step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) + (_fd.order - 1),
    ),
    mx,
]
get_mgrid(::vz, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] - 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) + (_fd.order - 1),
    ),
    my,
    mx,
]

get_mgrid(::vx, ::FdtdElastic, mz, mx) = [
    mz,
    range(
        mx[1] - 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) + (_fd.order - 1),
    ),
]
get_mgrid(::vz, ::FdtdElastic, mz, mx) = [
    range(
        mz[1] - 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) + (_fd.order - 1),
    ),
    mx,
]

get_mgrid(::dvxdx, ::FdtdElastic, mz, my, mx) = [mz, my, mx]
get_mgrid(::dvydy, ::FdtdElastic, mz, my, mx) = [mz, my, mx]
get_mgrid(::dvzdz, ::FdtdElastic, mz, my, mx) = [mz, my, mx]
get_mgrid(::dvxdy, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - 2(_fd.order - 1),
    ),
    range(
        my[1] + 0.5 * step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - (_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]
get_mgrid(::dvxdz, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        my[1] + step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - 2(_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]
get_mgrid(::dvydx, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - 2(_fd.order - 1),
    ),
    range(
        my[1] + 0.5 * step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - (_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]
get_mgrid(::dvydz, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        my[1] + 0.5 * step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - (_fd.order - 1),
    ),
    range(
        mx[1] + step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - 2(_fd.order - 1),
    ),
]
get_mgrid(::dvzdx, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        my[1] + step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - 2(_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]
get_mgrid(::dvzdy, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        my[1] + 0.5 * step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - (_fd.order - 1),
    ),
    range(
        mx[1] + step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - 2(_fd.order - 1),
    ),
]

get_mgrid(::dvxdx, ::FdtdElastic, mz, mx) = [mz, mx]
get_mgrid(::dvzdz, ::FdtdElastic, mz, mx) = [mz, mx]
get_mgrid(::dvxdz, ::FdtdElastic, mz, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]
get_mgrid(::dvzdx, ::FdtdElastic, mz, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]

get_mgrid(::dtauxxdx, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - 2(_fd.order - 1),
    ),
    range(
        my[1] + step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - 2(_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]
get_mgrid(::dtauyydy, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - 2(_fd.order - 1),
    ),
    range(
        my[1] + 0.5 * step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - (_fd.order - 1),
    ),
    range(
        mx[1] + step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - 2(_fd.order - 1),
    ),
]
get_mgrid(::dtauzzdz, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        my[1] + step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - 2(_fd.order - 1),
    ),
    range(
        mx[1] + step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - 2(_fd.order - 1),
    ),
]
get_mgrid(::dtauxydx, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - 2(_fd.order - 1),
    ),
    range(
        my[1] + 0.5 * step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - (_fd.order - 1),
    ),
    range(
        mx[1] + step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - 2(_fd.order - 1),
    ),
]
get_mgrid(::dtauxydy, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - 2(_fd.order - 1),
    ),
    range(
        my[1] + step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - 2(_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]
get_mgrid(::dtauxzdx, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        my[1] + step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - 2(_fd.order - 1),
    ),
    range(
        mx[1] + step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - 2(_fd.order - 1),
    ),
]
get_mgrid(::dtauxzdz, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - 2(_fd.order - 1),
    ),
    range(
        my[1] + step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - 2(_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]
get_mgrid(::dtauyzdy, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        my[1] + step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - 2(_fd.order - 1),
    ),
    range(
        mx[1] + step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - 2(_fd.order - 1),
    ),
]
get_mgrid(::dtauyzdz, ::FdtdElastic, mz, my, mx) = [
    range(
        mz[1] + step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - 2(_fd.order - 1),
    ),
    range(
        my[1] + 0.5 * step(my) * (_fd.order - 1),
        step = step(my),
        length = length(my) - (_fd.order - 1),
    ),
    range(
        mx[1] + step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - 2(_fd.order - 1),
    ),
]

get_mgrid(::dtauxxdx, ::FdtdElastic, mz, mx) = [
    range(
        mz[1] + step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - 2(_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]
get_mgrid(::dtauzzdz, ::FdtdElastic, mz, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        mx[1] + step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - 2(_fd.order - 1),
    ),
]
get_mgrid(::dtauxzdx, ::FdtdElastic, mz, mx) = [
    range(
        mz[1] + 0.5 * step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - (_fd.order - 1),
    ),
    range(
        mx[1] + step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - 2(_fd.order - 1),
    ),
]
get_mgrid(::dtauxzdz, ::FdtdElastic, mz, mx) = [
    range(
        mz[1] + step(mz) * (_fd.order - 1),
        step = step(mz),
        length = length(mz) - 2(_fd.order - 1),
    ),
    range(
        mx[1] + 0.5 * step(mx) * (_fd.order - 1),
        step = step(mx),
        length = length(mx) - (_fd.order - 1),
    ),
]


text:
# using SymPy

# x, y = @vars x y
# u = x*(1-x)*y*(1-y)
# v = x*(1-x)*y*(1-y)^2

# p = x*(1-x)*y*(1-y)

# ux = diff(u,x)
# uy = diff(u,y)
# vx = diff(v,x)
# vy = diff(v,y)
# px = diff(p,x)
# py = diff(p,y)
# f = u*ux + v*uy - (diff(ux,x)+diff(uy,y)) + px
# g = u*vx + v*vy - (diff(vx,x)+diff(vy,y)) + py
# h = diff(u, x) + diff(v, y)
# println(replace(replace(sympy.julia_code(simplify(f)), ".^"=>"^"), ".*"=>"*"))
# println(replace(replace(sympy.julia_code(simplify(g)), ".^"=>"^"), ".*"=>"*"))
# println(replace(replace(sympy.julia_code(simplify(h)), ".^"=>"^"), ".*"=>"*"))

function u_exact(x,y)
    x*(1-x)*y*(1-y)
end

function v_exact(x,y)
    x*(1-x)*y*(1-y)
end

function p_exact(x,y)
    x*(1-x)*y*(1-y)
end

function t_exact(x,y)
    x*(1-x)*y*(1-y)
end

function k_exact(x,y)
    1.0
end

function ffunc_(x, y)
    x*y*(1 - x)*(1 - y)*(-x*y*(1 - x) + x*(1 - x)*(1 - y)) + x*y*(1 - x)*(1 - y)*(-x*y*(1 - y) + y*(1 - x)*(1 - y)) - x*y*(1 - y) + 0.02*x*(1 - x) + y*(1 - x)*(1 - y) + 0.02*y*(1 - y)
end

function gfunc_(x, y)
    x*y*(1 - x)*(1 - y)*(-x*y*(1 - x) + x*(1 - x)*(1 - y)) + x*y*(1 - x)*(1 - y)*(-x*y*(1 - y) + y*(1 - x)*(1 - y)) - x*y*(1 - x) + x*(1 - x)*(1 - y) + 0.02*x*(1 - x) + 0.02*y*(1 - y)
end

function hfunc_(x,y)
    -x*y*(1 - x) - x*y*(1 - y) + x*(1 - x)*(1 - y) + y*(1 - x)*(1 - y)
end

function heat_source_func(x, y)
    x^2*y*(x - 1)^2*(y - 1)*(2*y - 1) + x*y^2*(x - 1)*(2*x - 1)*(y - 1)^2 - 2.0*x*(x - 1) - 2.0*y*(y - 1)
end

function q1_func_(t, x, y)
    -1.0*x*y*(1 - x)*(1 - y)
end

function q2_func_(t, x, y)
    -1.0*x*y*(1 - x)*(1 - y)
end

using LinearAlgebra
using MAT
using AdFem
using PyPlot
using SparseArrays

m = 50
n = 50
h = 1/n
nu = 0.01
NT_transport = 500
Δt = 1/NT_transport
κ1 = 1.0
κ2 = 1.0

NT = 5    # number of iterations for Newton's method

bd = bcnode("all", m, n, h)
bd = [bd; bd .+ (m+1)*(n+1);
    2*(m+1)*(n+1)+(n-1)*m+1:2*(m+1)*(n+1)+(n-1)*m+2;
    (2*(m+1)*(n+1)+m*n) .+ bd]
# only apply Dirichlet to velocity; set left bottom two points to zero to fix rank deficient problem for pressure

## STEP 1: NAVIER STOKES EQUATIONS
F1 = compute_fem_source_term1(eval_f_on_gauss_pts(ffunc_, m, n, h), m, n, h)
F2 = compute_fem_source_term1(eval_f_on_gauss_pts(gfunc_, m, n, h), m, n, h)
H = h^2*eval_f_on_fvm_pts(hfunc_, m, n, h)
B = constant(compute_interaction_matrix(m, n, h))

# compute F
Laplace = nu * constant(compute_fem_laplace_matrix1(m, n, h))
heat_source = eval_f_on_gauss_pts(heat_source_func, m, n, h)
heat_source = constant(compute_fem_source_term1(heat_source, m, n, h))
# The temperature term is
# T_xx + T_yy + q = u T_x + v T_y

function compute_residual(S)
    u, v, p, T = S[1:(m+1)*(n+1)],
        S[(m+1)*(n+1)+1:2(m+1)*(n+1)],
        S[2(m+1)*(n+1)+1:2(m+1)*(n+1)+m*n],
        S[2(m+1)*(n+1)+m*n+1:end]
    G = eval_grad_on_gauss_pts([u;v], m, n, h)
    ugauss = fem_to_gauss_points(u, m, n, h)
    vgauss = fem_to_gauss_points(v, m, n, h)
    ux, uy, vx, vy = G[:,1,1], G[:,1,2], G[:,2,1], G[:,2,2]

    interaction = compute_interaction_term(p, m, n, h) # julia kernel needed
    f1 = compute_fem_source_term1(ugauss.*ux, m, n, h)
    f2 = compute_fem_source_term1(vgauss.*uy, m, n, h)
    f3 = -interaction[1:(m+1)*(n+1)]
    f4 = Laplace*u
    f5 = -F1
    F = f1 + f2 + f3 + f4 + f5

    g1 = compute_fem_source_term1(ugauss.*vx, m, n, h)
    g2 = compute_fem_source_term1(vgauss.*vy, m, n, h)
    g3 = -interaction[(m+1)*(n+1)+1:end]
    g4 = Laplace*v
    g5 = -F2
    G = g1 + g2 + g3 + g4 + g5

    H0 = -B * [u;v] + H

    # T_xx + T_yy = u T_x + v T_y - heat_source
    # ugauss = zeros_like(ugauss)
    # vgauss = zeros_like(vgauss)
    A = constant(compute_fem_laplace_matrix1(m, n, h))
    T0 = A * T + compute_fem_advection_matrix1(ugauss,vgauss, m, n, h) * T - heat_source
    R = [F;G;H0;T0]
    return R
end

function compute_jacobian(S)
    u, v, p, T = S[1:(m+1)*(n+1)],
                 S[(m+1)*(n+1)+1:2(m+1)*(n+1)],
                 S[2(m+1)*(n+1)+1:2(m+1)*(n+1)+m*n],
                 S[2(m+1)*(n+1)+m*n+1:end]

    graduv = eval_grad_on_gauss_pts([u;v], m, n, h)
    ugauss = fem_to_gauss_points(u, m, n, h)
    vgauss = fem_to_gauss_points(v, m, n, h)
    ux, uy, vx, vy = graduv[:,1,1], graduv[:,1,2], graduv[:,2,1], graduv[:,2,2]

    M1 = constant(compute_fem_mass_matrix1(ux, m, n, h))
    M2 = constant(compute_fem_advection_matrix1(constant(ugauss), constant(vgauss), m, n, h)) # a julia kernel needed
    M3 = Laplace
    Fu = M1 + M2 + M3

    Fv = constant(compute_fem_mass_matrix1(uy, m, n, h))

    N1 = constant(compute_fem_mass_matrix1(vy, m, n, h))
    N2 = constant(compute_fem_advection_matrix1(constant(ugauss), constant(vgauss), m, n, h))
    N3 = Laplace
    Gv = N1 + N2 + N3

    Gu = constant(compute_fem_mass_matrix1(vx, m, n, h))

    M = constant(compute_fem_laplace_matrix1(m, n, h)) +
         compute_fem_advection_matrix1(ugauss,vgauss, m, n, h)

    gradT = eval_grad_on_gauss_pts1(T, m, n, h)
    Tx, Ty = gradT[:,1], gradT[:,2]
    DU_TX = constant(compute_fem_mass_matrix1(Tx, m, n, h))       # (m+1)*(n+1), (m+1)*(n+1)
    DV_TY = constant(compute_fem_mass_matrix1(Ty, m, n, h))       # (m+1)*(n+1), (m+1)*(n+1)

    J0 = [Fu Fv                         # Jacobian of [u; v]
          Gu Gv]

    J1 = [J0 -B'                        # Jacobian of [u; v; p]
          -B spdiag(zeros(size(B,1)))]

    N = 2*(m+1)*(n+1) + m*n                                     # length of [u; v; p]
    J = [J1 SparseTensor(spzeros(N,(m+1)*(n+1)))                # Jacobian of [u; v; p; T]
        DU_TX DV_TY SparseTensor(spzeros((m+1)*(n+1), m*n)) M]
end

function solve_steady_cavityflow_one_step(S)
    residual = compute_residual(S)
    J = compute_jacobian(S)

    J, _ = fem_impose_Dirichlet_boundary_condition1(J, bd, m, n, h)
    residual = scatter_update(residual, bd, zeros(length(bd)))    # residual[bd] .= 0.0 in Tensorflow syntax

    d = J\residual
    residual_norm = norm(residual)
    op = tf.print("residual norm", residual_norm)
    d = bind(d, op)
    S_new = S - d
    return S_new
end

function condition(i, S_arr)
    i <= NT + 1
end

function body(i, S_arr)
    S = read(S_arr, i-1)
    op = tf.print("i=",i)
    i = bind(i, op)
    S_new = solve_steady_cavityflow_one_step(S)
    S_arr = write(S_arr, i, S_new)
    return i+1, S_arr
end

xy = fem_nodes(m, n, h)
x, y = xy[:,1], xy[:,2]
u0 = @. u_exact(x,y)
v0 = @. v_exact(x,y)
T0 = @. t_exact(x,y)

xy = fvm_nodes(m, n, h)
x, y = xy[:,1], xy[:,2]
p0 = @. p_exact(x,y)

S_arr = TensorArray(NT+1)
S_arr = write(S_arr, 1, [u0; v0; p0; T0])
# S_arr = write(S_arr, 1, zeros(m*n+3*(m+1)*(n+1)))

i = constant(2, dtype=Int32)

_, S = while_loop(condition, body, [i, S_arr])
S = set_shape(stack(S), (NT+1, 2*(m+1)*(n+1)+m*n+(m+1)*(n+1)))

u = S[NT+1, 1:(m+1)*(n+1)]
v = S[NT+1, (m+1)*(n+1)+1:2*(m+1)*(n+1)]
p = S[NT+1, 2*(m+1)*(n+1)+1:2*(m+1)*(n+1)+m*n]
T = S[NT+1, 2*(m+1)*(n+1)+m*n+1:end]

## STEP 2: TRANSPORT EQUATION
# pre-compute source term
Q1 = zeros(NT_transport, (m+1)*(n+1))
Q2 = zeros(NT_transport, (m+1)*(n+1))
for i = 1:NT_transport
    t = i*Δt
    Q1[i, :] = eval_f_on_fem_pts((x,y)->q1_func_(t, x, y), m, n, h)
    Q2[i, :] = eval_f_on_fem_pts((x,y)->q2_func_(t, x, y), m, n, h)
end
Q1 = constant(Q1)
Q2 = constant(Q2)

function solve_transport_step(w1, w2, q1, q2)
    w1 = (1/Δt * w1 + κ1 * u + q1) / (1/Δt + κ1)
    w2 = (1/Δt * w2 + κ2 * v + q2) / (1/Δt + κ2)
    w1, w2
end

function transport_condition(i, w1_arr, w2_arr)
    i <= NT_transport
end

function transport_body(i, w1_arr, w2_arr)
    w1 = read(w1_arr, i)
    w2 = read(w2_arr, i)
    q1 = Q1[i]
    q2 = Q2[i]
    # op = tf.print("transport equation, step: ", i)
    # i = bind(i, op)
    w1, w2 = solve_transport_step(w1, w2, q1, q2)
    i+1, write(w1_arr, i+1, w1), write(w2_arr, i+1, w2)
end

i = constant(1, dtype = Int32)
w1_arr = TensorArray(NT_transport+1)
w2_arr = TensorArray(NT_transport+1)
w1_0 = eval_f_on_fem_pts((x,y)->x*y*(1-x)*(1-y), m, n, h)
w2_0 = eval_f_on_fem_pts((x,y)->x*y*(1-x)*(1-y), m, n, h)
w1_arr = write(w1_arr, 1, w1_0)
w2_arr = write(w2_arr, 1, w2_0)
_, w1, w2 = while_loop(transport_condition, transport_body, [i, w1_arr, w2_arr])
w1, w2 = stack(w1), stack(w2)
w1 = set_shape(w1, (NT_transport+1, (m+1)*(n+1)))
w2 = set_shape(w2, (NT_transport+1, (m+1)*(n+1)))

sess = Session(); init(sess)
W1, W2, u_, v_, p_ = run(sess, [w1, w2, u, v, p])

figure(figsize=(15,5))
subplot(131)
visualize_scalar_on_fem_points(W1[end,:], m, n, h)
title("Computed droplet x-velocity")
subplot(132)
visualize_scalar_on_fem_points(eval_f_on_fem_pts((x,y)->exp(-1)*(1-x)*x*(1-y)*y, m, n, h), m, n, h)
title("Exact droplet x-velocity")
subplot(133)
visualize_scalar_on_fem_points(W1[end,:]-eval_f_on_fem_pts((x,y)->exp(-1)*(1-x)*x*(1-y)*y, m, n, h), m, n, h)
title("Difference")

figure(figsize=(15,5))
subplot(131)
visualize_scalar_on_fem_points(W2[end,:], m, n, h)
title("Computed droplet y-velocity")
subplot(132)
visualize_scalar_on_fem_points(eval_f_on_fem_pts((x,y)->exp(-1)*(1-x)*x*(1-y)*y, m, n, h), m, n, h)
title("Exact droplet y-velocity")
subplot(133)
visualize_scalar_on_fem_points(W2[end,:]-eval_f_on_fem_pts((x,y)->exp(-1)*(1-x)*x*(1-y)*y, m, n, h), m, n, h)
title("Difference")
tight_layout()

# # S = output
# # # out_v = output[:, 1:2*(m+1)*(n+1)]
# # # out_p = output[:, 2*(m+1)*(n+1)+1:end]

matwrite("covid_data.mat", Dict(
        "W1"=>W1, "W2"=>W2))


text:
@testset "Dynamics" begin

"Tests that SphericalMidpoint conserves energy for simple forces to a certain tolerance."
function test_spherical_midpoint()
    crystal = Sunny.diamond_crystal()

    interactions = [
        diamond_test_exchanges()...,
        external_field([0, 0, 1])
    ]

    sys = SpinSystem(crystal, interactions, (5, 5, 5))
    Random.seed!(0)
    rand!(sys)

    NITERS = 5_000
    Δt     = 0.01
    energies = Float64[]

    integrator = Sunny.SphericalMidpoint(sys)
    for it in 1:NITERS
        evolve!(integrator, Δt)
        # Compute the energy
        push!(energies, energy(sys))
    end

    # Check that the energy hasn't fluctuated much. Expected fluctuations should
    # scale like square-root of system size.
    sqrt_size = sqrt(length(sys.sites))
    ΔE = (maximum(energies) - minimum(energies)) / sqrt_size
    @test ΔE < 1e-2
end

test_spherical_midpoint()

"Tests that SphericalMidpoint conserves energy for dipole forces (w/ Fourier) to a certain tolerance."
function test_dipole_ft()
    crystal = Sunny.diamond_crystal()
    interactions = [dipole_dipole()]
    sys = SpinSystem(crystal, interactions, (5, 5, 5))
    Random.seed!(0)
    rand!(sys)

    NITERS = 5_000
    Δt     = 0.01
    energies = Float64[]

    integrator = Sunny.SphericalMidpoint(sys)
    for it in 1:NITERS
        evolve!(integrator, Δt)
        # Compute the energy
        push!(energies, energy(sys))
    end

    # Check that the energy hasn't fluctuated much
    sqrt_size = sqrt(length(sys.sites))
    ΔE = (maximum(energies) - minimum(energies)) / sqrt_size
    @test ΔE < 2e-2
end

test_dipole_ft()

"Tests that set_temp!/get_temp behave as expected"
function test_set_get_temp_langevin()
    system = produce_example_system()

    Random.seed!(1111)
    rand_kTs = rand(100)

    sampler = LangevinSampler(system, 1.0, 1.0, 1.0, 1)
    for kT in rand_kTs
        set_temp!(sampler, kT)
        if get_temp(sampler) != kT
            return false
        end
    end

    return true
end

@test test_set_get_temp_langevin()

end

# == These should not be @test-ed, but are for manual inspection == #

"Measure timings for speed disparity between real and fourier-space dipole interactions"
function time_real_fourier_dipole()
    lat_vecs = [1 0 0; 0 1 0; 0 0 1]
    positions = [[0, 0, 0], [1, 1, 1]/2]
    crystal = Crystal(lat_vecs, positions)

    Ls = [2, 6, 10, 14, 18, 22, 26, 30]
    real_times = Float64[]
    ft_times = Float64[]
    for L in Ls
        println("Measuring L = $L")

        # We'll manually set interactions, since DipoleRealCPU
        #  isn't actually a create-able force in the current
        #  SpinSystem/HamiltonianCPU constructor.
        interactions = Sunny.Interaction[]
        latsize = (L, L, L)
        sys = SpinSystem(crystal, interactions, latsize)
        Random.seed!(0)
        rand!(sys)

        # Since we just want to time integration speed,
        #  it's alright for the Ewald summation to be
        #  very inaccurate (truncated at a single cell)
        dip_dip = dipole_dipole(;extent=1)
        empty_h = Sunny.HeisenbergCPU{3}[]
        empty_d = Sunny.DiagonalCouplingCPU{3}[]
        empty_g = Sunny.GeneralCouplingCPU{3}[]

        if L <= 14
            println("\tSetting up real-space...")
            dipr = Sunny.DipoleRealCPU(dip_dip, crystal, latsize)
            ℋ = Sunny.HamiltonianCPU(
                nothing, empty_h, empty_d, empty_g, dipr
            )
            sys.hamiltonian = ℋ
            integrator = HeunP(sys)

            println("\tStarting real-space...")
            evolve!(integrator, 0.001)
            _, t, _, _ = @timed for _ in 1:100
                evolve!(integrator, 0.001)
            end
            push!(real_times, t)
        end

        println("\tSetting up Fourier-space...")
        dipf = Sunny.DipoleFourierCPU(dip_dip, crystal, latsize)
        ℋ = Sunny.HamiltonianCPU(
            nothing, empty_h, empty_d, empty_g, dipf
        )
        sys.hamiltonian = ℋ
        integrator = HeunP(sys)

        println("\tStarting Fourier-space...")
        evolve!(integrator, 0.001)
        _, t, _, _ = @timed for _ in 1:100
            evolve!(integrator, 0.001)
        end
        push!(ft_times, t)
    end
    return real_times, ft_times
end


text:
module BranchLineCouplers

using Printf
using Polynomials
using GenericLinearAlgebra
using GLMakie

export Coupler,
    analysis_functions,
    synthesise_coupler,
    plot_coupler,
    interactive_coupler_design

include("coupler.jl")
include("synthesis.jl")
include("analysis.jl")
include("visualization.jl")
include("interactive.jl")

end # module