text:
# This scripts shows how to train a network with feedback operations (as well as
# a comparable forward network) on MNIST.
usegpu = false

using Base.Iterators: partition
using Flux, Statistics
using Flux: onehotbatch, onecold, crossentropy
using MLDatasets
using BSON: @save
using FeedbackNets
if usegpu
    using CuArrays
end

include("DataHelpers.jl")
using .DataHelpers

# load MNIST
@info("Preparing data ...")
batchsize=200
images, labels = MNIST.traindata()
images = DataHelpers.standardize(Float32.(images))
trainimgs, trainlbls = DataHelpers.makebatches(images[:, :, 1:55000], labels[1:55000], batchsize=batchsize)
valimgs, vallbls = DataHelpers.makebatches(images[:, :, 55001:end], labels[55001:end], batchsize=batchsize)
trainimgs = map(x -> reshape(x, size(x, 1), size(x, 2), 1, size(x, 3)), trainimgs)
valimgs = map(x -> reshape(x, size(x, 1), size(x, 2), 1, size(x, 3)), valimgs)
trainlbls = map(x -> onehotbatch(x, 0:9), trainlbls)
vallbls = map(x -> onehotbatch(x, 0:9), vallbls)

if usegpu
    trainimgs = gpu.(trainimgs)
    trainlbls = gpu.(trainlbls)
    valimgs = gpu.(valimgs)
    vallbls = gpu.(vallbls)
end

# construct models
@info("Building models ...")
# Forward network: this is essentially LeNet5, with the difference that ReLUs
# are used and that the first layer uses a padding of 2 in order to accomodate
# input images of size 28x28 instead of 32x32.
forwardnet = Chain(
    Conv((5,5), 1=>6, relu, pad=2),
    MaxPool((2,2), stride=(2,2)),
    Conv((5,5), 6=>16, relu),
    MaxPool((2,2), stride=(2,2)),
    Conv((5,5), 16=>120, relu),
    x -> reshape(x, 120, size(x, 4)),
    Dense(120, 84, relu),
    Dense(84, 10),
    softmax
)
if usegpu
    forwardnet = forwardnet |> gpu
end
# Feedback network: the same structure as the forward network, but with feedback
# from the output of the second convolution to before the first pooling, and
# from the output of the first fully connected layer (here implemented as a
# convolution) to before the second pooling layer.
feedbacknet = FeedbackChain(
    Conv((5,5), 1=>6, relu, pad=2),
    Merger("conv2", ConvTranspose((10,10), 16=>6, relu, stride=2), +),
    MaxPool((2,2), stride=(2,2)),
    Conv((5,5), 6=>16, relu),
    Splitter("conv2"),
    Merger("fc1", ConvTranspose((10,10), 120=>16, relu), +),
    MaxPool((2,2), stride=(2,2)),
    Conv((5,5), 16=>120, relu),
    Splitter("fc1"),
    x -> reshape(x, 120, size(x, 4)),
    Dense(120, 84, relu),
    Dense(84, 10),
    softmax
)
# generate initial state and wrap feedback net in Recur
h = Dict(
    "conv2" => zeros(Float32, 10, 10, 16, batchsize),
    "fc1" => zeros(Float32, 1, 1, 120, batchsize)
)
if usegpu
    feedbacknet = feedbacknet |> gpu
    h = Dict(key => gpu(val) for (key, val) in pairs(h))
end
feedbacknet = Flux.Recur(feedbacknet, h)

# Training the feedforward model works with the standard Flux workflows. We
# generate a loss function using the pre-implemented crossentropy, ...
function loss(x, y, model)
    error = crossentropy(model(x), y)
    # This reset is necessary for the feedback network later. The way Flux Recur
    # handles recurrency requires a reset after each loss calculation.
    Flux.reset!(model)
    return(error)
end
# ... and an accuracy function to monitor our learning progress.
function accuracy(valset, model)
    accum = 0.0
    for (x, y) in valset
        accum += mean(x == y for (x, y) in zip(onecold(model(x)), onecold(y)))
    end
    accum = accum / length(valset)
end
# We also wrap training in a function, as executing the for-loop etc. in script
# mode may be slow.
function trainmodel(model, trainset, valset; epochs=1, opt=Nesterov(0.0001),
                    verbose=true, loss=loss, accuracy=accuracy)
    # training loop
    verbose ? @info("initial accuracy: $(accuracy(valset, model))") : nothing
    for e in 1:epochs
        Flux.train!((x, y) -> loss(x, y, model), params(model), trainset, opt)
        verbose ? @info("accuracy after episode $e: $(accuracy(valset, model))") : nothing
    end
end
# Now we can do the actual training.
@info("Training forward network ...")
trainmodel(forwardnet, zip(trainimgs, trainlbls), zip(valimgs, vallbls))
@save "forward_MNIST.bson"

# Training the feedback network involves one additional complication: we need to
# unroll the network over time. If we work with sequence data, this amounts to
# broadcasting our model over the sequence. Since we are using MNIST digits here
# we first need to replicate them over timesteps.
#
# This helper function takes care of this replication and calculates the loss.
function lossunroll(x, y, model; steps=4)
    error = sum(crossentropy(model(x) .+ eps(Float32), y) for i in 1:steps)
    Flux.reset!(model)
    return error
end
# Similarly, we need to unroll for several steps to calculate the accuracy. We
# will compare the last prediction of the network to the ground truth.
function accuracyunroll(valset, model; steps=4)
    accum = 0.0
    for (x, y) in valset
        prediction = model.(x for i in 1:steps)[end]
        Flux.reset!(model)
        accum += mean(x == y for (x, y) in zip(onecold(prediction), onecold(y)))
    end
    accum = accum / length(valset)
end

@info("Training feedback network ...")
trainmodel(feedbacknet, zip(trainimgs, trainlbls), zip(valimgs, vallbls),
           loss=lossunroll, accuracy=accuracyunroll)
@save "feedback_MNIST.bson"


text:
module GraphTransformer

using Flux
using Flux:@functor, glorot_uniform, Zygote
using GeometricFlux
using GeometricFlux:_view, edge_index_table
using GraphSignals
using LightGraphs:edges
using GraphLaplacians
using Statistics

export GraphTransformerLayer
export MultiHeadAttentionLayer
export message, apply_batch_message, aggregate_neighbors, propagate

include("multi_head_attn.jl")
include("graph_transformer_layer.jl")

end

ext: jl

max_stars_repo_name: DwaraknathT/Flux-Graph-Transformer

max_issues_repo_name: DwaraknathT/Flux-Graph-Transformer

max_forks_repo_name: DwaraknathT/Flux-Graph-Transformer


text:
@testset "weak scatterers" begin
    # background medium
    medium = Medium(1.0,1.0+0.0im)

    # angular frequencies
    ωs = [0.001,20.]

    species = [
        Specie(ρ=10.,r=0.01, c=12., volfrac=0.05),
        Specie(ρ=3., r=0.2, c=2.0, volfrac=0.04)
    ]

    # wavenumbers
    eff_medium = effective_medium(medium, species)

    k_eff_lows = ωs./eff_medium.c
    k_eff_φs = wavenumber_low_volfrac(ωs, medium, species; tol=1e-9)

    k_effs = [
        wavenumbers(ω, medium, species;
            num_wavenumbers=1, tol = 1e-8)
    for ω in ωs]
    inds = [argmin(abs.(k_effs[i] .- k_eff_φs[i])) for i in eachindex(ωs)]
    k_effs2 = [k_effs[i][inds[i]] for i in eachindex(inds)]

    @test norm(k_effs2 - k_eff_φs)/norm(k_eff_φs) < 0.002
    @test norm(k_effs2 - k_eff_lows)/norm(k_eff_lows) < 0.01
    @test norm(k_effs2[1] - k_eff_lows[1])/norm(k_effs[1]) < 4e-7

    wave_effs_2 = [EffectiveWave(ωs[i], k_effs2[i], medium, species; tol = 1e-9) for i in eachindex(ωs)]
    wave_effs_φs = [EffectiveWave(ωs[i], k_eff_φs[i], medium, species; tol = 1e-9) for i in eachindex(ωs)]

    # reflection coefficient
    Rs2 = reflection_coefficients(ωs, medium, species; tol=1e-9, num_wavenumbers=1)

    # if the wavenumber k_effs are already calculated, the below is faster
    # pairs each ω in ωs with each wave in wave_effs2 to calculate the refleciton coefficients

    Rs = map(eachindex(ωs)) do i
        w_effs = [EffectiveWave(ωs[i], k_eff, medium, species; tol = 1e-9) for k_eff in k_effs[i]]
        reflection_coefficient(ωs[i], w_effs, medium, species; tol = 1e-9)
    end

    @test norm(Rs - Rs2)/norm(Rs) < 5e-6

    Rs1 = map(eachindex(ωs)) do i
        w_effs = EffectiveWave(ωs[i], k_effs[i][1], medium, species; tol = 1e-9)
        reflection_coefficient(ωs[i], w_effs, medium, species)
    end

    # Direct incidence
    R_low = reflection_coefficient_halfspace(medium, eff_medium; tol=1e-9)
    Rs_φs = reflection_coefficients(ωs, wave_effs_φs, medium, species; tol=1e-9)
    # the below takes a low-volfrac expansion for both the wavenumber and reflection coefficient
    Rs_φs2 = reflection_coefficient_low_volfrac(ωs, medium, species; tol=1e-9)

    len = length(ωs)
    @test maximum(abs.(Rs_φs - Rs1)) < 2e-4 # already relative to incident wave amplitude = 1
    @test maximum(abs.(Rs_φs2 - Rs1)) < 1e-3
    @test abs(R_low - Rs1[1]) < 1e-7

    # Vary angle of incidence θin. Beware of extreme angles θ ~= -pi/2 or pi/2.
    θs = (-π/2.1):0.3:(π/2.1)
    R_low = [
        reflection_coefficient_halfspace(medium, eff_medium; θin = θ, tol=1e-9)
    for θ in θs]

    Rs = map(θs) do θ  # , hankel_order =7
        wave_effs_2 = [
            EffectiveWave(ωs[i], k_effs2[i], medium, species;
                tol = 1e-9, θin = θ
            )
        for i in eachindex(ωs)]
        reflection_coefficients(ωs, wave_effs_2, medium, species; θin = θ)
    end
    Rs_φs = [
        reflection_coefficient_low_volfrac(ωs, medium, species; θin = θ, tol = 1e-9, hankel_order =7)
    for θ in θs];

    @test maximum(abs(R_low[i] - Rs[i][1]) for i in 1:length(R_low)) < 1e-7
    @test maximum(norm(R)/len for R in (Rs_φs - Rs)[2:end-2]) < 4e-3
end


text:
"""
Solution to a problem.

    (::Solution)(r, t)
    (::Solution)(ϕ)

Evaluate the solution.

# Properties
- `i`: initial value.
- `b`: boundary value.
- `d_dϕb`: ϕ-derivative at the boundary, where ϕ is the Boltzmann variable.
- `ϕb`: boundary constant. See also [`rb`](@ref).
- `ϕi`: for `ϕ≥ϕi`, the solution evaluates to the initial value.
- `iterations`: number of iterations needed to find this solution.

Type `\\phi<tab>` to obtain the `ϕ` symbol.
"""
struct Solution{_Todesol,_Teq,_T,_Td_dϕ,_Tϕ} <: TransformedFunction
    _odesol::_Todesol
    _eq::_Teq
    i::_T
    b::_T
    d_dϕb::_Td_dϕ
    ϕb::_Tϕ
    ϕi::_Tϕ
    iterations::Int

    function Solution(odesol, eq, iterations)
        new{typeof(odesol),typeof(eq),typeof(odesol.u[1][1]),typeof(odesol.u[1][2]),eltype(odesol.t)}(odesol, eq, odesol.u[end][1], odesol.u[1][1], odesol.u[1][2], odesol.t[1], odesol.t[end], iterations)
    end
end

function (sol::Solution)(ϕ)
    if ϕ > sol.ϕi
        return sol.i
    elseif ϕ < sol.ϕb
        return typeof(sol.b)(NaN)
    end

    return sol._odesol(ϕ, idxs=1)
end

"""
    d_dϕ(::Solution, r, t)
    d_dϕ(::Solution, ϕ)

ϕ-derivative of the solution, where ϕ is the Boltzmann variable.

Type `\\phi<tab>` to obtain the `ϕ` symbol.

See also: [`ϕ`](@ref)
"""
function d_dϕ(sol::Solution, ϕ)
    if ϕ > sol.ϕi
        return zero(sol.d_dϕb)
    elseif ϕ < sol.ϕb
        return typeof(sol.d_dϕb)(NaN)
    end

    return sol._odesol(ϕ, idxs=2)
end

"""
    ∂_∂r(::Solution, r, t)

Spatial derivative of the solution.

Type `\\partial<tab>` to obtain the `∂` symbol.

---

    ∂_∂r(::Solution, :b, t)

Spatial derivative of the solution at the boundary.

"""
function ∂_∂r(sol::Solution, symbol::Symbol, t)
    @argcheck symbol === :b
    ∂_∂r(sol, rb(sol,t), t)
end

"""
    ∂_∂t(::Solution, r, t)

Time derivative of the solution sol.

Type `\\partial<tab>` to obtain the `∂` symbol.

---

    ∂_∂t(::Solution, :b, t)

Time derivative of the solution at the boundary.
"""
function ∂_∂t(sol::Solution, symbol::Symbol, t)
    @argcheck symbol === :b
    ∂_∂t(sol, rb(sol,t), t)
end

"""
    flux(::Solution, r, t)

Diffusive flux of the solution.
"""
flux(sol::Solution, r, t) = flux(sol._eq, sol, r, t)

"""
    flux(::Solution, :b, t)

Diffusive flux of the solution at the boundary.
"""
function flux(sol::Solution, symbol::Symbol, t)
    @argcheck symbol === :b
    flux(sol, rb(sol,t), t)
end

"""
    rb(::Solution, t)

Location of the boundary in the solution at time `t`, equal to `ϕb*√t`.
"""
rb(sol::Solution, t) = r(sol.ϕb, t)

Base.broadcastable(sol::Solution) = Ref(sol)

function Base.show(io::IO, sol::Solution)
    println(io, "Solution $(sol._eq.symbol) obtained after $(sol.iterations) iterations")
    println(io, "$(sol._eq.symbol)b = $(sol.b)")
    println(io, "d$(sol._eq.symbol)/dϕ|b = $(sol.d_dϕb)")
    if !iszero(sol.ϕb)
        println(io, "ϕb = $(sol.ϕb)")
    end
    print(io, "$(sol._eq.symbol)i = $(sol.i)")
end

# Plot recipe
@recipe function plot(sol::Solution; label=string(sol._eq.symbol),
                                     legend=false,
                                     xguide="ϕ=r/√t",
                                     yguide=label)
    vars := 1  # Plot only u[1] = sol
    sol._odesol  # Delegate the rest to the contained ODESolution
end


text:
@symbol_func function K_CD(cur_reactor::AbstractReactor)
  cur_K = K_F(cur_reactor)

  cur_K *= K_n(cur_reactor)

  cur_K *= cur_reactor.eta_CD

  cur_K /= cur_reactor.Q

  cur_K
end


text:
type BamIOStream <: IO
    name::AbstractString
    handle::Ptr{Void}
    phdr::Ptr{Header}
    precord::Ptr{Record}
    pkstr::Ptr{KString}
    eof::Bool
    function BamIOStream(name::AbstractString,handle::Ptr{Void},phdr::Ptr{Header},
                         prec::Ptr{Record},pkstr::Ptr{KString})
        new(name,handle,phdr,prec,pkstr,false)
    end
end

@doc """ general file open function
""" ->
function open(fname::AbstractString,mode::AbstractString,ty::AbstractString)
    fname = realpath(fname)
    ty == "bam" ? bam_open(fname,mode) :
    ty == "sam" ? sam_open(fname,mode) : open(fname,mode)
end

@doc """ bam file name must endwith bam
""" ->
function bam_open(fname::AbstractString, mode::AbstractString="rb")
    endswith(fname, ".bam") || info("file name doesn't endwith bam")
    fname = realpath(fname)
    bam_fl = hts_open(fname, mode)
    if mode == "rb" || mode == "r"
        # bam_hdr_read(bam_fl->fp.bgzf)
        phdr = sam_hdr_read(bam_fl) #header info
    else
        phdr = bam_hdr_init()
    end
    prec = bam_init1()
    kstr = KString(0,0,C_NULL)
    pkstr = convert(Ptr{KString}, pointer_from_objref(kstr))

    BamIOStream(fname,bam_fl,phdr,prec,pkstr)
end

function bam_hdr_read!(bios::BamIOStream)
    bios.phdr = bam_hdr_read(bios.handle)
end
#=
@doc """ read a record from a bam iostream
""" ->
function readline(bios::BamIOStream)
    handle = bios.handle
    phdr = bios.phdr
    prec = bios.precord
    ret = sam_read!(handle,phdr,prec)
    if ret == 0
        bios.eof = true
    elseif ret == -1
        bios.eof = false
        info("error occured in sam read ")
    end
    pkstr = bios.pkstr
    sam_format!(phdr,prec,pkstr)

    strptr(pkstr)
end

@doc """ eof(bios) -> Bool
         Tests whether an Bam I/O stream is at end-of-file.
""" ->
function eof(bios::BamIOStream)
    bios.eof
end

@doc """ read all the records from a bam file
""" ->
function readlines(fname::AbstractString)
    fname = realpath(fname)
    bios = bam_open(fname,"rb")
    data = readlines(bios)
    close(bios)
    data
end

@doc """ read all the records from a BamIOStream
""" ->
function readlines(bios::BamIOStream)
    size = 256
    data = Array{ASCIIString,1}(size)
    i = 0
    while !eof(bios)
        i += 1
        if i > size
            resize!(data,2*size)
            size = 2*size
        end
        line = readline(bios)
        data[i] = line
    end
    data[1:i]
end

@doc """ write a record to a BamIOStream
""" ->
function writeline(bios::BamIOStream, str::ASCIIString)
#    error("Wait for hdr function finished")
    kstr = KString(str)
    pkstr = convert(Ptr{KString}, pointer_from_objref(kstr))
    handle = bios.handle
    phdr = bios.phdr
    prec = bios.precord
    #sam_parse1(pkstr, phdr, prec)
    sam_write1(handle,phdr,prec)
end

@doc """ write all the records in strs to a BamIOStream
""" ->
function writelines(bios::BamIOStream, strs::Array{ASCIIString,1})
    #sam_hdr_write will use bam_hdr_write
    #sam_hdr_write(bios.handle,bios.phdr)
    for str in strs
        writeline(bios, str)
    end
end

@doc """ close a BamIOStream
""" ->
function close(bios::BamIOStream)
    #may be header need close but not find function now
    bam_destroy1(bios.precord)
    bam_hdr_destroy(bios.phdr)
    info("prec,pdhr closedd")
    sam_close(bios.handle)
end
=#


text:
module UnumTests

using Unums
using FactCheck

# # write your own tests here
# @fact 1 --> 1

facts("Basics") do
  u = Unums.FixedUnum64(0)
  @fact_throws -u
  @fact_throws u+u
end

FactCheck.exitstatus()
end # module


text:
import Base: +, -, *, /, ^, abs, inv, zero, one, convert, float, ==, <, <=, >, >=, numerator, denominator

struct RationalNumber{T <: Integer} <: Real
    # numerator / denominator = quotient
    numerator::T
    denominator::T

    function RationalNumber(numerator::T, denominator::T) where {T <: Integer}
        if numerator == 0 && denominator == 0
            throw(ArgumentError("Invalid number"))
        end
        cdv = gcd(numerator, denominator)
        sgn = xor(numerator < 0, denominator < 0) ? -1 : 1  # any but not both are negative
        return new{T}(abs(numerator) ÷ cdv * sgn, abs(denominator) ÷ cdv)
    end
end

RationalNumber(numerator) = RationalNumber(numerator, 1)

float(x::RationalNumber) = float(x.numerator / x.denominator)

numerator(x::RationalNumber) = x.numerator
denominator(x::RationalNumber) = x.denominator

show(io::IO, x::RationalNumber) = print(io, "$(x.numerator)//$(x.denominator)")

zero(::Type{RationalNumber{T}}) where {T <: Integer} = RationalNumber(0)
one(::Type{RationalNumber{T}}) where {T <: Integer} = RationalNumber(1)

# Math

function +(x::RationalNumber, y::RationalNumber)
    num = x.numerator * y.denominator + y.numerator * x.denominator
    den = x.denominator * y.denominator
    return RationalNumber(num, den)
end

function -(x::RationalNumber, y::RationalNumber)
    num = x.numerator * y.denominator - y.numerator * x.denominator
    den = x.denominator * y.denominator
    return RationalNumber(num, den)
end

function *(x::RationalNumber, y::RationalNumber)
    num = x.numerator * y.numerator
    den = x.denominator * y.denominator
    return RationalNumber(num, den)
end

function /(x::RationalNumber, y::RationalNumber)
    num = x.numerator * y.denominator
    den = y.numerator * x.denominator
    return RationalNumber(num, den)
end

^(x::RationalNumber, n::Integer) = RationalNumber(x.numerator ^ abs(n), x.denominator ^ abs(n))
^(x::Real, n::RationalNumber) = x ^ float(n)

abs(x::RationalNumber) = RationalNumber(abs(x.numerator), abs(x.denominator))

-(x::RationalNumber) = RationalNumber(-x.numerator, x.denominator)

inv(x::RationalNumber) = RationalNumber(x.denominator, x.numerator)

# Comparisons

==(x::RationalNumber, y::RationalNumber) = (x.numerator == y.numerator) && (x.denominator == y.denominator)
==(x::RationalNumber, y::Real) = float(x) == y
<(x::RationalNumber, y::RationalNumber) = float(x) < float(y)
<(x::RationalNumber, y::Real) = float(x) < y
>(x::RationalNumber, y::RationalNumber) = float(x) > float(y)
>(x::RationalNumber, y::Real) = float(x) > y
<=(x::RationalNumber, y::RationalNumber) = !(x > y)
<=(x::RationalNumber, y::Real) = !(x > y)
>=(x::RationalNumber, y::RationalNumber) = !(x < y)
>=(x::RationalNumber, y::Real) = !(x < y)

max_stars_repo_name: twolodzko/Learning

max_issues_repo_name: twolodzko/Learning

max_forks_repo_name: twolodzko/Learning


text:
# call stack:
# run/replay
#   > calls apply(GreensIterator, group, dqmc, model, sweep)
#       > calls measure(LatticeIterator, measurement, dqmc, model, sweep, GreensResults...)
#           > calls apply!(LatticeIterator, measurement, dqmc, model, GreensResults...)
#               > calls kernel(measurement, dqmc, model, MaskResult..., GreensResults...)

# Does:
# creates groups, runs simulation
#   > resolves GreensIterator (calclates G, Gkk, Gkl or whatever)
#       > resolves sweep (skipping, frequency based measurements) and commits
#           > resolves Mask (to Greens indices)
#               > calculate element from Wicks theorem

# TODO
# try replacing global function rather than creatin new local ones

struct DQMCMeasurement{GI, LI, F <: Function, OT, T} <: AbstractMeasurement
    greens_iterator::GI
    lattice_iterator::LI
    kernel::F
    observable::OT
    temp::T
end

missing_kernel(args...) = error("kernel couldn't be loaded.")

function DQMCMeasurement(
        m::DQMCMeasurement;
        greens_iterator = m.greens_iterator, lattice_iterator = m.lattice_iterator,
        kernel = m.kernel,
        observable = m.observable, temp = m.temp,
        capacity = nothing
    )
    if capacity === nothing
        DQMCMeasurement(greens_iterator, lattice_iterator, kernel, observable, temp)
    else
        binner = rebuild(observable, capacity)
        DQMCMeasurement(greens_iterator, lattice_iterator, kernel, binner, temp)
    end
end
rebuild(B::LogBinner, capacity) = LogBinner(B, capacity=capacity)
rebuild(B::T, capacity) where T = T(B, capacity=capacity)

function Measurement(
        dqmc, _model, greens_iterator, lattice_iterator, kernel;
        capacity = _default_capacity(dqmc), eltype = geltype(dqmc),
        temp = let
            shape = _get_temp_shape(dqmc, _model, lattice_iterator)
            shape === nothing ? nothing : Array{eltype}(undef, shape)
        end,
        obs = let
            shape = _get_final_shape(dqmc, _model, lattice_iterator)
            _zero = shape === nothing ? zero(eltype) : zeros(eltype, shape)
            LogBinner(_zero, capacity=capacity)
        end
    )
    DQMCMeasurement(greens_iterator, lattice_iterator, kernel, obs, temp)
end

################################################################################
### DQMCMeasurement utilities
################################################################################

function Base.show(io::IO, ::MIME"text/plain", m::DQMCMeasurement)
    max = applicable(capacity, m.observable) ? capacity(m.observable) : Inf
    current = length(m.observable)
    GI = m.greens_iterator
    LI = if m.lattice_iterator isa Type
        m.lattice_iterator
    else
        typeof(m.lattice_iterator)
    end
    print(io, "[$current/$max] DQMCMeasurement($GI, $LI)")
end

MonteCarloObservable.mean(m::DQMCMeasurement) = mean(m.observable)
MonteCarloObservable.var(m::DQMCMeasurement) = var(m.observable)
MonteCarloObservable.std_error(m::DQMCMeasurement) = std_error(m.observable)
MonteCarloObservable.tau(m::DQMCMeasurement) = tau(m.observable)
Base.length(m::DQMCMeasurement) = length(m.observable)
Base.isempty(m::DQMCMeasurement) = isempty(m.observable)
Base.empty!(m::DQMCMeasurement) = empty!(m.observable)

"""
    _default_capacity(mc::DQMC)

Returns a default capacity based in the number of sweeps and the measure rate.
"""
_default_capacity(mc::DQMC) = 2 * ceil(Int, mc.parameters.sweeps / mc.parameters.measure_rate)

# _get_temp_shape is the shape of the temporary array
# nothing is interpreted as "It's not needed"
_get_temp_shape(mc, model, li) = _get_shape(mc, model, li)
_get_temp_shape(mc, model, ::Type{<: Sum}) = 1
function _get_temp_shape(mc, model, ::LatticeIterationWrapper{LI}) where {LI}
    _get_shape(mc, model, LI)
end

# final_shape refers to the shape of what the observable saves
# here `nothing` means saving the eltype instead of an array
_get_final_shape(mc, model, li) = _get_shape(mc, model, li)
_get_final_shape(mc, model, ::Type{<: Sum}) = nothing
_get_final_shape(mc, model, ::SuperfluidDensity) = nothing
function _get_final_shape(mc, model, s::ApplySymmetries{LI, N}) where {LI, N}
    if LI <: EachLocalQuadByDistance || LI <: EachLocalQuadBySyncedDistance
        shape = _get_shape(mc, model, LI)
        (first(shape), N)
    else
        throw(MethodError(_get_final_shape, (mc, model, s)))
    end
end

_get_shape(model, mask::RawMask) = (mask.nsites, mask.nsites)
_get_shape(model, mask::DistanceMask) = length(mask)
_get_shape(mc, model, ::Nothing) = nothing

_get_shape(mc, model, LI::Type) = _get_shape(LI(mc, model))
_get_shape(mc, model, ::Type{Nothing}) = nothing
_get_shape(mc, model, ::Type{EachSite}) = length(lattice(model))
_get_shape(mc, model, ::Type{EachSiteAndFlavor}) = nflavors(model) * length(lattice(model))
_get_shape(mc, model, ::Type{EachSitePair}) = (length(lattice(model)), length(lattice(model)))
_get_shape(iter::DeferredLatticeIterator) = ndirections(iter)

# Some type piracy to make things easier
Base.Nothing(::DQMC, ::Model) = nothing

# To identify the requirement of equal-time Greens functions
struct Greens <: AbstractGreensIterator end
Greens(::DQMC, ::Model)= Greens()

struct GreensAt <: AbstractUnequalTimeGreensIterator
    k::Int
    l::Int
end
GreensAt(l::Integer) = GreensAt(l, l)
GreensAt(k::Integer, l::Integer) = GreensAt(k, l)
# GreensAt{k, l}(::DQMC, ::Model) where {k, l} = GreensAt(k, l)
Base.:(==)(a::GreensAt, b::GreensAt) = (a.k == b.k) && (a.l == b.l)

# maybe we want to some stuff to get custom indexing?
# How would this integrate with greens iterators? maybe:
# greens(mc::DQMC, model::Model) = greens(mc)

requires(::AbstractMeasurement) = (Nothing, Nothing)
requires(m::DQMCMeasurement) = (m.greens_iterator, m.lattice_iterator)
# TODO ApplySymmetries{LI}

@bm function generate_groups(mc, model, measurements)
    # maybe instead:
    requirements = requires.(measurements)
    GIs = unique(first.(requirements))
    LIs = unique(last.(requirements))
    lattice_iterators = map(T -> T(mc, model), LIs)

    if any(x -> (x isa Type ? x : typeof(x)) <: AbstractUnequalTimeGreensIterator, GIs)
        initialize_stack(mc, mc.ut_stack)
    end

    # greens_iterator => [
    #     (lattice_iterator, measurement),
    #     (lattice_iterator, measurement),
    #     ...
    # ]
    output = map(enumerate(GIs)) do (i, G)
        ms = filter(m -> G == requires(m)[1], measurements)
        group = map(ms) do m
            LI = requires(m)[2]
            j = findfirst(==(LI), LIs)
            @assert j !== nothing
            (lattice_iterators[j], m)
        end
        (G isa Type ? G(mc, model) : G) => group
    end

    if length(measurements) != mapreduce(x -> length(x[2]), +, output, init = 0)
        for (G, group) in output
            println(G)
            for (li, m) in group
                println("\t", typeof(li), typeof(m.kernel))
            end
        end
        N = length(measurements)
        M = mapreduce(x -> length(x[2]), +, output, init = 0)
        error("Oh no. We lost some measurements. $N-> $M")
    end

    return output
end

lattice_iterator(m::DQMCMeasurement, mc, model) = m.lattice_iterator(mc, model)

# TODO
# Saving the kernel function as a symbol is kinda risky because the function
# definition is not guaranteed to be the same
# though maybe that's a good thing - changing function definitions doesn't
# break stuff this way
function _save(file::JLDFile, m::DQMCMeasurement, key::String)
    write(file, "$key/VERSION", 1)
    write(file, "$key/tag", "DQMCMeasurement")
    write(file, "$key/GI", m.greens_iterator)
    write(file, "$key/LI", m.lattice_iterator)
    # maybe add module for eval?
    write(file, "$key/kernel", Symbol(m.kernel))
    write(file, "$key/obs", m.observable)
    write(file, "$key/temp", m.temp)
end

function _load(data, ::Val{:DQMCMeasurement})
    temp = haskey(data, "temp") ? data["temp"] : data["output"]

    kernel = try
        eval(data["kernel"])
    catch e
        @warn "Failed to load kernel in module MonteCarlo." exception=e
        missing_kernel
    end
    DQMCMeasurement(data["GI"], data["LI"], kernel, data["obs"], temp)
end

to_tag(::Type{<: DQMCMeasurement}) = Val(:DQMCMeasurement)

################################################################################
### Greens function related
################################################################################

@bm function apply!(::Nothing, combined::Vector{<: Tuple}, mc::DQMC, model, sweep)
    for (lattice_iterator, measurement) in combined
        # Clear temp if necessary
        prepare!(lattice_iterator, model, measurement)
        # Write measurement to ouput
        measure!(lattice_iterator, measurement, mc, model, sweep)
        # Finalize computation (temp) and commit
        finish!(lattice_iterator, model, measurement)
    end

    nothing
end

@bm function apply!(::Greens, combined::Vector{<: Tuple}, mc::DQMC, model, sweep)
    G = greens!(mc)
    for (lattice_iterator, measurement) in combined
        prepare!(lattice_iterator, model, measurement)
        measure!(lattice_iterator, measurement, mc, model, sweep, G)
        finish!(lattice_iterator, model, measurement)
    end
    nothing
end

@bm function apply!(g::GreensAt, combined::Vector{<: Tuple}, mc::DQMC, model, sweep)
    G = greens!(mc, g.k, g.l)
    for (lattice_iterator, measurement) in combined
        prepare!(lattice_iterator, model, measurement)
        measure!(lattice_iterator, measurement, mc, model, sweep, G)
        finish!(lattice_iterator, model, measurement)
    end
    nothing
end

@bm function apply!(iter::CombinedGreensIterator, combined::Vector{<: Tuple}, mc::DQMC, model, sweep)
    for (lattice_iterator, measurement) in combined
        prepare!(lattice_iterator, model, measurement)
    end

    G00 = greens!(mc)
    for (G0l, Gl0, Gll) in init(mc, iter)
        for (lattice_iterator, measurement) in combined
            measure!(lattice_iterator, measurement, mc, model, sweep, (G00, G0l, Gl0, Gll))
        end
    end

    for (lattice_iterator, measurement) in combined
        finish!(lattice_iterator, model, measurement, mc.parameters.delta_tau)
    end
    nothing
end

@bm function apply!(iter::AbstractGreensIterator, combined::Vector{<: Tuple}, mc::DQMC, model, sweep)
    for (lattice_iterator, measurement) in combined
        prepare!(lattice_iterator, model, measurement)
    end

    for packed_greens in init(mc, iter)
        for (lattice_iterator, measurement) in combined
            measure!(lattice_iterator, measurement, mc, model, sweep, packed_greens)
        end
    end

    for (lattice_iterator, measurement) in combined
        finish!(lattice_iterator, model, measurement, 1.0)
    end
    nothing
end

################################################################################
### measure!
################################################################################

@bm function measure!(lattice_iterator, measurement, mc::DQMC, model, sweep, packed_greens)
    # ignore sweep
    apply!(lattice_iterator, measurement, mc, model, packed_greens)
    nothing
end

# Lattice irrelevant
@bm function measure!(::Nothing, measurement, mc::DQMC, model, sweep, packed_greens)
    push!(measurement.observable, measurement.kernel(mc, model, packed_greens))
    nothing
end

################################################################################
### apply Lattice Iterators
################################################################################

# Call kernel for each site (linear index)
@bm function apply!(iter::DirectLatticeIterator, measurement, mc::DQMC, model, packed_greens)
    for i in iter
        measurement.temp[i] += measurement.kernel(mc, model, i, packed_greens)
    end
    nothing
end

# Call kernel for each pair (src, trg) (Nsties² total)
@bm function apply!(iter::EachSitePair, measurement, mc::DQMC, model, packed_greens)
    for (i, j) in iter
        measurement.temp[i, j] += measurement.kernel(mc, model, (i, j), packed_greens)
    end
    nothing
end

# Call kernel for each pair (site, site) (i.e. on-site)
@bm function apply!(iter::OnSite, measurement, mc::DQMC, model, packed_greens)
    for (i, j) in iter
        measurement.temp[i] += measurement.kernel(mc, model, (i, j), packed_greens)
    end
    nothing
end

@bm function apply!(iter::DeferredLatticeIterator, measurement, mc::DQMC, model, packed_greens)
    @inbounds for idxs in iter
        measurement.temp[first(idxs)] += measurement.kernel(mc, model, idxs[2:end], packed_greens)
    end
    nothing
end

# Sums
@bm function apply!(iter::Sum{<: DirectLatticeIterator}, measurement, mc::DQMC, model, packed_greens)
    @inbounds for idxs in iter
        measurement.temp[1] += measurement.kernel(mc, model, idxs, packed_greens)
    end
    nothing
end
@bm function apply!(iter::Sum{<: DeferredLatticeIterator}, measurement, mc::DQMC, model, packed_greens)
    @inbounds for idxs in iter
        measurement.temp[1] += measurement.kernel(mc, model, idxs[2:end], packed_greens)
    end
    nothing
end

@inline apply!(s::LatticeIterationWrapper, m, mc, model, pg) = apply!(s.iter, m, mc, model, pg)

################################################################################
### LatticeIterator preparation and finalization
################################################################################

# If LatticeIterator is Nothing, then things should be handled in measure!
@inline prepare!(::Nothing, model, m) = nothing
@inline prepare!(::AbstractLatticeIterator, model, m) = m.temp .= zero(eltype(m.temp))
@inline prepare!(s::Sum, args...) = prepare!(s.iter, args...)

@inline finish!(::Nothing, args...) = nothing # handled in measure!
@inline function finish!(li, model, m, factor=1.0)
    finalize_temp!(li, model, m, factor)
    commit!(li, m)
end

@inline function finalize_temp!(::AbstractLatticeIterator, model, m, factor)
    m.temp .*= factor
end
@inline function finalize_temp!(::DeferredLatticeIterator, model, m, factor)
    m.temp .*= factor / length(lattice(model))
end
@inline function finalize_temp!(s::LatticeIterationWrapper, model, m , factor)
    finalize_temp!(s.iter, model, m, factor)
end

@inline commit!(::AbstractLatticeIterator, m) = push!(m.observable, m.temp)
@inline commit!(::Sum, m) = push!(m.observable, m.temp[1])

function commit!(s::_ApplySymmetries{<: EachLocalQuadByDistance}, m)
    final = zeros(eltype(m.temp), size(m.temp, 1), length(s.symmetries))
    # This calculates
    # ∑_{a a'} O(Δr, a, a') f_ζ(a) f_ζ(a')
    # where a, a' are typically nearest neighbor directions and
    # f_ζ(a) is the weight in direction a for a symmetry ζ
    for (i, sym) in enumerate(s.symmetries)
        for k in 1:length(sym), l in 1:length(sym)
            @. final[:, i] += m.temp[:, k, l] * sym[k] * sym[l]
        end
    end
    push!(m.observable, final)
end
function commit!(s::_ApplySymmetries{<: EachLocalQuadBySyncedDistance}, m)
    final = zeros(eltype(m.temp), size(m.temp, 1), length(s.symmetries))
    # Same as above but with a = a'
    for (i, sym) in enumerate(s.symmetries)
        for k in 1:length(sym)
            @. final[:, i] += m.temp[:, k] * sym[k] * sym[k]
        end
    end
    push!(m.observable, final)
end

function commit!(s::_SuperfluidDensity{<: EachLocalQuadBySyncedDistance}, m)
    final = zero(eltype(m.observable))
    for (shift, long, trans) in zip(s.dir_idxs, s.long_qs, s.trans_qs)
        for (i, dir) in enumerate(s.dirs)
            final += m.temp[i, shift] * (cis(-dot(dir, long)) - cis(-dot(dir, trans)))
        end
    end
    push!(m.observable, final)
end


text:
using Cairo
using CairoScript

using Compat, Colors
import Compat.String

@compat import Base.show

using Test

include("test_painting.jl")

pkg_dir = dirname(dirname(@__FILE__))

# Image Surface
@testset "Interpreter Run" begin

    # run a script that includes a writing operation
    surf = CairoImageSurface(256, 256, Cairo.FORMAT_ARGB32)

    GC.@preserve surf h c ci begin

    testfile = joinpath(pkg_dir,"data","a1.cs");

    h = CairoScript.InterpreterHooks()
    h.surface_create = CairoScript.surf_create_c[]
    h.closure = surf.ptr
    ci = CairoScript.Interpreter()
    c = CairoScript.interpreter_install_hooks(ci,h)

    status = CairoScript.interpreter_run(c,testfile)

    @test status == 0
    finish(surf)
    #surf.ptr = C_NULL;
    end

    outputfile = "out.png"
    @test isfile(outputfile)
    rm(outputfile)

    # run a script that just does painting
    surf = CairoImageSurface(256, 256, Cairo.FORMAT_ARGB32)
    testfile = joinpath(pkg_dir,"data","a2.cs");

    h = CairoScript.InterpreterHooks()
    h.surface_create = CairoScript.surf_create_c[]
    h.closure = surf.ptr

    c = CairoScript.Interpreter()
    c = CairoScript.interpreter_install_hooks(c,h)

    GC.@preserve c begin
        status = CairoScript.interpreter_run(c,testfile)
    end

    @test status == 0

    outputfile = "out.png"
    Cairo.write_to_png(surf,outputfile)
    finish(surf)
    surf.ptr = C_NULL;

    @test isfile(outputfile)
    rm(outputfile)

    # script from string
    surf = CairoImageSurface(256, 256, Cairo.FORMAT_ARGB32)
    testdata = """
%!CairoScript
<< /content //COLOR_ALPHA /width 400 /height 300 >> surface context
0 0 1 rgb set-source
n 128 25.601562 m 230.398438 230.398438 l 128 230.398438 l 51.199219 230.398438 51.199219 128 128 128 c h 64 25.601562 m 115.199219 76.800781 l 64 128 l 12.800781 76.800781 l h 64 25.601562 m
fill+
0 g set-source
10 set-line-width
stroke
"""

    h = CairoScript.InterpreterHooks()
    h.surface_create = CairoScript.surf_create_c[]
    h.closure = surf.ptr
    c = CairoScript.Interpreter()
    c = CairoScript.interpreter_install_hooks(c,h)
    GC.@preserve c begin
        status = CairoScript.interpreter_feed_string(c,testdata)
    end

    #display(testdata)
    @test status == 0

    outputfile = "out.png"
    Cairo.write_to_png(surf,outputfile)
    finish(surf)
    surf.ptr = C_NULL;

    @test isfile(outputfile)
    rm(outputfile)

end

@testset "Content Test" begin

    # test if surface contains color
    surf = CairoImageSurface(256, 256, Cairo.FORMAT_ARGB32)
    testfile = joinpath(pkg_dir,"data","a2.cs");

    h = CairoScript.InterpreterHooks()
    h.surface_create = CairoScript.surf_create_c[]
    h.closure = surf.ptr
    c = CairoScript.Interpreter()
    c = CairoScript.interpreter_install_hooks(c,h)

    GC.@preserve c begin
        status = CairoScript.interpreter_run(c,testfile)
    end

    @test status == 0

    d = simple_hist(matrix_read(surf))

    @test length(d) == 1
    @test collect(keys(d))[1] == 0xffffff00

    finish(surf)
    surf.ptr = C_NULL;

    # short setup, same data and test
    surf = CairoImageSurface(256, 256, Cairo.FORMAT_ARGB32)
    testfile = joinpath(pkg_dir,"data","a2.cs");
    h = CairoScript.InterpreterHooks(
            closure = surf.ptr, surface_create = CairoScript.surf_create_c[])

    c = CairoScript.Interpreter()
    c = CairoScript.interpreter_install_hooks(c,h)

    GC.@preserve c begin
        status = CairoScript.interpreter_run(c,testfile)
    end

    @test status == 0

    d = simple_hist(matrix_read(surf))

    @test length(d) == 1
    @test collect(keys(d))[1] == 0xffffff00

    finish(surf)

    surf.ptr = C_NULL;

end


text:
function read_puzzle(fname)
    lines = readlines("data/$fname")
    npieces = length(lines)
    puzzle = zeros(Int, (npieces, 5))
    for i in 1:npieces
        puzzle[i, 1] = i
        parts = split(lines[i], " ")
        puzzle[i, 2:end] = parse.(Int, parts)
    end
    return puzzle
end

function get_rotations(puzzle)
    npieces = size(puzzle)[1]
    rotations = zeros(Int, (npieces * 4, 5))
    rotation_indices = [[2, 3, 4, 5], [3, 4, 5, 2], [4, 5, 2, 3], [5, 2, 3, 4]]
    for i in 1:npieces
        j = 1
        for rotation in rotation_indices
            rotations[(i - 1) * 4 + j, 1] = i
            rotations[(i - 1) * 4 + j, 2:end] = puzzle[i, rotation]
            j += 1
        end
    end

    return rotations
end

@testset "Eternity 6x5" begin
    width = 6
    height = 5
    puzzle = read_puzzle("eternity_6x5")
    rotations = get_rotations(puzzle)
    npieces = size(puzzle)[1]

    ncolors = maximum(puzzle[:, 2:end])

    m = Model(optimizer_with_attributes(CS.Optimizer, "keep_logs" => true, "logging" => []))
    @variable(m, 1 <= p[1:height, 1:width] <= npieces, Int)
    @variable(m, 0 <= pu[1:height, 1:width] <= ncolors, Int)
    @variable(m, 0 <= pr[1:height, 1:width] <= ncolors, Int)
    @variable(m, 0 <= pd[1:height, 1:width] <= ncolors, Int)
    @variable(m, 0 <= pl[1:height, 1:width] <= ncolors, Int)

    @constraint(m, p[:] in CS.AllDifferentSet())
    for i in 1:height, j in 1:width
        @constraint(
            m,
            [p[i, j], pu[i, j], pr[i, j], pd[i, j], pl[i, j]] in CS.TableSet(rotations)
        )
    end

    # borders
    # up and down
    for j in 1:width
        @constraint(m, pu[1, j] == 0)
        @constraint(m, pd[height, j] == 0)

        if j != width
            @constraint(m, pr[1, j] == pl[1, j + 1])
            @constraint(m, pr[height, j] == pl[height, j + 1])
        end
    end

    # right and left
    for i in 1:height
        @constraint(m, pl[i, 1] == 0)
        @constraint(m, pr[i, width] == 0)

        if i != height
            @constraint(m, pd[i, 1] == pu[i + 1, 1])
            @constraint(m, pd[i, width] == pu[i + 1, width])
        end
    end

    for i in 1:(height - 1), j in 1:(width - 1)
        @constraint(m, pd[i, j] == pu[i + 1, j])
        @constraint(m, pr[i, j] == pl[i, j + 1])
    end

    if width == height
        start_piece = findfirst(i -> count(c -> c == 0, puzzle[i, :]) == 2, 1:npieces)
        @constraint(m, p[1, 1] == start_piece)
    end

    optimize!(m)
    com = CS.get_inner_model(m)
    @test is_solved(com)

    status = JuMP.termination_status(m)
    @test status == MOI.OPTIMAL

    @test general_tree_test(com)
end

@testset "Eternity 5x5 Indicator" begin
    width = 5
    height = 5
    puzzle = read_puzzle("eternity_5x5")
    rotations = get_rotations(puzzle)
    npieces = size(puzzle)[1]

    ncolors = maximum(puzzle[:, 2:end])

    m = Model(optimizer_with_attributes(CS.Optimizer, "keep_logs" => true, "logging" => []))
    @variable(m, 1 <= p[1:height, 1:width] <= npieces, Int)
    @variable(m, 0 <= pu[1:height, 1:width] <= ncolors, Int)
    @variable(m, 0 <= pr[1:height, 1:width] <= ncolors, Int)
    @variable(m, 0 <= pd[1:height, 1:width] <= ncolors, Int)
    @variable(m, 0 <= pl[1:height, 1:width] <= ncolors, Int)
    @variable(m, indicator, Bin)

    @constraint(m, p[:] in CS.AllDifferentSet())
    for i in 1:height, j in 1:width
        @constraint(
            m,
            indicator => {
                [p[i, j], pu[i, j], pr[i, j], pd[i, j], pl[i, j]] in CS.TableSet(rotations),
            }
        )
    end

    # borders
    # up and down
    for j in 1:width
        @constraint(m, pu[1, j] == 0)
        @constraint(m, pd[height, j] == 0)

        if j != width
            @constraint(m, pr[1, j] == pl[1, j + 1])
            @constraint(m, pr[height, j] == pl[height, j + 1])
        end
    end

    # right and left
    for i in 1:height
        @constraint(m, pl[i, 1] == 0)
        @constraint(m, pr[i, width] == 0)

        if i != height
            @constraint(m, pd[i, 1] == pu[i + 1, 1])
            @constraint(m, pd[i, width] == pu[i + 1, width])
        end
    end

    for i in 1:(height - 1), j in 1:(width - 1)
        @constraint(m, pd[i, j] == pu[i + 1, j])
        @constraint(m, pr[i, j] == pl[i, j + 1])
    end

    if width == height
        start_piece = findfirst(i -> count(c -> c == 0, puzzle[i, :]) == 2, 1:npieces)
        @constraint(m, p[1, 1] == start_piece)
    end

    @objective(m, Max, indicator)

    optimize!(m)
    com = CS.get_inner_model(m)
    # check that indicator is set to 1
    @test JuMP.objective_value(m) ≈ 1.0
    @test is_solved(com)

    status = JuMP.termination_status(m)
    @test status == MOI.OPTIMAL

    @test general_tree_test(com)
end

@testset "Eternity 5x5 Reified" begin
    width = 5
    height = 5
    puzzle = read_puzzle("eternity_5x5")
    rotations = get_rotations(puzzle)
    npieces = size(puzzle)[1]

    ncolors = maximum(puzzle[:, 2:end])

    m = Model(optimizer_with_attributes(CS.Optimizer, "keep_logs" => true, "logging" => []))
    @variable(m, 1 <= p[1:height, 1:width] <= npieces, Int)
    @variable(m, 0 <= pu[1:height, 1:width] <= ncolors, Int)
    @variable(m, 0 <= pr[1:height, 1:width] <= ncolors, Int)
    @variable(m, 0 <= pd[1:height, 1:width] <= ncolors, Int)
    @variable(m, 0 <= pl[1:height, 1:width] <= ncolors, Int)
    @variable(m, reified[1:height, 1:width], Bin)

    @constraint(m, p[:] in CS.AllDifferentSet())
    for i in 1:height, j in 1:width
        @constraint(
            m,
            reified[i, j] := {
                [p[i, j], pu[i, j], pr[i, j], pd[i, j], pl[i, j]] in CS.TableSet(rotations),
            }
        )
    end

    # borders
    # up and down
    for j in 1:width
        @constraint(m, pu[1, j] == 0)
        @constraint(m, pd[height, j] == 0)

        if j != width
            @constraint(m, pr[1, j] == pl[1, j + 1])
            @constraint(m, pr[height, j] == pl[height, j + 1])
        end
    end

    # right and left
    for i in 1:height
        @constraint(m, pl[i, 1] == 0)
        @constraint(m, pr[i, width] == 0)

        if i != height
            @constraint(m, pd[i, 1] == pu[i + 1, 1])
            @constraint(m, pd[i, width] == pu[i + 1, width])
        end
    end

    for i in 1:(height - 1), j in 1:(width - 1)
        @constraint(m, pd[i, j] == pu[i + 1, j])
        @constraint(m, pr[i, j] == pl[i, j + 1])
    end

    if width == height
        start_piece = findfirst(i -> count(c -> c == 0, puzzle[i, :]) == 2, 1:npieces)
        @constraint(m, p[1, 1] == start_piece)
    end

    @objective(m, Max, sum(reified))

    optimize!(m)
    com = CS.get_inner_model(m)
    # check that all reified variables are set to 1
    @test JuMP.objective_value(m) ≈ width * height
    @test is_solved(com)

    status = JuMP.termination_status(m)
    @test status == MOI.OPTIMAL

    @test general_tree_test(com)
end


text:
using LinearAlgebra
using Plots
using IJulia
using Printf

function meshgrid(x, y)
    x = reshape(x, 1, length(x))
    y = reshape(y, length(y), 1)
    X = repeat(x, length(y), 1)
    Y = repeat(y, 1, length(x))
    return X, Y
end

# choose y label to be either y or ϕ, depending on the geometry
function plot_basic_state(grid, background, geometry, file)

    if geometry == Cartesian()
        y = grid.y/1e3
        xlabel = "y (km)"
    elseif geometry == Spherical()
        y = rad2deg.(grid.ϕ)
        xlabel = "ϕ (°)"
    end

    solution = background[geometry]

     p1 = plot(
         y,
         solution.u,
         linewidth = 3,
         label = "u",
         title = "Velocity",
         ylabel = "(m/s)",
         xlims = (y[end], y[1])
         )

     p2 = plot(
         y,
         solution.η,
         linewidth = 3,
         label = "η",
         xlabel = xlabel,
         ylabel = "(m)",
         title = "Free-Surface",
         xlims = (y[end], y[1])
         )

     plt1 = plot(p1, p2, layout = (2,1))
     savefig(plt1, file)
end

function plot_growth_rates(ks, σ, phys, geometry, Nmodes, file)

    if geometry == Cartesian()
        k = ks[geometry]*1e3
        xlabel = "k (1/km)"
    elseif geometry == Spherical()
        k = ks[geometry]/1e3
        xlabel = "k (10^3/rad)"
    end

    titles = ["1st mode", "2nd mode"]
    plt = plot()
    for (index,value) in enumerate(1:Nmodes)
        σnormalized = [σ[(k, geometry)][index] for k in ks[geometry]]/phys[geometry].f₀
        plot!(plt,
        k,
        σnormalized,
        linewidth = 3,
        label=titles[index],
        xlabel=xlabel,
        ylabel="σ/f₀",
        title="Growth Rates",
        xlims=(k[1], k[end])
        )
    end
    savefig(plt,file)

end

#ticks = collect(0:0.2:1)
#ticklabels = [ @sprintf("%5.1f",x) for x in ticks ]
#plot(x,y)
#plot!(xticks=(ticks,ticklabels))

function  make_tick_labels(fs, Nfs)
    fs_max = maximum([real(fs); imag(fs)])
    fs_min = minimum([real(fs); imag(fs)])
    dfs = (fs_max - fs_min)/Nfs

    ticks = collect(fs_min:dfs:fs_max)
    ticklabels = [ @sprintf("%f",f) for f in ticks]

    return ticks, ticklabels
end

function plot_1D_streamfunction(k_index, k, phys, σmodes, geometry, mode_number, file)

    if geometry == Cartesian()
        y = phys[geometry].grid.y/1e3
        xlabel = "y (km)"
    elseif geometry == Spherical()
        y = rad2deg.(phys[geometry].grid.ϕ)
        xlabel = "ϕ (°)"
    end

    N = phys[geometry].grid.N
    Lx = 2*pi/k;
    Nx = N;
     x = LinRange(0, Lx, Nx+1)
  X, Y = meshgrid(x,y)

  uvec = σmodes[(k,geometry)][    1:N+1,   mode_number];
  vvec = σmodes[(k,geometry)][  N+2:2*N+2, mode_number];
  ηvec = σmodes[(k,geometry)][2*N+1:3*N+1, mode_number];

    p1 = plot(
      y,
      real(uvec),
      linewidth = 3,
      label = "real",
      title = "u",
      yticks=make_tick_labels(uvec, 3),
      xlims = (y[end], y[1])
      )
  plot!(
      p1,
      y,
      linewidth = 3,
      imag(uvec),
      label = "imag")
  p2 = plot(
      y,
      real(vvec),
      linewidth = 3,
      label = "real",
      title = "v",
      yticks=make_tick_labels(vvec, 3),
       xlims = (y[end], y[1])
      )
  plot!(
      p2,
      y,
      linewidth = 3,
      imag(vvec),
      label = "imag")
  p3 = plot(
      y,
      real(ηvec),
      linewidth = 3,
      label = "real",
      title = "η",
      yticks=make_tick_labels(ηvec, 3),
      xlabel = xlabel,
      xlims = (y[end], y[1])
      )
  plot!(
      p3,
      y,
      linewidth = 3,
      imag(ηvec),
      label = "imag")
  plt = plot(p1, p2, p3, layout = (3,1))
  savefig(plt, file)
  #display(plt)

end

function plot_2D_streamfunction(k_index, k, phys, σmodes, geometry, mode_number, file)

    N = phys[geometry].grid.N
    Nx = N;
    Lx = 2*pi/k;
    x = LinRange(0, Lx, Nx+1)

    if geometry == Cartesian()
        y = phys[geometry].grid.y/1e3
        ylabel = "y (km)"
        x = x./1e3
        kscaled = k*1e3
        xlabel = "x (km)"
    elseif geometry == Spherical()
        y = rad2deg.(phys[geometry].grid.ϕ)
        ylabel = "ϕ (°)"
        x = rad2deg.(x)
        kscaled = k * π / 180
        xlabel = "λ (°)"
     end

  X, Y = meshgrid(x,y)

  uvec = σmodes[(k,geometry)][    1:N+1,   mode_number];
  vvec = σmodes[(k,geometry)][  N+2:2*N+2, mode_number];
  ηvec = σmodes[(k,geometry)][2*N+1:3*N+1, mode_number];

    u =     repeat(real(uvec),1,Nx+1).*cos.(kscaled*X) -    repeat(imag(uvec),1,Nx+1).*sin.(kscaled*X);
    v = -k.*repeat(imag(vvec),1,Nx+1).*cos.(kscaled*X) - k.*repeat(real(vvec),1,Nx+1).*sin.(kscaled*X);
    η =     repeat(real(ηvec),1,Nx+1).*cos.(kscaled*X) -    repeat(imag(ηvec),1,Nx+1).*sin.(kscaled*X);

    kwargs = (
        xlabel = "x (km)",
        ylabel = ylabel,
          fill = true,
        levels = 20,
     linewidth = 0,
         color = :balance,
      colorbar = true,
          xlim = (x[1],   x[end]),
          ylim = (y[end], y[1])
    )

    #print("x ticks = ", make_tick_labels(x, 2), "\n")

    u_plt = contour(x, y[end:-1:1], u[:, end:-1:1], title="u"; kwargs...)
    v_plt = contour(x, y[end:-1:1], v[:, end:-1:1], title="v"; kwargs...)
    η_plt = contour(x, y[end:-1:1], η[:, end:-1:1], title="η"; kwargs...)

    plt1 = plot(u_plt, v_plt, η_plt, layout = (1,3))
    savefig(plt1, file)

end

function plot_2D_vorticity(phys, k_index, k, σmodes, geometry, mode_number, file)

    N = phys[geometry].grid.N
    Lx = 2*pi/k;
    Nx = N;
    x = LinRange(0, Lx, Nx+1)

    if geometry == Cartesian()
        y = phys[geometry].grid.y/1e3
        ylabel = "y (km)"
        Dy = phys[geometry].grid.D
        x = x./1e3
        kscaled = k*1e3
        xlabel = "x (km)"
    elseif geometry == Spherical()
        y = rad2deg.(phys[geometry].grid.ϕ)
        ylabel = "ϕ (°)"
        Dy = phys[geometry].grid.D/phys[geometry].grid.a
        x = rad2deg.(x)
        kscaled = k * π / 180
        xlabel = "λ (°)"
    end

    X, Y = meshgrid(x,y)

  uvec = σmodes[(k,geometry)][    1:N+1,   mode_number];
  vvec = σmodes[(k,geometry)][  N+2:2*N+2, mode_number];

    ζvec = im * k * vvec - Dy * uvec
    ζ    = repeat(real(ζvec),1,Nx+1).*cos.(kscaled*X) - repeat(imag(ζvec),1,Nx+1).*sin.(kscaled*X);

    kwargs = (
        xlabel = xlabel,
        ylabel = ylabel,
          fill = true,
        levels = 20,
     linewidth = 0,
         color = :balance,
      colorbar = true,
          xlim = (  x[1], x[end]),
         ylim = (y[end],  y[1])
    )

    plt = contour(x, y[end:-1:1], ζ[:, end:-1:1], title="ζ"; kwargs...)

    savefig(plt, file)

end

#string.(fieldnames(typeof(gridC)))


text:
__precompile__()
module TSSM

push!(LOAD_PATH, dirname(@__FILE__))

import Base.copy!
import LinearAlgebra.norm
import LinearAlgebra.scale!
import LinearAlgebra.axpy!
import LinearAlgebra.normalize!

export WaveFunction, WaveFunction1D, WaveFunction2D, WaveFunction3D
export WaveFunctionComplex, WaveFunctionComplex1D, WaveFunctionComplex2D, WaveFunctionComplex3D
export WaveFunctionReal, WaveFunctionReal1D, WaveFunctionReal2D, WaveFunctionReal3D

export TimeSplittingSpectralMethod, TimeSplittingSpectralMethod1D, TimeSplittingSpectralMethod2D, TimeSplittingSpectralMethod3D
export TimeSplittingSpectralMethodComplex, TimeSplittingSpectralMethodComplex1D, TimeSplittingSpectralMethodComplex2D, TimeSplittingSpectralMethodComplex3D
export TimeSplittingSpectralMethodReal, TimeSplittingSpectralMethodReal1D, TimeSplittingSpectralMethodReal2D, TimeSplittingSpectralMethodReal3D

export dim, is_real

export periodic, dirichlet, neumann

export Fourier1D, WfFourier1D
export Fourier2D, WfFourier2D
export Fourier3D, WfFourier3D
export FourierReal1D, WfFourierReal1D
export FourierReal2D, WfFourierReal2D
export FourierReal3D, WfFourierReal3D

export Schroedinger1D, WfSchroedinger1D
export Schroedinger2D, WfSchroedinger2D
export Schroedinger3D, WfSchroedinger3D
export SchroedingerReal1D, WfSchroedingerReal1D
export SchroedingerReal2D, WfSchroedingerReal2D
export SchroedingerReal3D, WfSchroedingerReal3D

export SchroedingerRotating2D, WfSchroedingerRotating2D
export SchroedingerRotating3D, WfSchroedingerRotating3D
export SchroedingerRotatingReal2D, WfSchroedingerRotatingReal2D
export SchroedingerRotatingReal3D, WfSchroedingerRotatingReal3D

export SchroedingerHermite1D, WfSchroedingerHermite1D
export SchroedingerHermite2D, WfSchroedingerHermite2D
export SchroedingerHermite3D, WfSchroedingerHermite3D
export SchroedingerHermiteReal1D, WfSchroedingerHermiteReal1D
export SchroedingerHermiteReal2D, WfSchroedingerHermiteReal2D
export SchroedingerHermiteReal3D, WfSchroedingerHermiteReal3D

export SchroedingerGeneralizedLaguerre2D, WfSchroedingerGeneralizedLaguerre2D
export SchroedingerGeneralizedLaguerreHermite3D, WfSchroedingerGeneralizedLaguerreHermite3D
export SchroedingerGeneralizedLaguerreReal2D, WfSchroedingerGeneralizedLaguerreReal2D
export SchroedingerGeneralizedLaguerreHermiteReal3D, WfSchroedingerGeneralizedLaguerreHermiteReal3D

export gauss, radau, lobatto

export FourierBessel2D, WfFourierBessel2D
export FourierBesselReal2D, WfFourierBesselReal2D
export BesselRotSym1D, WfBesselRotSym1D
export BesselRotSymReal1D, WfBesselRotSymReal1D

export wave_function, clone
export is_real_space, is_frequency_space, to_real_space!, to_frequency_space!
export is_real_space_x, is_frequency_space_x, to_real_space_x!, to_frequency_space_x!
export is_real_space_y, is_frequency_space_y, to_real_space_y!, to_frequency_space_y!
export is_real_space_z, is_frequency_space_z, to_real_space_z!, to_frequency_space_z!
export set_time!, get_time, propagate_time!
export set_propagate_time_together_with_A!, get_propagate_time_together_with_A
export propagate_A!, propagate_B!, propagate_C!, add_apply_A!, add_apply_B!, add_apply_C!
export add_phi_A!
export propagate_A_derivative!, propagate_B_derivative!, propagate_C_derivative!
export norm_in_frequency_space, distance
export inner_product, eigen_function!, evaluate
export save, load!, get_data, set!
export get_eigenvalues, get_nodes, get_weights, get_transformation_matrices
export get_nx, get_ny, get_nz
export get_xmin, get_xmax
export get_ymin, get_ymax
export get_zmin, get_zmax
export get_omega_x,  get_omega_y,  get_omega_z
export get_omega_r, get_Omega
export get_nr, get_nfr, get_ntheta

export get_hbar, get_mass, get_cubic_coupling, set_cubic_coupling
export set_potential!, get_potential, load_potential!, save_potential
export set_potential_t!, get_potential_t
export set_potential_t_derivative!, get_potential_t_derivative
export imaginary_time_propagate_A!, imaginary_time_propagate_B!
export kinetic_energy, potential_energy, interaction_energy, observable
export get_energy_expectation_deviation, get_realspace_observables
export kinetic_matrix_element, potential_matrix_element
export selfconsistent_nonlinear_step!

## abstract types ############################################################################
abstract type TimeSplittingSpectralMethod{T<:AbstractFloat} end
abstract type TimeSplittingSpectralMethodReal{T<:AbstractFloat}<:TimeSplittingSpectralMethod{T} end
abstract type TimeSplittingSpectralMethodReal1D{T<:AbstractFloat}<:TimeSplittingSpectralMethodReal{T} end
abstract type TimeSplittingSpectralMethodReal2D{T<:AbstractFloat}<:TimeSplittingSpectralMethodReal{T} end
abstract type TimeSplittingSpectralMethodReal3D{T<:AbstractFloat}<:TimeSplittingSpectralMethodReal{T} end
abstract type TimeSplittingSpectralMethodComplex{T<:AbstractFloat}<:TimeSplittingSpectralMethod{T} end
abstract type TimeSplittingSpectralMethodComplex1D{T<:AbstractFloat}<:TimeSplittingSpectralMethodComplex{T} end
abstract type TimeSplittingSpectralMethodComplex2D{T<:AbstractFloat}<:TimeSplittingSpectralMethodComplex{T} end
abstract type TimeSplittingSpectralMethodComplex3D{T<:AbstractFloat}<:TimeSplittingSpectralMethodComplex{T} end
TimeSplittingSpectralMethod1D = Union{TimeSplittingSpectralMethodReal1D, TimeSplittingSpectralMethodComplex1D}
TimeSplittingSpectralMethod2D = Union{TimeSplittingSpectralMethodReal2D, TimeSplittingSpectralMethodComplex2D}
TimeSplittingSpectralMethod3D = Union{TimeSplittingSpectralMethodReal3D, TimeSplittingSpectralMethodComplex3D}

abstract type WaveFunction{T<:AbstractFloat} end
abstract type WaveFunctionReal{T<:AbstractFloat}<:WaveFunction{T} end
abstract type WaveFunctionReal1D{T<:AbstractFloat}<:WaveFunctionReal{T} end
abstract type WaveFunctionReal2D{T<:AbstractFloat}<:WaveFunctionReal{T} end
abstract type WaveFunctionReal3D{T<:AbstractFloat}<:WaveFunctionReal{T} end
abstract type WaveFunctionComplex{T<:AbstractFloat}<:WaveFunction{T} end
abstract type WaveFunctionComplex1D{T<:AbstractFloat}<:WaveFunctionComplex{T} end
abstract type WaveFunctionComplex2D{T<:AbstractFloat}<:WaveFunctionComplex{T} end
abstract type WaveFunctionComplex3D{T<:AbstractFloat}<:WaveFunctionComplex{T} end
WaveFunction1D = Union{WaveFunctionReal1D, WaveFunctionComplex1D}
WaveFunction2D = Union{WaveFunctionReal2D, WaveFunctionComplex2D}
WaveFunction3D = Union{WaveFunctionReal3D, WaveFunctionComplex3D}

dim(psi::WaveFunction1D) = 1
dim(psi::WaveFunction2D) = 2
dim(psi::WaveFunction3D) = 3
dim(psi::TimeSplittingSpectralMethod1D) = 1
dim(psi::TimeSplittingSpectralMethod2D) = 2
dim(psi::TimeSplittingSpectralMethod3D) = 3

is_real(psi::WaveFunctionReal) = true
is_real(psi::WaveFunctionComplex) = false
is_real(psi::TimeSplittingSpectralMethodReal) = true
is_real(psi::TimeSplittingSpectralMethodComplex) = false

# Dummy propagators, to be overloaded
propagate_B!(psi::WaveFunction, dt::Number) = psi
propagate_C!(psi::WaveFunction, dt::Number) = psi

## Fourier types ############################################################################

mutable struct Fourier1D{T<:AbstractFloat} <: TimeSplittingSpectralMethodComplex1D{T}
    m::Ptr{Nothing}
end

mutable struct Fourier2D{T<:AbstractFloat} <: TimeSplittingSpectralMethodComplex2D{T}
    m::Ptr{Nothing}
end

mutable struct Fourier3D{T<:AbstractFloat} <: TimeSplittingSpectralMethodComplex3D{T}
    m::Ptr{Nothing}
end

mutable struct FourierReal1D{T<:AbstractFloat} <: TimeSplittingSpectralMethodReal1D{T}
    m::Ptr{Nothing}
end

mutable struct FourierReal2D{T<:AbstractFloat} <: TimeSplittingSpectralMethodReal2D{T}
    m::Ptr{Nothing}
end

mutable struct FourierReal3D{T<:AbstractFloat} <: TimeSplittingSpectralMethodReal3D{T}
    m::Ptr{Nothing}
end

mutable struct WfFourier1D{T<:AbstractFloat} <: WaveFunctionComplex1D{T}
    p::Ptr{Nothing}
    m::Fourier1D{T}
end

mutable struct WfFourier2D{T<:AbstractFloat} <: WaveFunctionComplex2D{T}
    p::Ptr{Nothing}
    m::Fourier2D{T}
end

mutable struct WfFourier3D{T<:AbstractFloat} <: WaveFunctionComplex3D{T}
    p::Ptr{Nothing}
    m::Fourier3D{T}
end

mutable struct WfFourierReal1D{T<:AbstractFloat} <: WaveFunctionReal1D{T}
    p::Ptr{Nothing}
    m::FourierReal1D{T}
end

mutable struct WfFourierReal2D{T<:AbstractFloat} <: WaveFunctionReal2D{T}
    p::Ptr{Nothing}
    m::FourierReal2D{T}
end

mutable struct WfFourierReal3D{T<:AbstractFloat} <: WaveFunctionReal3D{T}
    p::Ptr{Nothing}
    m::FourierReal3D{T}
end

## FourierBessel types ###########################################################################

mutable struct FourierBessel2D{T<:AbstractFloat} <: TimeSplittingSpectralMethodComplex2D{T}
    m::Ptr{Nothing}
end

mutable struct FourierBesselReal2D{T<:AbstractFloat} <: TimeSplittingSpectralMethodReal2D{T}
    m::Ptr{Nothing}
end

mutable struct BesselRotSym1D{T<:AbstractFloat} <: TimeSplittingSpectralMethodComplex1D{T}
    m::Ptr{Nothing}
end

mutable struct BesselRotSymReal1D{T<:AbstractFloat} <: TimeSplittingSpectralMethodReal1D{T}
    m::Ptr{Nothing}
end

mutable struct WfFourierBessel2D{T<:AbstractFloat} <: WaveFunctionComplex2D{T}
    p::Ptr{Nothing}
    m::FourierBessel2D{T}
end

mutable struct WfFourierBesselReal2D{T<:AbstractFloat} <: WaveFunctionReal2D{T}
    p::Ptr{Nothing}
    m::FourierBesselReal2D{T}
end

mutable struct WfBesselRotSym1D{T<:AbstractFloat} <: WaveFunctionComplex1D{T}
    p::Ptr{Nothing}
    m::BesselRotSym1D{T}
end

mutable struct WfBesselRotSymReal1D{T<:AbstractFloat} <: WaveFunctionReal1D{T}
    p::Ptr{Nothing}
    m::BesselRotSymReal1D{T}
end

## Schroedinger types ############################################################################

mutable struct Schroedinger1D{T<:AbstractFloat} <: TimeSplittingSpectralMethodComplex1D{T}
    m::Ptr{Nothing}
end

mutable struct Schroedinger2D{T<:AbstractFloat} <: TimeSplittingSpectralMethodComplex2D{T}
    m::Ptr{Nothing}
end

mutable struct Schroedinger3D{T<:AbstractFloat} <: TimeSplittingSpectralMethodComplex3D{T}
    m::Ptr{Nothing}
end

mutable struct SchroedingerReal1D{T<:AbstractFloat} <: TimeSplittingSpectralMethodReal1D{T}
    m::Ptr{Nothing}
end

mutable struct SchroedingerReal2D{T<:AbstractFloat} <: TimeSplittingSpectralMethodReal2D{T}
    m::Ptr{Nothing}
end

mutable struct SchroedingerReal3D{T<:AbstractFloat} <: TimeSplittingSpectralMethodReal3D{T}
    m::Ptr{Nothing}
end

mutable struct WfSchroedinger1D{T<:AbstractFloat} <: WaveFunctionComplex1D{T}
    p::Ptr{Nothing}
    m::Schroedinger1D{T}
end

mutable struct WfSchroedinger2D{T<:AbstractFloat} <: WaveFunctionComplex2D{T}
    p::Ptr{Nothing}
    m::Schroedinger2D{T}
end

mutable struct WfSchroedinger3D{T<:AbstractFloat} <: WaveFunctionComplex3D{T}
    p::Ptr{Nothing}
    m::Schroedinger3D{T}
end

mutable struct WfSchroedingerReal1D{T<:AbstractFloat} <: WaveFunctionReal1D{T}
    p::Ptr{Nothing}
    m::SchroedingerReal1D{T}
end

mutable struct WfSchroedingerReal2D{T<:AbstractFloat} <: WaveFunctionReal2D{T}
    p::Ptr{Nothing}
    m::SchroedingerReal2D{T}
end

mutable struct WfSchroedingerReal3D{T<:AbstractFloat} <: WaveFunctionReal3D{T}
    p::Ptr{Nothing}
    m::SchroedingerReal3D{T}
end

## SchroedingerRotating types ############################################################################

mutable struct SchroedingerRotating2D{T<:AbstractFloat} <: TimeSplittingSpectralMethodComplex2D{T}
    m::Ptr{Nothing}
end

mutable struct SchroedingerRotating3D{T<:AbstractFloat} <: TimeSplittingSpectralMethodComplex3D{T}
    m::Ptr{Nothing}
end

mutable struct WfSchroedingerRotating2D{T<:AbstractFloat} <: WaveFunctionComplex2D{T}
    p::Ptr{Nothing}
    m::SchroedingerRotating2D{T}
end

mutable struct WfSchroedingerRotating3D{T<:AbstractFloat} <: WaveFunctionComplex3D{T}
    p::Ptr{Nothing}
    m::SchroedingerRotating3D{T}
end

mutable struct SchroedingerRotatingReal2D{T<:AbstractFloat} <: TimeSplittingSpectralMethodReal2D{T}
    m::Ptr{Nothing}
end

mutable struct SchroedingerRotatingReal3D{T<:AbstractFloat} <: TimeSplittingSpectralMethodReal3D{T}
    m::Ptr{Nothing}
end

mutable struct WfSchroedingerRotatingReal2D{T<:AbstractFloat} <: WaveFunctionReal2D{T}
    p::Ptr{Nothing}
    m::SchroedingerRotatingReal2D{T}
end

mutable struct WfSchroedingerRotatingReal3D{T<:AbstractFloat} <: WaveFunctionReal3D{T}
    p::Ptr{Nothing}
    m::SchroedingerRotatingReal3D{T}
end

## SchroedingerHermite types ############################################################################

mutable struct SchroedingerHermite1D{T<:AbstractFloat} <: TimeSplittingSpectralMethodComplex1D{T}
    m::Ptr{Nothing}
end

mutable struct SchroedingerHermite2D{T<:AbstractFloat} <: TimeSplittingSpectralMethodComplex2D{T}
    m::Ptr{Nothing}
end

mutable struct SchroedingerHermite3D{T<:AbstractFloat} <: TimeSplittingSpectralMethodComplex3D{T}
    m::Ptr{Nothing}
end

mutable struct SchroedingerHermiteReal1D{T<:AbstractFloat} <: TimeSplittingSpectralMethodReal1D{T}
    m::Ptr{Nothing}
end

mutable struct SchroedingerHermiteReal2D{T<:AbstractFloat} <: TimeSplittingSpectralMethodReal2D{T}
    m::Ptr{Nothing}
end

mutable struct SchroedingerHermiteReal3D{T<:AbstractFloat} <: TimeSplittingSpectralMethodReal3D{T}
    m::Ptr{Nothing}
end

mutable struct WfSchroedingerHermite1D{T<:AbstractFloat} <: WaveFunctionComplex1D{T}
    p::Ptr{Nothing}
    m::SchroedingerHermite1D{T}
end

mutable struct WfSchroedingerHermite2D{T<:AbstractFloat} <: WaveFunctionComplex2D{T}
    p::Ptr{Nothing}
    m::SchroedingerHermite2D{T}
end

mutable struct WfSchroedingerHermite3D{T<:AbstractFloat} <: WaveFunctionComplex3D{T}
    p::Ptr{Nothing}
    m::SchroedingerHermite3D{T}
end

mutable struct WfSchroedingerHermiteReal1D{T<:AbstractFloat} <: WaveFunctionReal1D{T}
    p::Ptr{Nothing}
    m::SchroedingerHermiteReal1D{T}
end

mutable struct WfSchroedingerHermiteReal2D{T<:AbstractFloat} <: WaveFunctionReal2D{T}
    p::Ptr{Nothing}
    m::SchroedingerHermiteReal2D{T}
end

mutable struct WfSchroedingerHermiteReal3D{T<:AbstractFloat} <: WaveFunctionReal3D{T}
    p::Ptr{Nothing}
    m::SchroedingerHermiteReal3D{T}
end

## SchroedingerGeneralizedLaguerre(Hermite) types ########################################################

mutable struct SchroedingerGeneralizedLaguerre2D{T<:AbstractFloat} <: TimeSplittingSpectralMethodComplex2D{T}
    m::Ptr{Nothing}
end

mutable struct SchroedingerGeneralizedLaguerreHermite3D{T<:AbstractFloat} <: TimeSplittingSpectralMethodComplex3D{T}
    m::Ptr{Nothing}
end

mutable struct SchroedingerGeneralizedLaguerreReal2D{T<:AbstractFloat} <: TimeSplittingSpectralMethodReal2D{T}
    m::Ptr{Nothing}
end

mutable struct SchroedingerGeneralizedLaguerreHermiteReal3D{T<:AbstractFloat} <: TimeSplittingSpectralMethodReal3D{T}
    m::Ptr{Nothing}
end

mutable struct WfSchroedingerGeneralizedLaguerre2D{T<:AbstractFloat} <: WaveFunctionComplex2D{T}
    p::Ptr{Nothing}
    m::SchroedingerGeneralizedLaguerre2D{T}
end

mutable struct WfSchroedingerGeneralizedLaguerreHermite3D{T<:AbstractFloat} <: WaveFunctionComplex3D{T}
    p::Ptr{Nothing}
    m::SchroedingerGeneralizedLaguerreHermite3D{T}
end

mutable struct WfSchroedingerGeneralizedLaguerreReal2D{T<:AbstractFloat} <: WaveFunctionReal2D{T}
    p::Ptr{Nothing}
    m::SchroedingerGeneralizedLaguerreReal2D{T}
end

mutable struct WfSchroedingerGeneralizedLaguerreHermiteReal3D{T<:AbstractFloat} <: WaveFunctionReal3D{T}
    p::Ptr{Nothing}
    m::SchroedingerGeneralizedLaguerreHermiteReal3D{T}
end

const periodic = 0
const dirichlet = 1
const neumann = 2

const gauss = 1
const radau = 2
const lobatto = 3

const MEASURE         = UInt32(0)
const PATIENT         = UInt32(1 << 5)
const ESTIMATE        = UInt32(1 << 6)

none_1D(x)=zero(x)
none_2D(x,y)=zero(x)
none_3D(x,y,z)=zero(x)
none_4D(x,y,z,t)=zero(x)

using Libdl

const libtssm = joinpath(dirname(@__FILE__),  "..", "deps", "usr", "lib",
                     string("libtssm.", dlext))
const libtssm_debug = joinpath(dirname(@__FILE__),  "..", "deps", "usr", "lib",
                     string("libtssm_debug.", dlext))
const libtssmq = joinpath(dirname(@__FILE__),  "..", "deps", "usr", "lib",
                     string("libtssmq.", dlext))

# __use_Float128 = false
#try
#   h=dlopen(libtssmq);
#   using Quadmath
#   __use_Float128 = true
#   dlclose(h);
#finally
#end
#const use_Float128 = __use_Float128
const use_Float128 = false

function __init__()
    global tssm_handle
    if "TSSM_DEBUG" in keys(ENV) && ENV["TSSM_DEBUG"]=="1"
        warn("using TSSM debug version")
        tssm_handle = dlopen(libtssm_debug)
    else
        tssm_handle = dlopen(libtssm)
    end
    ccall( dlsym(tssm_handle, "tssm_initialize"), Nothing, ())
    ccall( dlsym(tssm_handle, "tssm_fourier_initialize"), Nothing, ())

    global tssmq_handle
    if use_Float128
        tssmq_handle = dlopen(libtssmq)
        ccall( dlsym(tssmq_handle, "tssmq_initialize"), Nothing, ())
        ccall( dlsym(tssmq_handle, "tssmq_fourier_initialize"), Nothing, ())
    end
    set_fftw_planning_rigor(ESTIMATE)
end

function set_fftw_planning_rigor(flag::Integer=ESTIMATE)
   if !(flag in [ ESTIMATE, PATIENT, MEASURE])
       error("wrong planning rigor flag")
   end
   ccall( dlsym(tssm_handle, "tssm_set_fftw_planning_rigor"), Nothing, (Int32,), flag )
    if use_Float128
       ccall( dlsym(tssmq_handle, "tssmq_set_fftw_planning_rigor"), Nothing, (Int32,), flag )
   end
end

function cfunction_check_return_type(f, r, a)
    rt = Base.return_types(f, a)
    @assert(length(rt)==1 && rt[1]==r, "wrong return type of function")
    cfunction(f, r, a)
end

T = :Float64
TSSM_HANDLE = :tssm_handle
PRE = :tssm_
include("tssm_fourier.jl")
include("tssm_fourier_bessel.jl")
include("tssm_schroedinger.jl")
include("tssm_schroedinger_rotating.jl")
include("tssm_schroedinger_hermite.jl")
include("tssm_schroedinger_generalized_laguerre.jl")
include("tssm_common.jl")
include("tssm_schroedinger_common.jl")

if use_Float128
    T = :Float128
    TSSM_HANDLE = :tssmq_handle
    PRE = :tssmq_
    include("tssm_fourier.jl")
    include("tssm_fourier_bessel.jl")
    include("tssm_schroedinger.jl")
    include("tssm_schroedinger_rotating.jl")
    include("tssm_schroedinger_hermite.jl")
    include("tssm_schroedinger_generalized_laguerre.jl")
    include("tssm_common.jl")
    include("tssm_schroedinger_common.jl")
end

end # TSSM


text:
##

using ACE, Random
using Printf, Test, LinearAlgebra, ACE.Testing, StaticArrays
using ACEbase.Testing: dirfdtest, fdtest, print_tf, test_fio
using ACE: evaluate, evaluate_d, Rn1pBasis, Ylm1pBasis,
      PositionState, Product1pBasis, O3

##

@info("Basic test of PIBasis construction and evaluation")

maxdeg = 6
ord = 3
Bsel = SimpleSparseBasis(ord, maxdeg)

φ = ACE.Invariant()

B1p = ACE.Utils.RnYlm_1pbasis(; maxdeg=maxdeg)

pibasis = PIBasis(B1p, O3(), Bsel; property = φ)
pibasis_r = PIBasis(B1p, O3(), Bsel; property = φ, isreal=true)

# generate a configuration
nX = 10
Xs = rand(PositionState{Float64}, B1p.bases[1], nX)
cfg = ACEConfig(Xs)

AA = evaluate(pibasis, cfg)
AA_r = evaluate(pibasis_r, cfg)

println(@test(length(pibasis) == length(AA)))

spec = ACE.get_spec(pibasis)
println(@test all(length(b) > 0 for b in spec))

spec_naive = [
    [ ACE.get_spec(B1p, pibasis.spec.iAA2iA[iAA, t])
      for t = 1:pibasis.spec.orders[iAA] ]   for iAA = 1:length(pibasis)
    ]

println(@test spec == spec_naive)

# get inverse Aspec
inv_spec1 = Dict{Any, Int}()
for (i, b1) in enumerate(ACE.get_spec(B1p))
  inv_spec1[b1] = i
end

## a really naive implementation of PIBasis to check correctness
A = evaluate(B1p, cfg)
AA_naive =  [ prod( A[ inv_spec1[ b1 ] ] for b1 in b ) for b in spec ]
println(@test( AA_naive ≈ AA ))

AA_r_naive = real.(AA_naive)
println(@test( AA_r_naive ≈ AA_r ))

## FIO tests

@info("FIO Test")
println(@test( all(test_fio(pibasis)) ))
println(@test( all(test_fio(pibasis_r)) ))

## Testing derivatives

@info("Derivatives of PIbasis")
for (pibasis, AA) in [(pibasis, AA), (pibasis_r, AA_r)]
  AA1, dAA = ACE.evaluate_ed(pibasis, cfg)
  println(@test AA1 ≈ AA)

  for ntest = 1:30
    _rrval(x::ACE.XState) = x.rr
    Us = randn(SVector{3, Float64}, length(Xs))
    c = randn(length(pibasis))
    F = t -> sum(c .* ACE.evaluate(pibasis, ACEConfig(Xs + t[1] * Us)))
    dF = t -> [ Us' * _rrval.(sum(c .* ACE.evaluate_ed(pibasis, ACEConfig(Xs + t[1] * Us))[2], dims=1)[:]) ]
    print_tf(@test fdtest(F, dF, [0.0], verbose=false))
  end
  println()
end
##


text:
module TestNAArrayOperators

using FactCheck
using MultidimensionalTables
using MultidimensionalTables.AbstractArrayWrapper

facts("NAArrayOperators tests") do
  context("AbstractArrayWrapper tests") do
    @fact 1 .+ AbstractArrayWrapper(nalift([1,2,3])) --> AbstractArrayWrapper(nalift([2,3,4]))
    @fact AbstractArrayWrapper(nalift([1,2,3])) .- 1 --> AbstractArrayWrapper(nalift([0,1,2]))
    @fact AbstractArrayWrapper(nalift([1,2,3])) .*
            AbstractArrayWrapper(nalift([5.0,6.0,7.0])) --> AbstractArrayWrapper(nalift([5.0,12.0,21.0]))
    @fact map(x->Nullable(x.value*2), AbstractArrayWrapper(nalift([1,2,3]))) --> AbstractArrayWrapper(nalift([2,4,6]))
    @fact (x=AbstractArrayWrapper([5,1,3,2,4]);sort!(x);x) --> AbstractArrayWrapper([1,2,3,4,5])
    @fact sort(AbstractArrayWrapper([5 1 3 2 4]),2) --> AbstractArrayWrapper([1 2 3 4 5])
    @fact sort(AbstractArrayWrapper([5,1,3,2,4])) --> AbstractArrayWrapper([1,2,3,4,5])
    @fact map((x,y)->MultidimensionalTables.naop_plus(x,y), nalift([1,2,3,4,5]), @nalift([1,2,NA,4,5])) --> @nalift([2,4,NA,8,10])
    @fact @nalift([1,2,NA]) .> 1 --> @nalift([false, true, NA])
    @fact 1 .< @nalift([1,2,NA]) --> @nalift([false, true, NA])
    @fact @nalift([1,2,3,NA]) .< @nalift([NA,3,1,2]) --> @nalift([NA,true,false,NA])
    @fact @nalift([1.0,2.0,NA]) .- 1 --> @nalift([0.0,1.0,NA])
    @fact 1 .- @nalift([1.0,2.0,NA]) --> @nalift([0.0,-1.0,NA])
    @fact @nalift([1.0,2.0,NA]) .* @nalift([NA,2.0,3.0]) --> @nalift([NA,4.0,NA])
    @fact -@nalift([1.0,2.0,NA]) --> @nalift([-1.0,-2.0,NA])
    @fact ~@nalift([1,2,NA]) --> @nalift([~1,~2,NA])
    @fact AbstractArrayWrapper([1,2,3]) .> 1 --> AbstractArrayWrapper([false,true,true])
    @fact AbstractArrayWrapper([1,2,3]) .> 1.5 --> AbstractArrayWrapper([false,true,true])
    @fact 1 .< AbstractArrayWrapper([1,2,3]) --> AbstractArrayWrapper([false,true,true])
    @fact AbstractArrayWrapper([1,2,3]) .> Nullable(1) --> nalift([false,true,true])
    @fact AbstractArrayWrapper([1,2,3]) + 1.0 --> AbstractArrayWrapper([2.0,3.0,4.0])
    @fact AbstractArrayWrapper([1,2,3]) + Nullable(1.0) --> nalift([2.0,3.0,4.0])
    @fact Nullable(1) .< AbstractArrayWrapper([1,2,3]) --> nalift([false,true,true])
    @fact AbstractArrayWrapper([3.0,2.0,1.0]) .< AbstractArrayWrapper([1.0,2.0,3.0]) --> AbstractArrayWrapper([false,false,true])
    @fact AbstractArrayWrapper([3,2,1]) .< nalift([1,2,3]) --> nalift([false,false,true])
    @fact @nalift([3,NA,1]) .< AbstractArrayWrapper([1.0,2.0,3.0]) --> @nalift([false,NA,true])
    @fact nalift([3,2,1]) .< nalift([1.0,2.0,3.0]) --> nalift([false,false,true])
    @fact nalift([3,2,1]) + nalift([1.0,2.0,3.0]) --> nalift([4.0,4.0,4.0])
    @fact nalift([3,2,1]) + AbstractArrayWrapper([1.0,2.0,3.0]) --> nalift([4.0,4.0,4.0])
    @fact AbstractArrayWrapper([3,2,1]) + AbstractArrayWrapper([1.0,2.0,3.0]) --> AbstractArrayWrapper([4.0,4.0,4.0])
    @fact AbstractArrayWrapper([3,2,1]) + 5.0 --> AbstractArrayWrapper([8.0,7.0,6.0])
    @fact AbstractArrayWrapper([3,2,1]) + Nullable(5.0) --> nalift([8.0,7.0,6.0])
    @fact -AbstractArrayWrapper([3,2,1]) --> AbstractArrayWrapper([-3,-2,-1])
    @fact -AbstractArrayWrapper([3.0,2.0,1.0]) --> AbstractArrayWrapper([-3.0,-2.0,-1.0])
    @fact -nalift([3.0,2.0,1.0]) --> nalift([-3.0,-2.0,-1.0])
    @fact -darr(a=[3.0,2.0,1.0]) --> darr(a=[-3.0,-2.0,-1.0])
    @fact -larr(a=[3.0,2.0,1.0], axis1=[:x,:y,:z]) --> larr(a=[-3.0,-2.0,-1.0], axis1=[:x,:y,:z])
    @fact nalift([1.0,2.0,3.0]) + nalift([1,2,3]) --> nalift([2.0,4.0,6.0])
    @fact nalift([1.0,2.0,3.0]) + nalift([1.0,2.0,3.0]) --> nalift([2.0,4.0,6.0])
    @fact AbstractArrayWrapper(map(Nullable, [1.0,2.0,3.0])) + nalift([1.0,2.0,3.0]) --> nalift([2.0,4.0,6.0])
    @fact AbstractArrayWrapper(map(Nullable, [1.0,2.0,3.0])) + nalift([1,2,3]) --> nalift([2.0,4.0,6.0])
    @fact AbstractArrayWrapper(map(Nullable, [1.0,2.0,3.0])) + AbstractArrayWrapper([1,2,3]) --> nalift([2.0,4.0,6.0])
    @fact nalift([1.0,2.0,3.0]) + 1.0 --> nalift([2.0,3.0,4.0])
    @fact 1.0 + nalift([1.0,2.0,3.0]) --> nalift([2.0,3.0,4.0])
    @fact nalift([1.0,2.0,3.0]) + 1 --> nalift([2.0,3.0,4.0])
    @fact 1 + nalift([1.0,2.0,3.0]) --> nalift([2.0,3.0,4.0])
    @fact 1.0 + nalift([1,2,3]) --> nalift([2.0,3.0,4.0])
    @fact nalift([1,2,3]) + 1.0 --> nalift([2.0,3.0,4.0])
    @fact nalift([1.0,2.0,3.0]) + Nullable(1.0) --> nalift([2.0,3.0,4.0])
    @fact Nullable(1.0) + nalift([1.0,2.0,3.0]) --> nalift([2.0,3.0,4.0])
    @fact nalift([1.0,2.0,3.0]) + Nullable(1) --> nalift([2.0,3.0,4.0])
    @fact Nullable(1) + nalift([1.0,2.0,3.0]) --> nalift([2.0,3.0,4.0])
    @fact Nullable(1.0) + nalift([1,2,3]) --> nalift([2.0,3.0,4.0])
    @fact nalift([1,2,3]) + Nullable(1.0) --> nalift([2.0,3.0,4.0])
    @fact nalift([1.0,2.0,3.0]) * 1.0 --> nalift([2.0,3.0,4.0]) - 1
    @fact 1.0 * nalift([1.0,2.0,3.0]) --> nalift([2.0,3.0,4.0]) - 1
    @fact nalift([1.0,2.0,3.0]) * 1 --> nalift([2.0,3.0,4.0]) - 1
    @fact 1 * nalift([1.0,2.0,3.0]) --> nalift([2.0,3.0,4.0]) - 1
    @fact 1.0 * nalift([1,2,3]) --> nalift([2.0,3.0,4.0]) - 1
    @fact nalift([1,2,3]) * 1.0 --> nalift([2.0,3.0,4.0]) - 1
    @fact nalift([1.0,2.0,3.0]) * Nullable(1.0) --> nalift([2.0,3.0,4.0]) - 1
    @fact Nullable(1.0) * nalift([1.0,2.0,3.0]) --> nalift([2.0,3.0,4.0]) - 1
    @fact nalift([1.0,2.0,3.0]) * Nullable(1) --> nalift([2.0,3.0,4.0]) - 1
    @fact Nullable(1) * nalift([1.0,2.0,3.0]) --> nalift([2.0,3.0,4.0]) - 1
    @fact Nullable(1.0) * nalift([1,2,3]) --> nalift([2.0,3.0,4.0]) - 1
    @fact nalift([1,2,3]) * Nullable(1.0) --> nalift([2.0,3.0,4.0]) - 1
  end
end

end


text:
# look at https://github.com/nacs-lab/yyc-data/blob/d082032d075070b133fe909c724ecc405e80526a/lib/NaCsCalc/src/utils.jl#L120-L142
# https://discourse.julialang.org/t/poor-performance-on-cluster-multithreading/12248/39
# https://discourse.julialang.org/t/two-questions-about-multithreading/14564/2
# https://discourse.julialang.org/t/question-about-multi-threading-performance/12075/3

# check out this??
# https://discourse.julialang.org/t/anyone-developing-multinomial-logistic-regression/23222

# also
# https://github.com/nacs-lab/yyc-data/blob/d082032d075070b133fe909c724ecc405e80526a/lib/NaCsCalc/src/utils.jl#L120-L142

struct DrillingTmpVars{T<:AbstractFloat} <: AbstractTmpVars
    ubv::Vector{T}
    dubv::Matrix{T}
    llj::Vector{T}
    grad::Vector{T}
    gradJ::Matrix{T}
    function DrillingTmpVars(ubv, dubv, llj, grad, gradJ)
        length(grad)==size(gradJ,1) == size(dubv,1) || throw(DimensionMismatch())
        length(ubv) == size(dubv,2) || throw(DimensionMismatch())
        T = eltype(ubv)
        return new{T}(ubv, dubv, llj, grad, gradJ)
    end
end

_ubv(  dtv::DrillingTmpVars) = dtv.ubv
_dubv( dtv::DrillingTmpVars) = dtv.dubv
_llj(  dtv::DrillingTmpVars) = dtv.llj
_grad( dtv::DrillingTmpVars) = dtv.grad
_gradJ(dtv::DrillingTmpVars) = dtv.gradJ

function DrillingTmpVars(J::Integer, maxchoices::Integer, k::Integer, T::Type=Float64)
    ubv   = zeros(T, maxchoices)
    dubv  = zeros(T, k, maxchoices)
    llj   = zeros(T, J)
    grad  = zeros(T, k)
    gradJ = zeros(T, k, J)
    return DrillingTmpVars(ubv, dubv, llj, grad, gradJ)
end

function DrillingTmpVars(J::Integer, model::AbstractDrillModel, T::Type=Float64)
    maxchoices = num_choices(model)
    k = _nparm(model)
    return DrillingTmpVars(J, maxchoices, k, T)
end


text:
module FunsTest

using Funs
using Base.Test

function inc(i::Integer)
    i+1
end

immutable Inc
    i::Int
    Inc(i) = new(i+1)
end

const incf = Fun{Int}(inc)

function main()
    # untyped functions
    let
        # Function
        v1 = fcall(inc, 1)
        @test v1 == 2

        # DataType
        v2 = fcall(Inc, 2)
        @test v2.i == 3

        # Fun
        v3 = fcall(incf, 3)
        @test v3 == 4
    end

    # typed function
    let
        # Function
        v1 = fcall(R=Int, inc, 1)
        @test v1 == 2

        # DataType
        v2 = fcall(R=Inc, Inc, 2)
        @test v2.i == 3

        # Fun
        v3 = fcall(R=Int, incf, 3)
        @test v3 == 4
    end

    # untyped macros
    let
        # Function
        v1 = @fcall inc(1)
        @test v1 == 2

        # DataType
        v2 = @fcall Inc(2)
        @test v2.i == 3

        # Fun
        v3 = @fcall incf(3)
        @test v3 == 4
    end

    # typed macros
    let
        # Function
        v1 = @fcall Int inc(1)
        @test v1 == 2

        # DataType
        v2 = @fcall Inc Inc(2)
        @test v2.i == 3

        # Fun
        v3 = @fcall Int incf(3)
        @test v3 == 4
    end
end

end


text:
@def Ax_max begin #(U)

	# maximum longitudinal acceleration for given speed
	Ax_max = Array(Float64,(length(U),1))
	for i in eachindex(U)
		Ax_max[i,1] = AXC[1]*U[i]^3 + AXC[2]*U[i]^2 + AXC[3]*U[i] + AXC[4]
	end
	Ax_max
end


text:
include("components/CO2emissions.jl")
include("components/CO2cycle.jl")
include("components/CO2forcing.jl")
include("components/CH4emissions.jl")
include("components/CH4cycle.jl")
include("components/CH4forcing.jl")
include("components/N2Oemissions.jl")
include("components/N2Ocycle.jl")
include("components/N2Oforcing.jl")
include("components/LGemissions.jl")
include("components/LGcycle.jl")
include("components/LGforcing.jl")
include("components/SulphateForcing.jl")
include("components/TotalForcing.jl")
include("components/ClimateTemperature.jl")

m = Model()
set_dimension!(m, :time, [2009, 2010, 2020, 2030, 2040, 2050, 2075, 2100, 2150, 2200])
set_dimension!(m, :region, ["EU", "USA", "OECD","USSR","China","SEAsia","Africa","LatAmerica"])

#add all the components
add_comp!(m,co2emissions)
add_comp!(m, co2cycle)
add_comp!(m, co2forcing)
add_comp!(m, ch4emissions)
add_comp!(m, ch4cycle)
add_comp!(m, ch4forcing)
add_comp!(m, n2oemissions)
add_comp!(m, n2ocycle)
add_comp!(m, n2oforcing)
add_comp!(m, LGemissions)
add_comp!(m, LGcycle)
add_comp!(m, LGforcing)
add_comp!(m, SulphateForcing)
add_comp!(m, TotalForcing)
add_comp!(m, ClimateTemperature)

#connect parameters together

set_param!(m, :co2cycle, :y_year, [2009.,2010.,2020.,2030.,2040.,2050.,2075.,2100.,2150.,2200.])
set_param!(m, :co2cycle, :y_year_0, 2008.)
connect_param!(m, :co2cycle => :e_globalCO2emissions, :co2emissions => :e_globalCO2emissions)
connect_param!(m, :co2cycle => :rt_g0_baseglobaltemp, :ClimateTemperature => :rt_g0_baseglobaltemp)
connect_param!(m, :co2cycle => :rt_g_globaltemperature, :ClimateTemperature => :rt_g_globaltemperature)

connect_param!(m, :co2forcing => :c_CO2concentration, :co2cycle => :c_CO2concentration)

set_param!(m, :ch4cycle, :y_year, [2009.,2010.,2020.,2030.,2040.,2050.,2075.,2100.,2150.,2200.])
set_param!(m, :ch4cycle, :y_year_0, 2008.)
connect_param!(m, :ch4cycle => :e_globalCH4emissions, :ch4emissions => :e_globalCH4emissions)
connect_param!(m, :ch4cycle => :rtl_g0_baselandtemp, :ClimateTemperature => :rtl_g0_baselandtemp)
connect_param!(m, :ch4cycle => :rtl_g_landtemperature, :ClimateTemperature => :rtl_g_landtemperature)

connect_param!(m, :ch4forcing => :c_CH4concentration, :ch4cycle => :c_CH4concentration)
connect_param!(m, :ch4forcing => :c_N2Oconcentration, :n2ocycle => :c_N2Oconcentration)

set_param!(m, :n2ocycle, :y_year, [2009.,2010.,2020.,2030.,2040.,2050.,2075.,2100.,2150.,2200.])
set_param!(m, :n2ocycle, :y_year_0, 2008.)
connect_param!(m, :n2ocycle => :e_globalN2Oemissions, :n2oemissions => :e_globalN2Oemissions)
connect_param!(m, :n2ocycle => :rtl_g0_baselandtemp, :ClimateTemperature => :rtl_g0_baselandtemp)
connect_param!(m, :n2ocycle => :rtl_g_landtemperature, :ClimateTemperature => :rtl_g_landtemperature)

connect_param!(m, :n2oforcing => :c_CH4concentration, :ch4cycle => :c_CH4concentration)
connect_param!(m, :n2oforcing => :c_N2Oconcentration, :n2ocycle => :c_N2Oconcentration)

set_param!(m, :LGcycle, :y_year, [2009.,2010.,2020.,2030.,2040.,2050.,2075.,2100.,2150.,2200.])
set_param!(m, :LGcycle, :y_year_0, 2008.)
connect_param!(m, :LGcycle => :e_globalLGemissions, :LGemissions => :e_globalLGemissions)
connect_param!(m, :LGcycle => :rtl_g0_baselandtemp, :ClimateTemperature => :rtl_g0_baselandtemp)
connect_param!(m, :LGcycle => :rtl_g_landtemperature, :ClimateTemperature => :rtl_g_landtemperature)

connect_param!(m, :LGforcing => :c_LGconcentration, :LGcycle => :c_LGconcentration)

connect_param!(m, :TotalForcing => :f_CO2forcing, :co2forcing => :f_CO2forcing)
connect_param!(m, :TotalForcing => :f_CH4forcing, :ch4forcing => :f_CH4forcing)
connect_param!(m, :TotalForcing => :f_N2Oforcing, :n2oforcing => :f_N2Oforcing)
connect_param!(m, :TotalForcing => :f_lineargasforcing, :LGforcing => :f_LGforcing)

set_param!(m, :ClimateTemperature, :y_year, [2009.,2010.,2020.,2030.,2040.,2050.,2075.,2100.,2150.,2200.])
set_param!(m, :ClimateTemperature, :y_year_0, 2008.)
connect_param!(m, :ClimateTemperature => :ft_totalforcing, :TotalForcing => :ft_totalforcing)
connect_param!(m, :ClimateTemperature => :fs_sulfateforcing, :SulphateForcing => :fs_sulphateforcing)

# next: add vector and panel example
p = load_parameters(m)
p["y_year_0"] = 2008.
p["y_year"] = Mimi.dim_keys(m.md, :time)
set_leftover_params!(m, p)

run(m)


text:
using BlockArrays, Test, Base64

# useds to force SubArray return
bview(a, b) = Base.invoke(view, Tuple{AbstractArray,Any}, a, b)

@testset "Block Views" begin
    @testset "block slice" begin
        A = PseudoBlockArray(1:6,1:3)
        b = parentindices(bview(A, Block(2)))[1] # A BlockSlice

        @test first(b) == 2
        @test last(b) == 3
        @test length(b) == 2
        @test step(b) == 1
        @test axes(b) == (Base.OneTo(2),)
        @test Base.axes1(b) == Base.OneTo(2)
        @test size(b) == (2,)
        @test collect(b) == [2,3]
        @test b[1] == 2
        @test b[1:2] == 2:3
    end

    @testset "block view" begin
        A = BlockArray(collect(1:6), 1:3)
        @test view(A, Block(2)) == [2,3]
        view(A, Block(2))[2] = -1
        @test A[3] == -1

        # backend tests
        @test_throws ArgumentError Base.to_index(A, Block(1))

        A = PseudoBlockArray(collect(1:6), 1:3)
        @test view(A, Block(2)) == [2,3]
        view(A, Block(2))[2] = -1
        @test A[3] == -1

        # backend tests
        @test_throws ArgumentError Base.to_index(A, Block(1))

        A = PseudoBlockArray(reshape(collect(1:(6*12)),6,12), 1:3, 3:5)
        V = view(A, Block(2), Block(3))
        @test size(V) == (2, 5)
        V[1,1] = -1
        @test A[2,8] == -1
        @test A[Block(4)] == A[Block(1),Block(2)]
        @test_throws BlockBoundsError A[Block(10)]

        V = view(A, Block(3, 2))
        @test size(V) == (3, 4)
        V[2,1] = -2
        @test A[5,4] == -2

        # test mixed blocks and other indices
        @test view(A, Block(2), 2) == [8,9]
        @test similar(A, (Base.OneTo(5), axes(A,2))) isa PseudoBlockArray{Int}
        @test view(A, Block(2), :) == A[2:3,:]

        @test view(A, 2, Block(1)) == [2,8,14]
        @test view(A, :, Block(1)) == A[:,1:3]

        @test view(V, Block(1, 1)) ≡ V

        @test_throws BlockBoundsError view(V, Block(1,2))
        @test_throws BlockBoundsError view(V, Block(2,1))

        A = PseudoBlockArray(reshape(collect(1:(6^3)),6,6,6), 1:3, 1:3, 1:3)
        V = view(A, Block(2), Block(3), Block(1))
        @test size(V) == (2, 3, 1)
        V[1,1,1] = -3
        @test A[2,4,1] == -3

        V = view(A,Block(1,1,1))
        @test size(V) == (1,1,1)
        V[1,1,1] = -4
        @test A[1,1,1] == -4

        # blocks mimic CartesianIndex in views
        V = view(A,Block(1,1),Block(2))
        @test size(V) == (1,1,2)
        V[1,1,1] = -5
        @test A[1,1,2] == -5

        V = view(A,Block(2),Block(1,1))
        @test size(V) == (2,1,1)
        V[1,1,1] = -6
        @test A[2,1,1] == -6

        # test mixed blocks and other indices
        @test view(A, Block(1), Block(2), 1) == A[1:1,2:3,1]
        @test view(A, Block(1,2), 1) == A[1:1,2:3,1]
        @test view(A, Block(1), 2, 1) == A[1:1,2,1]
        @test view(A, 1, Block(2), 1) == A[1,2:3,1]
        @test view(A, 1, 2, Block(2)) == A[1,2,2:3]
    end

    @testset "block view pointers" begin
        A = BlockArray(reshape(Vector{Float64}(1:(6^2)),6,6), 1:3, 1:3)

        V = view(A, Block(1,2))
        @test Base.unsafe_convert(Ptr{Float64}, V) == Base.unsafe_convert(Ptr{Float64}, A.blocks[1,2])
        @test unsafe_load(pointer(V)) == V[1,1]

        A = PseudoBlockArray(reshape(Vector{Float64}(1:(6^2)),6,6), 1:3, 1:3)

        V = view(A, Block(1,2))
        @test Base.unsafe_convert(Ptr{Float64}, V) == Base.unsafe_convert(Ptr{Float64}, view(A.blocks, 1:1, 2:3))
        @test unsafe_load(pointer(V)) == V[1,1]

        V = view(A, Block(2), 2:3)
        @test Base.unsafe_convert(Ptr{Float64}, V) == Base.unsafe_convert(Ptr{Float64}, view(A.blocks, 2:3, 2:3))
        @test unsafe_load(pointer(V)) == V[1,1]

        V = view(A, 2:3, Block(2))
        @test Base.unsafe_convert(Ptr{Float64}, V) == Base.unsafe_convert(Ptr{Float64}, view(A.blocks, 2:3, 2:3))
        @test unsafe_load(pointer(V)) == V[1,1]
    end

    @testset "block indx range of block range" begin
        A = PseudoBlockArray(collect(1:6), 1:3)
        V = view(A, Block.(1:2))
        @test V == 1:3
        @test axes(V,1) isa BlockArrays.BlockedUnitRange
        @test blockaxes(V,1) == Block.(1:2)
        @test view(V, Block(2)[1:2]) == [2,3]
        V = view(A, Block.(2:3))
        @test V == 2:6
        @test view(V, Block(2)[1:2]) == [4,5]
    end

    @testset "subarray implements block interface" begin
        A = PseudoBlockArray(reshape(Vector{Float64}(1:(6^2)),6,6), 1:3, 1:3)

        V = view(A, Block(2,3))
        @test PseudoBlockArray(V) isa PseudoBlockArray
        @test BlockArray(V) isa BlockArray
        @test PseudoBlockArray(V) == BlockArray(V) == V

        V = view(A, Block(2), Block.(2:3))
        @test PseudoBlockArray(V) isa PseudoBlockArray
        @test BlockArray(V) isa BlockArray
        @test PseudoBlockArray(V) == BlockArray(V) == V
        @test blocksize(V) == (1,2)

        V = view(A, Block.(2:3), Block(3))
        @test PseudoBlockArray(V) isa PseudoBlockArray
        @test BlockArray(V) isa BlockArray
        @test PseudoBlockArray(V) == BlockArray(V) == V
        @test blocksize(V) == (2,1)

        V = view(A, Block.(2:3), Block.(1:2))
        @test PseudoBlockArray(V) isa PseudoBlockArray
        @test BlockArray(V) isa BlockArray
        @test PseudoBlockArray(V) == BlockArray(V) == V
        @test blocksize(V) == (2,2)
    end

    @testset "Block-BlockRange blocks" begin
        A = BlockArray([1 2 3; 4 5 6; 7 8 9], 1:2, 1:2)
        V = view(A,Block(1),Block.(1:2))
        W = view(A,Block.(1:2),Block(1))
        @test blocks(V) == blocks(A)[1:1,1:2]
        @test blocks(W) == blocks(A)[1:2,1:1]
        @test stringmime("text/plain", V) == "1×3 view(::BlockMatrix{$Int, Matrix{Matrix{$Int}}, $(typeof(axes(A)))}, BlockSlice(Block(1),1:1), BlockSlice(Block{1, $Int}[Block(1), Block(2)],1:1:3)) with eltype $Int with indices Base.OneTo(1)×1:1:3:\n 1  │  2  3"
        @test stringmime("text/plain", W) == "3×1 view(::BlockMatrix{$Int, Matrix{Matrix{$Int}}, $(typeof(axes(A)))}, BlockSlice(Block{1, $Int}[Block(1), Block(2)],1:1:3), BlockSlice(Block(1),1:1)) with eltype $Int with indices 1:1:3×Base.OneTo(1):\n 1\n ─\n 4\n 7"
    end

    @testset "getindex with BlockRange" begin
        A = BlockArray(randn(6), 1:3)
        @test A[Block.(2:3)] isa BlockArray
        @test A[Block.(2:3)] == A[2:end]
        A = PseudoBlockArray(randn(6), 1:3)
        @test A[Block.(2:3)] isa PseudoBlockArray
        @test A[Block.(2:3)] == A[2:end]
    end

    @testset "non-allocation blocksize" begin
        A = BlockArray(randn(5050), 1:100)
        @test blocksize(A) == (100,)
        @test @allocated(blocksize(A)) ≤ 40
        V = view(A, Block(3))
        @test blocksize(V) == (1,)
        @test @allocated(blocksize(V)) ≤ 40
        V = view(A, Block.(1:30))
        @test blocksize(V) == (30,)
        @test @allocated(blocksize(V)) ≤ 40
        V = view(A, 3:43)
        @test blocksize(V) == (1,)
        V = view(A, 5)
        @test blocksize(V) == ()

        A = BlockArray(randn(5050,21), 1:100, 1:6)
        @test blocksize(A) == (100,6)
        @test @allocated(blocksize(A)) ≤ 40
        V = view(A, Block(3,2))
        @test blocksize(V) == (1,1)
        @test @allocated(blocksize(V)) ≤ 40
        V = view(A, Block.(1:30), Block(3))
        @test blocksize(V) == (30,1)
        @test @allocated(blocksize(V)) ≤ 40
        V = view(A, Block.(1:30), Block.(1:3))
        @test blocksize(V) == (30,3)
        @test @allocated(blocksize(V)) ≤ 40
        V = view(A, 3:43,1:3)
        @test blocksize(V) == (1,1)
        V = view(A, 5, 1:3)
        @test blocksize(V) == (1,)
    end

    @testset "hasmatchingblocks" begin
        A = BlockArray{Int}(undef, 1:20, 1:20)
        B = BlockArray{Int}(undef, 1:3, fill(3,2))
        V = view(A,Block.(1:10),Block.(1:10))

        @test BlockArrays.hasmatchingblocks(A)
        @test BlockArrays.hasmatchingblocks(V)
        @test @allocated(BlockArrays.hasmatchingblocks(V)) == 0
        @test !BlockArrays.hasmatchingblocks(view(A,Block.(1:2),1:3))
        @test !BlockArrays.hasmatchingblocks(view(A,Block.(1:2),Block.(2:3)))

        @test BlockArrays.hasmatchingblocks(view(B,Block.(3:3),Block.(2:2)))
    end

    @testset "sub_materialize cases" begin
        a = BlockArray(randn(6), 1:3)
        b = PseudoBlockArray(randn(6), 1:3)
        @test a[Block.(1:2)] isa BlockArray
        @test b[Block.(1:2)] isa PseudoBlockArray
        @test a[1:3] isa Array
        @test b[1:3] isa Array
        A = BlockArray(randn(6,6), 1:3, fill(3,2))
        B = PseudoBlockArray(randn(6,6), 1:3, fill(3,2))
        @test A[Block.(1:2),Block.(1:2)] isa BlockArray
        @test B[Block.(1:2),Block.(1:2)] isa PseudoBlockArray
        @test A[Block.(1:2),1:3] isa PseudoBlockArray
        @test B[Block.(1:2),1:3] isa PseudoBlockArray
        @test A[1:3,Block.(1:2)] isa PseudoBlockArray
        @test B[1:3,Block.(1:2)] isa PseudoBlockArray
        @test A[Block.(1:2),:] isa BlockArray
        @test B[Block.(1:2),:] isa PseudoBlockArray
        @test blockisequal(axes(A,2),axes(A[Block.(1:2),:],2))
        @test blockisequal(axes(B,2),axes(B[Block.(1:2),:],2))
        @test A[:,Block.(1:2)] isa BlockArray
        @test B[:,Block.(1:2)] isa PseudoBlockArray
        @test blockisequal(axes(A,1),axes(A[:,Block.(1:2)],1))
        @test blockisequal(axes(B,1),axes(B[:,Block.(1:2)],1))
        @test A[:,:] isa BlockArray
        @test B[:,:] isa PseudoBlockArray
        @test blockisequal(axes(A),axes(A[:,:]))
        @test blockisequal(axes(B),axes(B[:,:]))
        @test A[1:3,1:3] isa Array
        @test B[1:3,1:3] isa Array
        A = BlockArray(randn(6,6,6), 1:3, fill(3,2),1:3)
        B = PseudoBlockArray(randn(6,6,6), 1:3, fill(3,2),1:3)
        @test A[Block.(1:2),Block.(1:2),Block.(1:2)] isa BlockArray
        @test B[Block.(1:2),Block.(1:2),Block.(1:2)] isa PseudoBlockArray
        @test A[1:3,Block.(1:2),1:3] isa BlockArray
        @test B[1:3,Block.(1:2),1:3] isa PseudoBlockArray
    end

    @testset "BlockArray BlockRange view" begin
        a = BlockArray(randn(6), 1:3)
        v = view(a, Block(3)[1:2])
        v[1] = 5
        @test a[4] == 5
    end

    @testset "types with custom views" begin
        a = mortar([7:9,5:6])
        v = view(a,Block.(1:2))
        @test a[Block(1)[1:3]] ≡ view(a,Block(1)[1:3]) ≡ view(v,Block(1)[1:3]) ≡ 7:9
    end

    @testset "blockrange-of-blockreange" begin
        a = mortar([7:9,5:6])
        v = view(a,Block.(1:2))
        @test view(v, Block(1)) ≡ 7:9
        @test bview(v, Block(1)) == 7:9
        @test view(v,Block.(1:2)) == v

        A = BlockArray([1 2 3; 4 5 6; 7 8 9], 1:2, 1:2)
        V = view(A, Block.(1:2), Block.(1:2))
        @test view(V, Block(1,1)) ≡ view(A, Block(1,1))
        @test bview(V, Block(1,1)) == view(A, Block(1,1))
    end

    @testset "view of BlockSlice unwinds" begin
        a = mortar([7:9,5:6])
        v = bview(a,Block(1))
        @test view(a, parentindices(v)...) ≡ 7:9
    end
end


text:
issquare(A) = size(A, 1) == size(A, 2) && ndims(A) == 2

perf_hvcat(A, B) =  [A B; B A]

function perf_hvcat_setind(A, B)
    @assert issquare(A) && size(A) == size(B)
    n = size(A, 1)
    C = Matrix{Float64}(undef, 2n, 2n)
    C[1:n, 1:n] = A
    C[1:n, (n+1):end] = B
    C[(n+1):end, 1:n] = B
    C[(n+1):end, (n+1):end] = A
end

perf_hcat(A, B) = [A B B A]

function perf_hcat_setind(A, B)
    @assert issquare(A) && size(A) == size(B)
    n = size(A, 1)
    C = Matrix{Float64}(undef, n, 4n)
    C[:, 1:n] = A
    C[:, (n+1):2n] = B
    C[:, (2n+1):3n] = B
    C[:, (3n+1):end] = A
    return C
end

perf_vcat(A, B) = [A; B; B; A]

function perf_vcat_setind(A, B)
    @assert issquare(A) && size(A) == size(B)
    n = size(A, 1)
    C = Matrix{Float64}(undef, 4n, n)
    C[1:n, :] = A
    C[(n+1):2n, :] = B
    C[(2n+1):3n, :] = B
    C[(3n+1):4n, :] = A
    return C
end

function perf_catnd(n)
    A = samerand(1, n, n, 1)
    B = samerand(1, n, n)
    return cat(A, B, B, A; dims = 3)
end

function perf_catnd_setind(n)
    A = samerand(1, n, n, 1)
    B = samerand(1, n, n)
    C = Array{Float64}(undef, 1, n, 4n, 1)
    C[1, :, 1:n, 1] = A
    C[1, :, (n+1):2n, 1] = B
    C[1, :, (2n+1):3n, 1] = B
    C[1, :, (3n+1):4n, 1] = A
    return C
end

function perf_cat_4467()
    a, b, c, d, e = 1.0, 2.0, 3.0, 4.0, 5
    return [a, b, c, d, e]
end


text:
module LanguageDynamics

# Imports
import Distributions

# Auxiliaries
include("common.jl")

# Abstract types
include("AbstractPopulations.jl")
include("AbstractInteractors.jl")
include("AbstractReplicators.jl")

# Replicators
include("Replicators/BinaryVariable.jl")

# Interactors
include("Interactors/Interactors.jl")
include("Interactors/EmptyInteractor.jl")
include("Interactors/MomentumSelector.jl")

# Populations
include("Populations/Populations.jl")
include("Populations/PoolPopulation.jl")
include("Populations/ZipfTravellerPopulation.jl")

# Exports: auxiliaries
export Point
export zipf
export manhattan
export chebyshev

# Exports: abstracts
export AbstractInteractor
export AbstractPopulation

# Exports: populations
export PoolPopulation
export ZipfTravellerPopulation
export inject!
export eject!
export rendezvous!
export travel!
export characterize
export characterize_by_location

# Exports: interactors
export EmptyInteractor
export MomentumSelector

# Exports: replicators
export BinaryVariable

end