text:
export aff_ind_maximum_likelihood_vertices
export aff_ind_anti_guessing_vertices
export aff_ind_non_negativity_vertices
export aff_ind_ambiguous_guessing_vertices
export aff_ind_k_guessing_vertices
export aff_ind_coarse_grained_input_ambiguous_guessing_vertices

"""
    aff_ind_maximum_likelihood_vertices( N :: Int64, d :: Int64 ) :: Vector{Matrix{Int64}}

Enumerates an affinely independent set of vertices that maximize the [`maximum_likelihood_facet`](@ref).

A `DomainError` is thrown if `N > d > 1` is not satisfied.
"""
function aff_ind_maximum_likelihood_vertices(N :: Int64, d :: Int64) :: Vector{Matrix{Int64}}
    if !(N > d > 1)
        throw(DomainError(d, "input `d` does not satisfy `N > d > 1`."))
    end

    aff_ind_vecs = _aff_ind_vecs(d-1, N-d)

    matrices = Vector{Matrix{Int64}}(undef, N*(N-1))
    id = 1
    for row_id in 1:N
        for vec in aff_ind_vecs
            m = zeros(Int64,N,N)
            m[row_id,row_id] = 1

            for col in 1:N-1
                col_id = col < row_id ? col : col + 1

                if vec[col] == 0
                    m[col_id,col_id] = 1
                else
                    m[row_id,col_id] = 1
                end
            end

            matrices[id] = m

            id += 1
        end
    end

    return matrices
end

"""
    aff_ind_non_negativity_vertices(
        X :: Int64,
        Y :: Int64
    ) :: Vector{Matrix{Int64}}

Enumerates an affinely independent set of deterministic, vertices that maximize
the [`non_negativity_facet`](@ref)

A `DomainError` is thrown if the inputs do not satisfy `X > 1` and  `Y > 1`.
"""
function aff_ind_non_negativity_vertices(X :: Int64, Y :: Int64) :: Vector{Matrix{Int64}}
    if !(Y > 1)
        throw(DomainError(Y, "Inputs must satisfy `Y > 1`"))
    elseif !(X > 1)
        throw(DomainError(X, "Inputs must satisfy `X > 1`"))
    end

    matrices = Vector{Matrix{Int64}}(undef, (Y-1)*(X))

    matrix_id = 1
    for row_id in 1:(Y-1)
        m1 = zeros(Int64, Y, X)
        m1[row_id,:] .= 1

        matrices[matrix_id] = m1
        matrix_id += 1

        for col_id in 2:X
            m2 = zeros(Int64, Y, X)
            m2[row_id,:] .= 1
            m2[row_id,col_id] = 0
            m2[Y,col_id] = 1

            matrices[matrix_id] = m2
            matrix_id += 1
        end
    end

    matrices
end

"""
    aff_ind_ambiguous_guessing_vertices(
        Y :: Int64,
        d :: Int64
    ) :: Vector{Matrix{Int64}}

Enumerates an affinely independent set of vertices that maximize the [`ambiguous_guessing_facet`](@ref).

A `DomainError` is thrown if the inputs don't satisfy the following requirements:
* `Y ≥ 4`
* `(Y - 2) ≥ d ≥ 2`
"""
function aff_ind_ambiguous_guessing_vertices(Y :: Int64, d :: Int64) :: Vector{Matrix{Int64}}
    if !(Y ≥ 4)
        throw(DomainError(Y, "Y must satisfy `Y ≥ 4`"))
    elseif !(Y-2 ≥ d ≥ 2)
        throw(DomainError(d, "d must satisfy `Y-2 ≥ d ≥ 2`"))
    end

    matrices = Vector{Matrix{Int64}}(undef, (Y-1)*(Y-1))
    matrix_id = 1

    maximum_likelihood_facet_strats = aff_ind_maximum_likelihood_vertices(Y-1, d)

    for strat in maximum_likelihood_facet_strats
        m = cat(strat, zeros(Int64,(1,Y-1)), dims=1)

        matrices[matrix_id] = m
        matrix_id += 1
    end

    vecs = _aff_ind_vecs(d-1, Y-d)
    for vec in vecs
        m = zeros(Int64, (Y,Y-1))
        m[Y,:] = vec

        for id in 1:(Y-1)
            if vec[id] == 0
                m[id,id] = 1
            end
        end

        matrices[matrix_id] = m
        matrix_id += 1
    end

    matrices
end

"""
    aff_ind_coarse_grained_input_ambiguous_guessing_vertices(
        Y :: Int64,
        d :: Int64
    ) :: Vector{Matrix{Int64}}

Enumerates an affinely independent set of vertices for the
[`coarse_grained_input_ambiguous_guessing_facet`](@ref).

A `DomainError` is thrown if the inputs don't satisfy the following requirements:
* `Y ≥ 4`
* `(Y - 2) ≥ d ≥ 2`
"""
function aff_ind_coarse_grained_input_ambiguous_guessing_vertices(Y :: Int64, d :: Int64) :: Vector{Matrix{Int64}}
    if !(Y ≥ 4)
        throw(DomainError(Y, "Y must satisfy `Y ≥ 4`"))
    elseif !(Y-2 ≥ d ≥ 2)
        throw(DomainError(d, "d must satisfy `Y-2 ≥ d ≥ 2`"))
    end

    matrices = Vector{Matrix{Int64}}(undef, (Y*(Y-1)))
    matrix_id = 1

    for strat in aff_ind_ambiguous_guessing_vertices(Y,d)
        target_rows = findall(i -> i ≠ 0, sum.(eachrow(strat)))

        target_row = (Y-1) in target_rows ? (Y-1) : min(target_rows...)

        m = zeros(Int64, Y,Y)
        m[1:Y,1:Y-1] = strat
        m[target_row,Y] = 1

        matrices[matrix_id] = m
        matrix_id += 1
    end

    for target_row in filter(i -> i ≠ Y-1, 2:Y)
        m = zeros(Int64, Y,Y)

        max_d = target_row in 1:d-1 ? d : d-1
        for id in 1:max_d
            m[id,id] = 1
        end

        m[target_row,max_d+1:end] .= 1

        matrices[matrix_id] = m
        matrix_id += 1
    end

    m = zeros(Int64 ,Y,Y)

    for id in 1:d-2
        m[id,id] = 1
    end

    m[Y,d-1:Y-1] .= 1
    m[Y-1,Y] = 1

    matrices[matrix_id] = m

    matrices
end

"""
    aff_ind_anti_guessing_vertices(N :: Int64, d :: Int64, ε :: Int64) :: Vector{Matrix{Int64}}

Enumerates an affinely independent set of vertices for the [`anti_guessing_facet`](@ref).

Valid inputs are `N ≥ 4`, `(N - 1) > d > 1`, and `(N - d + 1) ≥ ε ≥ 3`.
"""
function aff_ind_anti_guessing_vertices(N :: Int64, d :: Int64, ε :: Int64) :: Vector{Matrix{Int64}}
    if !(N ≥ 4)
        throw(DomainError(N, "N must satisfy `N ≥ 4`."))
    elseif !(N - 1 > d > 1)
        throw(DomainError(d, "d must satisfy `(N - 1) > d > 1`."))
    elseif !((N - d + 1) ≥ ε ≥ 3)
        throw(DomainError(ε, "ε must satisfy `(N - d + 1) ≥ ε ≥ 3`"))
    end

    # affiely independent vectors used to enumerate error block strategies
    error_vecs = sort(_aff_ind_vecs(d-1, N - (ε + d - 2))) # sorting ensures the first element has a leading 0.

    # affinely independent vectors used to enumerate success block strategies
    success_vecs = _aff_ind_vecs(ε + d - 3, N - ε - d + 2)

    # initializing vector to hold matrix data
    matrices = Vector{Matrix{Int64}}(undef, N*(N-1))
    matrix_id = 1
    for row_id in 1:N

        # rows containing the error block are treated separately
        if row_id <= ε

            # creating a strategy for each column in the error block excludin
            # the zero column at col_id = row_id.
            for i in filter(j -> j != row_id, 1:ε)
                m = zeros(Int64, N,N)

                # maximize error block
                m[row_id, i] = 1
                m[i, filter(j -> j != i, 1:ε)] .= 1

                # filling success block to saturate inequality
                for col_id in (ε+1):N
                    el = error_vecs[1][col_id - (ε - 1)]

                    if el == 0
                        m[col_id,col_id] = 1
                    else
                        m[row_id,col_id] = 1
                    end
                end

                # add matrix m to matrices
                matrices[matrix_id] = m
                matrix_id += 1
            end

            # enumerating the affinely ind. strategies for columns greater than
            # error block
            for vec in error_vecs[2:end]
                m = zeros(Int64, N,N)

                # setting error block columns
                m[row_id, filter(j -> j != row_id, 1:ε)] .= 1
                if vec[1] == 0
                    r_id = row_id == 1 ? 2 : 1
                    m[r_id,row_id] = 1
                else
                    m[row_id,row_id] = 1
                end

                # setting success block columns
                for col_id in (ε+1):N
                    el = vec[col_id - (ε - 1)]

                    if el == 0
                        m[col_id,col_id] = 1
                    else
                        m[row_id,col_id] = 1
                    end
                end

                # add matrix m to matrices
                matrices[matrix_id] = m
                matrix_id += 1
            end
        else
            # enumerating strategies for success block rows
            for vec in success_vecs
                m = zeros(Int64, N,N)
                m[row_id,row_id] = 1

                c_id = findfirst(j -> j == 1, vec)

                # filling in error block columns
                if c_id <= ε
                    m[row_id, c_id] = 1
                    m[c_id, filter(j -> j != c_id, 1:ε)] .= 1
                else
                    c_id = 1
                    r_id = d > 2 ? 2 : 1

                    m[r_id, c_id] = 1
                    m[c_id, filter(j -> j != c_id, 1:ε)] .= 1
                end

                # filling in success block columns
                for i in ε+1:N-1
                    el = vec[i]

                    col_id = i < row_id ? i : i + 1
                    if el == 1
                        m[row_id,col_id] = 1
                    else
                        # if d == 2, we need to place all 0 elements are mapped
                        # to first row
                        r_id = d > 2 ? col_id : c_id

                        m[r_id,col_id] = 1
                    end
                end

                # add matrix m to marices
                matrices[matrix_id] = m
                matrix_id += 1
            end
        end
    end

    return matrices
end

"""
    aff_ind_k_guessing_vertices(
        Y :: Int64, d :: Int64, k :: Int64
    ) :: Vector{Matrix{Int64}}

Enumerates an affinely independent set of vertices for the [`k_guessing_facet`](@ref).

A `DomainError` is thrown if the inputs do not satisfy `Y ≥ 4`, `d ≥ 2` and `k ≥ 2`.
"""
function aff_ind_k_guessing_vertices(Y :: Int64, d :: Int64, k :: Int64) :: Vector{Matrix{Int64}}
    if !(Y == k + d)
        throw(DomainError(Y, "`Y = k + d` must be satisfied"))
    elseif !( Y ≥ 4 && d ≥ 2 && k ≥ 2 )
        throw(DomainError((Y, k, d), "`Y ≥ 4`, `d ≥ 2`, and `k ≥ 2` must be satisfied"))
    end

    if d == 2
        return _aff_ind_k_guessing_vertices_d2(Y, d, k)
    elseif k == 2
        return _aff_ind_k_guessing_vertices_k2(Y, d, k)
    else
        matrices = Vector{Matrix{Int64}}(undef, (Y-1)*binomial(Y,k))
        matrix_id = 1

        # getting affinely independent matrices in the bottom right subblock
        br_subblock_matrices = aff_ind_k_guessing_vertices(Y-1, d-1, k)

        # expanding bottom right subblock strategies with constant top row
        for br_subblock in br_subblock_matrices
            m = zeros(Int64, Y, binomial(Y,k))
            m[2:end, (binomial(Y-1,k-1)+1):end] = br_subblock
            m[1, 1:(binomial(Y-1,k-1))] .= 1

            matrices[matrix_id] = m
            matrix_id += 1
        end

        br_subblock_game = k_guessing_facet(Y-1,d-1,k)

        # placing a 1 in each 0 element of top row
        for combo in combinations(1:Y-1, d-1)
            m = zeros(Int64, Y, binomial(Y,k))
            m[1, 1:(binomial(Y-1,k-1))] .= 1

            for combo_id in combo
                m[1+combo_id, (binomial(Y-1,k-1)+1):end] = br_subblock_game[combo_id,:]

                for col_id in (binomial(Y-1,k-1)+1):binomial(Y,k)
                    if sum(m[:,col_id]) > 1
                        m[1+combo_id, col_id] = 0
                    end
                end
            end

            for col_id in (binomial(Y-1,k-1)+1):binomial(Y,k)
                if sum(m[:,col_id]) == 0
                    m[1,col_id] = 1
                end
            end

            matrices[matrix_id] = m
            matrix_id += 1
        end

        bl_subblock_game = k_guessing_facet(Y-1,d,k-1)

        # placing a 0 in each 1 element of the top row
        for col_id in 1:binomial(Y-1,k-1)
            m = zeros(Int64, Y, binomial(Y,k))
            m[1,filter(i -> i != col_id, 1:binomial(Y-1,k-1))] .= 1

            target_row_id = findfirst(i -> i == 1, bl_subblock_game[:,col_id]) + 1

            m[target_row_id, col_id] = 1

            max_row_id = (target_row_id <= d) ? d : target_row_id

            for row_id in [2:d-1..., max_row_id]
                m[row_id, (binomial(Y-1,k-1)+1):end] = br_subblock_game[row_id-1,:]

                for col_id in (binomial(Y-1,k-1)+1):binomial(Y,k)
                    if sum(m[:,col_id]) > 1
                        m[row_id, col_id] = 0
                    end
                end
            end

            for col_id in (binomial(Y-1,k-1)+1):binomial(Y,k)
                if sum(m[:,col_id]) == 0
                    m[max_row_id, col_id] = 1
                end
            end

            matrices[matrix_id] = m
            matrix_id += 1
        end

        # affinely independent matrices for bottom left subblock
        bl_subblock_matrices = aff_ind_k_guessing_vertices(Y-1,d,k-1)

        # enumerating strategies in bottom left subblock not using top row
        for bl_subblock in bl_subblock_matrices
            m = zeros(Int64, Y, binomial(Y,k))

            m[2:end, 1:binomial(Y-1,k-1)] = bl_subblock

            non_zero_row_ids = filter(i -> sum(bl_subblock[i,1:end]) > 0, 1:Y-1)

            for row_id in non_zero_row_ids
                m[1+row_id, (binomial(Y-1,k-1)+1):end] = br_subblock_game[row_id,:]

                for col_id in (binomial(Y-1,k-1)+1):binomial(Y,k)
                    if sum(m[:,col_id]) > 1
                        m[1+row_id, col_id] = 0
                    end
                end
            end

            for col_id in (binomial(Y-1,k-1)+1):binomial(Y,k)
                if sum(m[:,col_id]) == 0
                    m[non_zero_row_ids[end], col_id] = 1
                end
            end

            matrices[matrix_id] = m
            matrix_id += 1
        end

        return matrices
    end
end

"""
    _aff_ind_k_guessing_vertices_k2(
        N :: Int64, d :: Int64, k :: Int64
    ) :: Vector{Matrix{Int64}}

Helper method for aff_ind_k_guessing_vertices.
Enumerates an affinely independent set of vertices for the generalized error game.
The vertices are `N x binomial(N,k)`, rank-`d`, column stochastic matrices.
This function is restricted to `k = 2` and `N = d + k`.

Valid inputs are `N > 4`, `k = 2`, and `d = N - k`.
"""
function _aff_ind_k_guessing_vertices_k2(N :: Int64, d :: Int64, k :: Int64) :: Vector{Matrix{Int64}}
    if !(k == 2)
        throw(DomainError(k, "k > 2 not supported by this function"))
    end

    if N == 4 && d == 2
        return _aff_ind_k_guessing_vertices_d2(N, d, k)
    else
        matrices = Vector{Matrix{Int64}}(undef, (N-1)*binomial(N,k))
        matrix_id = 1

        # expanding bottom right subblock strategies with constant top row
        br_subblock_matrices = _aff_ind_k_guessing_vertices_k2(N-1, d-1, k)

        for br_subblock in br_subblock_matrices
            m = zeros(Int64, N, binomial(N,k))
            m[1,1:d+1] .= 1
            m[2:end, d+2:end] = br_subblock

            matrices[matrix_id] = m
            matrix_id += 1
        end

        # placing a zero in each non-zero top row index
        for col_id in 1:N-1
            m = zeros(Int64, N, binomial(N,k))
            m[1+col_id,col_id] = 1
            m[1, filter(i -> i != col_id, 1:N-1)] .= 1

            col_marker = N
            for row_id in 2:d
                if row_id != d
                    m[row_id, col_marker:col_marker+N-row_id-1] .= 1
                    col_marker += N - row_id
                else
                    if col_id <= d-1
                        m[row_id, col_marker:end] .= 1
                    else
                        m[1+col_id, col_marker:end] .= 1
                    end
                end
            end

            matrices[matrix_id] = m
            matrix_id += 1
        end

        # enumerating all d-success game strategies
        subblock_game = k_guessing_facet(N-1,d-1,k)

        for combo in combinations(1:N-1, d)
            m = zeros(Int64, N, binomial(N,k))

            for combo_id in combo
                m[1+combo_id, N:end] = subblock_game[combo_id,:]

                for col_id in N:binomial(N,k)
                    if sum(m[:,col_id]) > 1
                        m[1+combo_id, col_id] = 0
                    end
                end

                m[1+combo_id, combo_id] = 1
            end

            target_col_id = filter(i -> !(i in combo), 1:N-1)[1]
            for combo_id in combo
                m_copy = copy(m)

                m_copy[1 + combo_id, target_col_id] = 1

                matrices[matrix_id] = m_copy
                matrix_id += 1
            end
        end

        # placing a 1 in each 0 element of top row
        for combo in combinations(1:N-1, d-1)
            m = zeros(Int64, N, binomial(N,k))
            m[1,1:N-1] .= 1

            for combo_id in combo
                m[1+combo_id, N:end] = subblock_game[combo_id,:]

                for col_id in N:binomial(N,k)
                    if sum(m[:,col_id]) > 1
                        m[1+combo_id, col_id] = 0
                    end
                end
            end

            for col_id in N:binomial(N,k)
                if sum(m[:,col_id]) == 0
                    m[1,col_id] = 1
                end
            end

            matrices[matrix_id] = m
            matrix_id += 1
        end

        return matrices
    end
end

"""
    _aff_ind_k_guessing_vertices_d2(
        N :: Int64, d :: Int64, k :: Int64
    ) :: Vector{Matrix{Int64}}

Helper method for aff_ind_k_guessing_vertices.
Enumerates an affinely independent set of vertices for the generalized error game.
The vertices are `N x binomial(N,k)`, rank-`d`, column stochastic matrices.
This function is restricted to games where `d = 2` and `N = d + k`.

Valid inputs are `N > 4`, `d = 2`, and `k = N - d`.
"""
function _aff_ind_k_guessing_vertices_d2(N :: Int64, d :: Int64, k :: Int64) :: Vector{Matrix{Int64}}
    if !(d==2)
        throw(DomainError(d, "d > 2 not implemented yet"))
    end

    matrices = Vector{Matrix{Int64}}(undef, (N-1)*binomial(N,k))
    matrix_id = 1

    game = k_guessing_facet(N,d,k)

    considered_cols = []
    prev_target_row = 1
    for target_rows in combinations(1:N, d)
        target_row = target_rows[1]
        considered_cols = (prev_target_row == target_row) ? considered_cols : []

        for target_col_id in 1:binomial(N,k)
            if target_col_id in considered_cols
                continue
            end

            m = zeros(Int64, N, binomial(N,k))

            m[target_rows[1],:] = game[target_rows[1],:]
            m[target_rows[2],:] = game[target_rows[2],:]

            if m[target_rows[1], target_col_id] == 0 && m[target_rows[2], target_col_id] == 0
                for col_id in filter(i -> i != target_col_id, 1:binomial(N,k))
                    if m[target_rows[1], col_id] == m[target_rows[2], col_id]
                        m[target_rows[1], col_id] = 1
                        m[target_rows[2], col_id] = 0
                    end
                end

                m1 = copy(m)
                m1[target_rows[1], target_col_id] = 1

                matrices[matrix_id] = m1
                matrix_id += 1

                m[target_rows[2], target_col_id] = 1

                matrices[matrix_id] = m
                matrix_id += 1

                push!(considered_cols, target_col_id)
            elseif m[target_rows[1], target_col_id] == m[target_rows[2], target_col_id]
                m[target_rows[1], target_col_id] = 0

                for col_id in filter(i -> i != target_col_id, 1:binomial(N,k))
                    if m[target_rows[1], col_id] == m[target_rows[2], col_id]
                        m[target_rows[1], col_id] = 1
                        m[target_rows[2], col_id] = 0
                    end
                end

                matrices[matrix_id] = m
                matrix_id += 1

                push!(considered_cols, target_col_id)
            end
        end

        prev_target_row = target_row
    end

    return matrices
end

"""
Helper function for `aff_ind_maximum_likelihood_vertices()`. Generates a set of binary
vectors with `num_zeros` and `num_ones`. The set contains `n` affinely independent
vectors where `n` is the length of the binary vectors.
"""
function _aff_ind_vecs(num_zeros :: Int64, num_ones :: Int64) :: Vector{Vector{Int64}}
    n = num_ones + num_zeros

    aff_ind_rows = Vector{Vector{Int64}}(undef, n)
    if num_zeros == 0 || num_ones == 0
        val = (num_zeros == 0) ? 1 : 0
        aff_ind_rows = [fill(val, n)]
    else
        for i in 1:n
            row = zeros(Int64, n)
            if i <= num_zeros + 1
                row[i] = 1
                row[num_zeros+2:end] .= 1
            else
                row[i+1:end] .= 1
                row[num_zeros:i-1] .= 1
            end

            aff_ind_rows[i] = row
        end

    end

    return aff_ind_rows
end


text:
"""
nmr - Naive MapReduce module

NMR is a Naive MapReduce implementation.
It has a simple interface that allows the user to input a filename and process it in a map-reduce fashion.
"""
module nmr
  using JSON

  pkgdir = Pkg.dir("nmr")
  datadir = "NMR_DATADIR" in keys(ENV) ? ENV["NMR_DATADIR"] : joinpath(pkgdir, "test", "data")
  splitdir = "NMR_SPLITDIR" in keys(ENV) ? ENV["NMR_SPLITDIR"] : joinpath(pkgdir, "test", "split")
  outputdir = "NMR_OUTPUTDIR" in keys(ENV) ? ENV["NMR_OUTPUTDIR"] : joinpath(pkgdir, "test", "output")

  ctx = Dict{String, Any}()
  default_ctx = Dict{String, Any}()

  # required for correct module operation
  include("logging_config.jl")
  include("types.jl")
  include("io.jl")
  include("phases.jl")

  # user defined functions and types
  include("udf/types.jl")
  include("udf/mappers.jl")
  include("udf/joiners.jl")
  include("udf/reducers.jl")
  include("udf/combiners.jl")

  function runjob(j)
    info("Running job: $j")
    split_raw_data(j)
    create_job_area(j)
    phase_run(j)
    ctx = copy(default_ctx)

    return true
  end

end # module

ext: jl

max_stars_repo_name: deyandyankov/nmr

max_issues_repo_name: deyandyankov/nmr

max_forks_repo_name: deyandyankov/nmr


text:
##############################################################################
"""
    plt = plotscheck("simlation/path/_output"::AbstractString; vartype::Symbol=:surface, AMRlevel=[], runup=true, region="", kwargs...)

Quick checker of the spatial distribution
"""
function plotscheck(simdir::AbstractString; vartype::Symbol=:surface, AMRlevel=[], runup::Bool=true, region="", testplot::Bool=false, kwargs...)

    !any([vartype==s for s in [:surface, :storm, :current]]) && error("Invalid input argument vartype: $vartype")
    ## define the filepath & filename
    if vartype==:surface
        fnamekw = r"^fort\.q\d+$"
        loadfunction = VisClaw.loadsurface
        kwargs_load = Dict([(:runup, runup)])
    elseif vartype==:current
        fnamekw = r"^fort\.q\d+$"
        loadfunction = VisClaw.loadcurrent
        kwargs_load = Dict([])
    elseif vartype==:storm
        fnamekw = r"^fort\.a\d+$"
        loadfunction = VisClaw.loadstorm
        kwargs_load = Dict([])
    end

    ## parse keyword args
    kwdict = KWARG(kwargs)
    ## xlims, ylims
    xlims, kwdict = VisClaw.kwarg_default(kwdict, VisClaw.parse_xlims, (-Inf,Inf))
    ylims, kwdict = VisClaw.kwarg_default(kwdict, VisClaw.parse_ylims, (-Inf,Inf))
    ## set range
    if isa(region, VisClaw.AbstractTopo); xlims=extrema(region.x); ylims=extrema(region.y); end
    if isa(region, VisClaw.Region); xlims=region.xlims; ylims=region.ylims; end

    ## make a list
    isdir(simdir) || error("Not found: directory $simdir")
    flist = readdir(simdir)
    filter!(x->occursin(fnamekw, x), flist)
    isempty(flist) && error("File named $fnamekw was not found")

    ## load geoclaw.data
    params = VisClaw.geodata(simdir)

    ## the number of files
    nfile = length(flist)

    ### draw figures until nothing or invalid number is input
    println(">> Press ENTER with a blank to finish")
    println(">> Input a file sequence number you want to plot:")
    ex=0 # initial value
    cnt=0
    while ex==0
        ## accept input the step number of interest
        @printf("checkpoint time (1 to %d) = ", nfile)
        i = testplot ? "1" : readline(stdin)
        if testplot; ex=1; end

        ## check whether the input is integer
        if isempty(i); ex=1; continue; end;
        ## parse to interger
        i = try; parse(Int64,i)
        catch; "cannot be parsed to integer"; ex=1; continue;
        end
        ## check whether the input is valid number
        if ( (i>nfile) || (i<1) ); println("Invalid number"); ex=1; continue; end

        ## load data
        amrs = loadfunction(simdir, i; AMRlevel=AMRlevel, xlims=xlims, ylims=ylims, kwargs_load...)
        runup && (coarsegridmask!(amrs))

        ## draw figure
        plt = VisClaw.plotsamr2d(amrs.amr[1]; AMRlevel=AMRlevel, xlims=xlims, ylims=ylims, kwdict...)
        plt = Plots.plot!(plt, title=@sprintf("%8.1f",amrs.timelap[1])*" s")

        ## show
        #plt = Plots.plot!(plt, show=true)
        display(plt)
        cnt += 1
    end

    ## if no plot is done
    if cnt==0; plt = nothing; end

    ## return value
    return plt
end
##############################################################################


text:
"""
    Simulation(
        steps::Int
        models::SimulationModels,
        sequence::Union{Nothing, SimulationSequence},
        simulation_folder::String,
        name::String,
        internal::Union{Nothing, SimulationInternal},
    )
"""
mutable struct Simulation
    steps::Int
    models::SimulationModels
    initial_time::Union{Nothing, Dates.DateTime}
    sequence::SimulationSequence
    simulation_folder::String
    name::String
    internal::Union{Nothing, SimulationInternal}

    function Simulation(;
        sequence::SimulationSequence,
        name::String,
        steps::Int,
        models::SimulationModels,
        simulation_folder::String,
        initial_time=nothing,
    )
        for model in get_decision_models(models)
            if model.internal.simulation_info.sequence_uuid != sequence.uuid
                model_name = get_name(model)
                throw(
                    IS.ConflictingInputsError(
                        "The decision model definition for $model_name doesn't correspond to the simulation sequence",
                    ),
                )
            end
        end
        em = get_emulation_model(models)
        if em !== nothing
            if em.internal.simulation_info.sequence_uuid != sequence.uuid
                model_name = get_name(em)
                throw(
                    IS.ConflictingInputsError(
                        "The emulation model definition for $model_name doesn't correspond to the simulation sequence",
                    ),
                )
            end
        end
        new(steps, models, initial_time, sequence, simulation_folder, name, nothing)
    end
end

"""
    Simulation(directory::AbstractString)

Constructs Simulation from a serialized directory. Callers should pass any kwargs here that
they passed to the original Simulation.

# Arguments

  - `directory::AbstractString`: the directory returned from the call to serialize
  - `model_info::Dict`: Two-level dictionary containing model parameters that cannot be
    serialized. The outer dict should be keyed by the problem name. The inner dict must contain
    'optimizer' and may contain 'jump_model'. These should be the same values used for the
    original simulation.
"""
function Simulation(directory::AbstractString, model_info::Dict)
    return deserialize_model(Simulation, directory, model_info)
end

###################### Simulation Accessor Functions ####################
function get_base_powers(sim::Simulation)
    base_powers = Dict()
    for model in get_models(sim)
        base_powers[get_name(model)] = PSY.get_base_power(get_system(model))
    end
    return base_powers
end

get_initial_time(sim::Simulation) = sim.initial_time
get_sequence(sim::Simulation) = sim.sequence
get_steps(sim::Simulation) = sim.steps
get_current_time(sim::Simulation) = get_current_time(get_simulation_state(sim))
get_simulation_model(s::Simulation, name) = get_simulation_model(get_models(s), name)
get_models(sim::Simulation) = sim.models
get_simulation_dir(sim::Simulation) = dirname(sim.internal.logs_dir)
get_simulation_files_dir(sim::Simulation) = sim.internal.sim_files_dir
get_store_dir(sim::Simulation) = sim.internal.store_dir
get_simulation_status(sim::Simulation) = sim.internal.status
get_simulation_build_status(sim::Simulation) = sim.internal.build_status
get_simulation_state(sim::Simulation) = sim.internal.simulation_state
set_simulation_store!(sim::Simulation, store) = sim.internal.store = store
get_simulation_store(sim::Simulation) = sim.internal.store
get_results_dir(sim::Simulation) = sim.internal.results_dir
get_models_dir(sim::Simulation) = sim.internal.models_dir

get_interval(sim::Simulation, name::Symbol) = get_interval(sim.sequence, name)

function get_simulation_time(sim::Simulation, problem_number::Int)
    return sim.internal.date_ref[problem_number]
end

get_ini_cond_chronology(sim::Simulation) = get_sequence(sim).ini_cond_chronology
get_name(sim::Simulation) = sim.name
get_simulation_folder(sim::Simulation) = sim.simulation_folder
get_execution_order(sim::Simulation) = get_sequence(sim).execution_order
get_current_execution_index(sim::Simulation) = get_sequence(sim).current_execution_index
get_logs_folder(sim::Simulation) = sim.internal.logs_dir
get_recorder_folder(sim::Simulation) = sim.internal.recorder_dir
get_console_level(sim::Simulation) = sim.internal.console_level
get_file_level(sim::Simulation) = sim.internal.file_level

set_simulation_status!(sim::Simulation, status) = sim.internal.status = status
set_simulation_build_status!(sim::Simulation, status::BuildStatus) =
    sim.internal.build_status = status

function set_current_time!(sim::Simulation, val::Dates.DateTime)
    set_current_time!(get_simulation_state(sim), val)
    return
end

function _get_simulation_initial_times!(sim::Simulation)
    model_initial_times = OrderedDict{Int, Vector{Dates.DateTime}}()
    sim_ini_time = get_initial_time(sim)
    for (model_number, model) in enumerate(get_models(sim).decision_models)
        system = get_system(model)
        model_horizon = get_horizon(model)
        system_horizon = PSY.get_forecast_horizon(system)
        system_interval = PSY.get_forecast_interval(system)
        if model_horizon > system_horizon
            throw(
                IS.ConflictingInputsError(
                    "$(get_name(model)) model horizon ($model_horizon) and forecast horizon ($system_horizon) are not compatible",
                ),
            )
        end
        model_initial_times[model_number] = PSY.get_forecast_initial_times(system)
        for (ix, element) in enumerate(model_initial_times[model_number][1:(end - 1)])
            if !(element + system_interval == model_initial_times[model_number][ix + 1])
                throw(
                    IS.ConflictingInputsError(
                        "The sequence of forecasts in the model's systems are invalid",
                    ),
                )
            end
        end
        if sim_ini_time !== nothing &&
           !mapreduce(x -> x == sim_ini_time, |, model_initial_times[model_number])
            throw(
                IS.ConflictingInputsError(
                    "The specified simulation initial_time $sim_ini_time isn't contained in model $(get_name(model)).
Manually provided initial times have to be compatible with the specified interval and horizon in the models.",
                ),
            )
        end
    end
    if get_initial_time(sim) === nothing
        sim.initial_time = model_initial_times[1][1]
        @debug("Initial Simulation timestamp will be infered from the data. \\
               Initial Simulation timestamp set to $(sim.initial_time)")
    end
    if get_models(sim).emulation_model !== nothing
        em = get_models(sim).emulation_model
        system = get_system(get_models(sim).emulation_model)
        ini_time, ts_length =
            PSY.check_time_series_consistency(system, PSY.SingleTimeSeries)
        resolution = PSY.get_time_series_resolution(system)
        em_available_times = range(ini_time, step=resolution, length=ts_length)
        if get_initial_time(sim) ∉ em_available_times
            throw(
                IS.ConflictingInputsError(
                    "The simulation initial_time $sim_ini_time isn't contained in the
                    emulation model $(get_name(em)).",
                ),
            )
        else
            model_initial_times[length(model_initial_times) + 1] = [sim.initial_time]
        end
    end
    set_current_time!(sim, sim.initial_time)
    return model_initial_times
end

function _check_steps(
    sim::Simulation,
    model_initial_times::OrderedDict{Int, Vector{Dates.DateTime}},
)
    sequence = get_sequence(sim)
    execution_order = get_execution_order(sequence)
    for (model_number, model) in enumerate(get_models(sim).decision_models)
        execution_counts = get_executions(model)
        # Checks the consistency between two methods of calculating the number of executions
        total_model_executions = length(findall(x -> x == model_number, execution_order))
        @assert_op total_model_executions == execution_counts

        forecast_count = length(model_initial_times[model_number])
        if get_steps(sim) * execution_counts > forecast_count
            throw(
                IS.ConflictingInputsError(
                    "The number of available time series ($(forecast_count)) is not enough to perform the
desired amount of simulation steps ($(sim.steps*get_execution_count(model))).",
                ),
            )
        end
    end
    return
end

function _check_folder(sim::Simulation)
    folder = get_simulation_folder(sim)
    !isdir(folder) && throw(IS.ConflictingInputsError("Specified folder is not valid"))
    try
        mkdir(joinpath(folder, "fake"))
        rm(joinpath(folder, "fake"))
    catch e
        throw(IS.ConflictingInputsError("Specified folder does not have write access [$e]"))
    end
end

function _build_decision_models!(sim::Simulation)
    for (model_number, model) in enumerate(get_decision_models(get_models(sim)))
        @info("Building problem $(get_name(model))")
        initial_time = get_initial_time(sim)
        set_initial_time!(model, initial_time)
        output_dir = joinpath(get_models_dir(sim), string(get_name(model)))
        mkpath(output_dir)
        set_output_dir!(model, output_dir)
        try
            TimerOutputs.@timeit BUILD_PROBLEMS_TIMER "Problem $(get_name(model))" begin
                # TODO-PJ: Temporary while are able to switch from PJ to POI
                container = get_optimization_container(model)
                container.built_for_recurrent_solves = true
                build_impl!(model)
            end
            sim.internal.date_ref[model_number] = initial_time
            set_status!(model, BuildStatus.BUILT)
            # TODO: Disable check of variable bounds ?
            _pre_solve_model_checks(model)
        catch
            set_status!(model, BuildStatus.FAILED)
            rethrow()
        end
    end
    return
end

function _build_emulation_model!(sim::Simulation)
    model = get_emulation_model(get_models(sim))

    if model === nothing
        return
    end

    try
        initial_time = get_initial_time(sim)
        set_initial_time!(model, initial_time)
        output_dir = joinpath(get_models_dir(sim), string(get_name(model)))
        mkpath(output_dir)
        set_output_dir!(model, output_dir)
        TimerOutputs.@timeit BUILD_PROBLEMS_TIMER "Problem $(get_name(model))" begin
            build_impl!(model)
        end
        sim.internal.date_ref[length(sim.internal.date_ref) + 1] = initial_time
        set_status!(model, BuildStatus.BUILT)
    catch
        set_status!(model, BuildStatus.FAILED)
        rethrow()
    end
    return
end

function _initialize_simulation_state!(sim::Simulation)
    step_resolution = get_step_resolution(get_sequence(sim))
    simulation_models = get_models(sim)
    initialize_simulation_state!(
        get_simulation_state(sim),
        simulation_models,
        step_resolution,
        get_initial_time(sim),
    )
    return
end

function _get_model_store_requirements!(
    rules::CacheFlushRules,
    model::OperationModel,
    num_rows::Int,
)
    model_name = get_name(model)
    horizon = get_horizon(model)
    reqs = SimulationModelStoreRequirements()
    container = get_optimization_container(model)

    for (key, array) in get_duals(container)
        !should_write_resulting_value(key) && continue
        reqs.duals[key] = _calc_dimensions(array, key, num_rows, horizon)
        add_rule!(rules, model_name, key, true)
    end

    for (key, param_container) in get_parameters(container)
        !should_write_resulting_value(key) && continue
        array = get_parameter_array(param_container)
        reqs.parameters[key] = _calc_dimensions(array, key, num_rows, horizon)
        add_rule!(rules, model_name, key, false)
    end

    for (key, array) in get_variables(container)
        !should_write_resulting_value(key) && continue
        reqs.variables[key] = _calc_dimensions(array, key, num_rows, horizon)
        add_rule!(rules, model_name, key, true)
    end

    for (key, array) in get_aux_variables(container)
        !should_write_resulting_value(key) && continue
        reqs.aux_variables[key] = _calc_dimensions(array, key, num_rows, horizon)
        add_rule!(rules, model_name, key, true)
    end

    for (key, array) in get_expressions(container)
        !should_write_resulting_value(key) && continue
        reqs.expressions[key] = _calc_dimensions(array, key, num_rows, horizon)
        add_rule!(rules, model_name, key, false)
    end

    return reqs
end

function _get_emulation_store_requirements(sim::Simulation)
    sim_state = get_simulation_state(sim)
    system_state = get_system_states(sim_state)
    sim_time = get_steps(sim) * get_step_resolution(get_sequence(sim))
    reqs = SimulationModelStoreRequirements()

    for (key, state_values) in get_duals_values(system_state)
        !should_write_resulting_value(key) && continue
        dims = sim_time ÷ get_data_resolution(state_values)
        cols = get_column_names(key, state_values)
        reqs.duals[key] = Dict("columns" => cols, "dims" => (dims, length(cols)))
    end

    for (key, state_values) in get_parameters_values(system_state)
        !should_write_resulting_value(key) && continue
        dims = sim_time ÷ get_data_resolution(state_values)
        cols = get_column_names(key, state_values)
        reqs.parameters[key] = Dict("columns" => cols, "dims" => (dims, length(cols)))
    end

    for (key, state_values) in get_variables_values(system_state)
        !should_write_resulting_value(key) && continue
        dims = sim_time ÷ get_data_resolution(state_values)
        cols = get_column_names(key, state_values)
        reqs.variables[key] = Dict("columns" => cols, "dims" => (dims, length(cols)))
    end

    for (key, state_values) in get_aux_variables_values(system_state)
        !should_write_resulting_value(key) && continue
        dims = sim_time ÷ get_data_resolution(state_values)
        cols = get_column_names(key, state_values)
        reqs.aux_variables[key] = Dict("columns" => cols, "dims" => (dims, length(cols)))
    end

    for (key, state_values) in get_expression_values(system_state)
        !should_write_resulting_value(key) && continue
        dims = sim_time ÷ get_data_resolution(state_values)
        cols = get_column_names(key, state_values)
        reqs.expressions[key] = Dict("columns" => cols, "dims" => (dims, length(cols)))
    end
    return reqs
end

function _initialize_problem_storage!(
    sim::Simulation,
    cache_size_mib,
    min_cache_flush_size_mib,
)
    sequence = get_sequence(sim)
    executions_by_model = sequence.executions_by_model
    models = get_models(sim)
    decision_model_store_params = OrderedDict{Symbol, ModelStoreParams}()
    dm_model_req = Dict{Symbol, SimulationModelStoreRequirements}()
    rules = CacheFlushRules(
        max_size=cache_size_mib * MiB,
        min_flush_size=trunc(min_cache_flush_size_mib * MiB),
    )
    for model in get_decision_models(models)
        model_name = get_name(model)
        decision_model_store_params[model_name] = model.internal.store_parameters
        num_executions = executions_by_model[model_name]
        num_rows = num_executions * get_steps(sim)
        dm_model_req[model_name] = _get_model_store_requirements!(rules, model, num_rows)
    end

    em = get_emulation_model(models)
    if em === nothing
        base_params = last(collect(values(decision_model_store_params)))
        resolution = minimum([v.resolution for v in values(decision_model_store_params)])
        emulation_model_store_params = OrderedDict(
            :Emulator => ModelStoreParams(
                get_step_resolution(sequence) ÷ resolution, # Num Executions
                1,
                resolution, # Interval
                resolution, # Resolution
                get_base_power(base_params),
                get_system_uuid(base_params),
            ),
        )
    else
        emulation_model_store_params =
            OrderedDict(Symbol(get_name(em)) => em.internal.store_parameters)
    end

    em_model_req = _get_emulation_store_requirements(sim)

    simulation_store_params = SimulationStoreParams(
        get_initial_time(sim),
        get_step_resolution(sequence),
        get_steps(sim),
        decision_model_store_params,
        emulation_model_store_params,
    )
    @debug "initialized problem requirements" simulation_store_params
    store = get_simulation_store(sim)

    initialize_problem_storage!(
        store,
        simulation_store_params,
        dm_model_req,
        em_model_req,
        rules,
    )
    return simulation_store_params
end

function _build!(sim::Simulation, serialize::Bool)
    set_simulation_build_status!(sim, BuildStatus.IN_PROGRESS)
    problem_initial_times = _get_simulation_initial_times!(sim)
    sequence = get_sequence(sim)
    step_resolution = get_step_resolution(sequence)
    simulation_models = get_models(sim)
    for (ix, model) in enumerate(get_decision_models(simulation_models))
        problem_interval = get_interval(sequence, model)
        # Note to devs: Here we are setting the number of operations problem executions we
        # will see for every step of the simulation. The step of the simulation is determined
        # by the first decision problem interval
        if ix == 1
            set_executions!(model, 1)
        else
            if step_resolution % problem_interval != Dates.Millisecond(0)
                error(
                    "The $(get_name(model)) problem interval is not an integer fraction of the simulation step",
                )
            end
            set_executions!(model, Int(step_resolution / problem_interval))
        end
    end

    em = get_emulation_model(simulation_models)
    if em !== nothing
        system = get_system(em)
        em_resolution = PSY.get_time_series_resolution(system)
        set_executions!(em, get_steps(sim) * Int(step_resolution / em_resolution))
    end

    TimerOutputs.@timeit BUILD_PROBLEMS_TIMER "Check Steps" begin
        _check_steps(sim, problem_initial_times)
    end

    TimerOutputs.@timeit BUILD_PROBLEMS_TIMER "Build Problems" begin
        _build_decision_models!(sim)
        _build_emulation_model!(sim)
    end

    TimerOutputs.@timeit BUILD_PROBLEMS_TIMER "Initialize Simulation State" begin
        _initialize_simulation_state!(sim)
    end

    if serialize
        TimerOutputs.@timeit BUILD_PROBLEMS_TIMER "Serializing Simulation Files" begin
            serialize_simulation(sim)
        end
        for model in get_decision_models(simulation_models)
            serialize_problem(model)
        end
        if em !== nothing
            serialize_problem(em)
        end
    end
    return
end

function _set_simulation_internal!(
    sim::Simulation,
    output_dir,
    recorders,
    console_level,
    file_level,
)
    sim.internal = SimulationInternal(
        sim.steps,
        get_models(sim),
        get_simulation_folder(sim),
        get_name(sim),
        output_dir,
        recorders,
        console_level,
        file_level,
    )
    return
end

"""
    build!(sim::Simulation)

Build the Simulation, problems and the related folder structure

# Arguments

  - `sim::Simulation`: simulation object
  - `output_dir` = nothing: Name of the output directory for the simulation. If nothing, the
    folder will have the same name as the simulation
  - `serialize::Bool = true`: serializes the simulation objects in the simulation
  - `recorders::Vector{Symbol} = []`: recorder names to register
  - `console_level = Logging.Error`:
  - `file_level = Logging.Info`:

Throws an exception if name is passed and the directory already exists.
"""
function build!(
    sim::Simulation;
    output_dir=nothing,
    recorders=[],
    console_level=Logging.Error,
    file_level=Logging.Info,
    serialize=true,
    initialize_problem=false,
)
    TimerOutputs.reset_timer!(BUILD_PROBLEMS_TIMER)
    TimerOutputs.@timeit BUILD_PROBLEMS_TIMER "Build Simulation" begin
        TimerOutputs.@timeit BUILD_PROBLEMS_TIMER "Initialize Simulation Internal" begin
            _check_folder(sim)
            _set_simulation_internal!(sim, output_dir, recorders, console_level, file_level)
        end
        file_mode = "w"
        logger = configure_logging(sim.internal, file_mode)
        register_recorders!(sim.internal, file_mode)
        try
            Logging.with_logger(logger) do
                try
                    _build!(sim, serialize)
                    set_simulation_build_status!(sim, BuildStatus.BUILT)
                    set_simulation_status!(sim, RunStatus.READY)
                catch e
                    @error "Simulation build failed" exception = (e, catch_backtrace())
                    set_simulation_build_status!(sim, BuildStatus.FAILED)
                    set_simulation_status!(sim, RunStatus.NOT_READY)
                    rethrow(e)
                end
            end
        finally
            unregister_recorders!(sim.internal)
            close(logger)
        end
    end
    initialize_problem && _initial_conditions_problems!(sim)
    @info "\n$(BUILD_PROBLEMS_TIMER)\n"
    return get_simulation_build_status(sim)
end

function _apply_warm_start!(model::OperationModel)
    container = get_optimization_container(model)
    # If the model was used to retrieve duals from an MILP the logic has to be different and
    # the results need to be read from the primal cache
    if isempty(container.primal_values_cache)
        jump_model = get_jump_model(container)
        all_vars = JuMP.all_variables(jump_model)
        all_vars_value = jump_value.(all_vars)
        JuMP.set_start_value.(all_vars, all_vars_value)
    else
        for (var_key, variable_value) in container.primal_values_cache.variables_cache
            variable = get_variable(container, var_key)
            JuMP.set_start_value.(variable, variable_value)
        end
    end
    return
end

function _get_next_problem_initial_time(sim::Simulation, model_name::Symbol)
    current_time = get_current_time(sim)
    sequence = get_sequence(sim)
    current_exec_index = sequence.current_execution_index
    exec_order = get_execution_order(sequence)

    if length(exec_order) > 1 && (current_exec_index + 1 > length(exec_order)) # Moving to the next step
        next_initial_time = get_simulation_time(sim, exec_order[1])
    elseif length(exec_order) == 1 ||
           exec_order[current_exec_index + 1] == exec_order[current_exec_index] # Solving the same problem again
        current_model_interval = get_interval(sim.sequence, model_name)
        next_initial_time = current_time + current_model_interval
    else # Solving another problem next
        next_initial_time = get_simulation_time(sim, exec_order[current_exec_index + 1])
    end
    return next_initial_time
end

function _update_system_state!(sim::Simulation, model_name::Symbol)
    sim_state = get_simulation_state(sim)
    system_state = get_system_states(sim_state)
    decision_state = get_decision_states(sim_state)
    simulation_time = get_current_time(sim_state)
    next_stage_initial_time = _get_next_problem_initial_time(sim, model_name)

    for key in get_dataset_keys(decision_state)
        state_data = get_dataset(decision_state, key)
        last_update = get_update_timestamp(decision_state, key)

        if last_update > simulation_time
            error("Something went really wrong. Please report this error. \\
            last_update: $(last_update) \\
            simulation_time: $(simulation_time) \\
            key: $(encode_key_as_string(key))")
        end

        resolution = get_data_resolution(state_data)
        update_timestamp = max(next_stage_initial_time - resolution, simulation_time)
        if update_timestamp < get_update_timestamp(system_state, key)
            error("The update overwrites more recent data with past data")
        elseif update_timestamp > get_update_timestamp(system_state, key)
            update_system_state!(system_state, key, decision_state, update_timestamp)
        else
            @assert_op update_timestamp == get_update_timestamp(system_state, key)
        end
    end

    IS.@record :execution StateUpdateEvent(simulation_time, model_name, "SystemState")
    return
end

function _update_system_state!(sim::Simulation, model::DecisionModel)
    _update_system_state!(sim, get_name(model))
    return
end

function _update_system_state!(sim::Simulation, model::EmulationModel)
    sim_state = get_simulation_state(sim)
    simulation_time = get_current_time(sim)
    system_state = get_system_states(sim_state)
    store = get_simulation_store(sim)
    em_model_name = get_name(model)
    for key in get_container_keys(get_optimization_container(model))
        !should_write_resulting_value(key) && continue
        update_system_state!(system_state, key, store, em_model_name, simulation_time)
    end
    IS.@record :execution StateUpdateEvent(simulation_time, em_model_name, "SystemState")
    return
end

function _update_simulation_state!(sim::Simulation, model::EmulationModel)
    # Order of these operations matters. Do not reverse.
    # This will update the state with the results of the store first and then fill
    # the remaning values with the decision state.
    _update_system_state!(sim, model)
    _update_system_state!(sim, get_name(model))
    return
end

function _update_simulation_state!(sim::Simulation, model::DecisionModel)
    model_name = get_name(model)
    store = get_simulation_store(sim)
    simulation_time = get_current_time(sim)
    state = get_simulation_state(sim)
    model_params = get_decision_model_params(store, model_name)
    for field in fieldnames(DatasetContainer)
        for key in list_decision_model_keys(store, model_name, field)
            !has_dataset(get_decision_states(state), key) && continue
            res = read_result(DataFrames.DataFrame, store, model_name, key, simulation_time)
            update_decision_state!(state, key, res, simulation_time, model_params)
        end
    end
    IS.@record :execution StateUpdateEvent(simulation_time, model_name, "DecisionState")
    return
end

function _write_state_to_store!(store::SimulationStore, sim::Simulation)
    sim_state = get_simulation_state(sim)
    system_state = get_system_states(sim_state)
    model_name = get_last_decision_model(sim_state)
    em_store = get_em_data(store)
    simulation_time = get_current_time(sim)
    sim_ini_time = get_initial_time(sim)
    state_resolution = get_system_states_resolution(sim_state)
    for key in get_dataset_keys(system_state)
        store_update_time = get_last_updated_timestamp(em_store, key)
        state_update_time = get_update_timestamp(system_state, key)
        # If the store is outdated w.r.t to the state
        @assert store_update_time <= simulation_time
        if store_update_time < state_update_time
            _update_timestamp = max(store_update_time + state_resolution, sim_ini_time)
            while _update_timestamp <= state_update_time
                state_values = get_decision_state_value(sim_state, key, _update_timestamp)
                ix = get_last_recorded_row(em_store, key) + 1
                write_result!(store, model_name, key, ix, _update_timestamp, state_values)
                _update_timestamp += state_resolution
            end
        end
    end
    return
end

"""
Default problem update function for most problems with no customization
"""
function update_model!(model::OperationModel, sim::Simulation)
    if get_requires_rebuild(model)
        # TODO: Implement this case where the model is re-built
        # build_impl!(model)
    else
        update_model!(model, get_simulation_state(sim), get_ini_cond_chronology(sim))
    end
    return
end

function _execute!(
    sim::Simulation;
    cache_size_mib=DEFAULT_SIMULATION_STORE_CACHE_SIZE_MiB,
    min_cache_flush_size_mib=MIN_CACHE_FLUSH_SIZE_MiB,
    exports=nothing,
    enable_progress_bar=progress_meter_enabled(),
    disable_timer_outputs=false,
)
    @assert sim.internal !== nothing

    set_simulation_status!(sim, RunStatus.RUNNING)
    execution_order = get_execution_order(sim)
    steps = get_steps(sim)
    num_executions = steps * length(execution_order)
    store_params =
        _initialize_problem_storage!(sim, cache_size_mib, min_cache_flush_size_mib)
    status = RunStatus.RUNNING
    if exports !== nothing
        if !(exports isa SimulationResultsExport)
            exports = SimulationResultsExport(exports, store_params)
        end

        if exports.path === nothing
            exports.path = get_results_dir(sim)
        end
    end
    sequence = get_sequence(sim)
    models = get_models(sim)

    prog_bar = ProgressMeter.Progress(num_executions; enabled=enable_progress_bar)
    disable_timer_outputs && TimerOutputs.disable_timer!(RUN_SIMULATION_TIMER)
    store = get_simulation_store(sim)
    for step in 1:steps
        IS.@record :simulation_status SimulationStepEvent(
            get_current_time(sim),
            step,
            "start",
        )
        for (ix, model_number) in enumerate(execution_order)
            model = get_simulation_model(models, model_number)
            model_name = get_name(model)
            set_current_time!(sim, sim.internal.date_ref[model_number])
            sequence.current_execution_index = ix
            current_time = get_current_time(sim)
            IS.@record :simulation_status ProblemExecutionEvent(
                current_time,
                step,
                model_name,
                "start",
            )

            ProgressMeter.update!(
                prog_bar,
                (step - 1) * length(execution_order) + ix;
                showvalues=[
                    (:Step, step),
                    (:Problem, model_name),
                    (:("Simulation Timestamp"), get_current_time(sim)),
                ],
            )

            TimerOutputs.@timeit RUN_SIMULATION_TIMER "Execute $(model_name)" begin
                if !is_built(model)
                    error("$(model_name) status is not BuildStatus.BUILT")
                end

                # Is first run of first problem? Yes -> don't update problem

                TimerOutputs.@timeit RUN_SIMULATION_TIMER "Update $(model_name)" begin
                    !(step == 1 && ix == 1) && update_model!(model, sim)
                end

                TimerOutputs.@timeit RUN_SIMULATION_TIMER "Solve $(model_name)" begin
                    status = solve!(step, model, current_time, store; exports=exports)
                end # Run problem Timer

                TimerOutputs.@timeit RUN_SIMULATION_TIMER "Update State" begin
                    if status == RunStatus.SUCCESSFUL
                        # TODO: _update_simulation_state! can use performance improvements
                        _update_simulation_state!(sim, model)
                        if model_number == execution_order[end]
                            _update_system_state!(sim, model)
                            _write_state_to_store!(store, sim)
                        end
                    end
                end

                sim.internal.run_count[step][model_number] += 1
                sim.internal.date_ref[model_number] += get_interval(sequence, model_name)

                # _apply_warm_start! can only be called once all the operations that read solutions
                # from the optimization container have been called.
                # See https://github.com/NREL-SIIP/PowerSimulations.jl/pull/793#discussion_r761545526
                # for reference
                if warm_start_enabled(model)
                    _apply_warm_start!(model)
                end

                IS.@record :simulation_status ProblemExecutionEvent(
                    get_current_time(sim),
                    step,
                    model_name,
                    "done",
                )
            end #execution problem timer
        end # execution order for loop

        IS.@record :simulation_status SimulationStepEvent(
            get_current_time(sim),
            step,
            "done",
        )
    end # Steps for loop
    return
end

"""
    execute!(sim::Simulation; kwargs...)

Solves the simulation model for sequential Simulations.

# Arguments

  - `sim::Simulation=sim`: simulation object created by Simulation()

The optional keyword argument `exports` controls exporting of results to CSV files as
the simulation runs. Refer to [`export_results`](@ref) for a description of this argument.

# Example

sim = Simulation("Test", 7, problems, "/Users/folder")
execute!(sim::Simulation; kwargs...)
"""
function execute!(sim::Simulation; kwargs...)
    file_mode = "a"
    logger = configure_logging(sim.internal, file_mode)
    register_recorders!(sim.internal, file_mode)

    # Undocumented option for test & dev only.
    in_memory = get(kwargs, :in_memory, false)
    store_type = in_memory ? InMemorySimulationStore : HdfSimulationStore

    if (get_simulation_build_status(sim) != BuildStatus.BUILT) ||
       (get_simulation_status(sim) != RunStatus.READY)
        error("Simulation status is invalid, you need to rebuild the simulation")
    end
    try
        Logging.with_logger(logger) do
            open_store(store_type, get_store_dir(sim), "w") do store
                set_simulation_store!(sim, store)
                try
                    TimerOutputs.reset_timer!(RUN_SIMULATION_TIMER)
                    TimerOutputs.@timeit RUN_SIMULATION_TIMER "Execute Simulation" begin
                        _execute!(sim; [k => v for (k, v) in kwargs if k != :in_memory]...)
                    end
                    @info ("\n$(RUN_SIMULATION_TIMER)\n")
                    set_simulation_status!(sim, RunStatus.SUCCESSFUL)
                    log_cache_hit_percentages(store)
                catch e
                    set_simulation_status!(sim, RunStatus.FAILED)
                    @error "simulation failed" exception = (e, catch_backtrace())
                end
            end
        end
    finally
        _empty_problem_caches!(sim)
        unregister_recorders!(sim.internal)
        close(logger)
    end

    if !in_memory
        compute_file_hash(get_store_dir(sim), HDF_FILENAME)
    end

    serialize_status(sim)
    return get_simulation_status(sim)
end

struct SimulationSerializationWrapper
    steps::Int
    models::Vector{Symbol}
    initial_time::Union{Nothing, Dates.DateTime}
    sequence::Union{Nothing, SimulationSequence}
    simulation_folder::String
    name::String
end

function _empty_problem_caches!(sim::Simulation)
    models = get_models(sim)
    for model in get_decision_models(models)
        empty_time_series_cache!(model)
    end
    return
end

"""
    serialize_simulation(sim::Simulation, path = ".")

Serialize the simulation to a directory in path.

Return the serialized simulation directory name that is created.

# Arguments

  - `sim::Simulation`: simulation to serialize
  - `path = "."`: path in which to create the serialzed directory
  - `force = false`: If true, delete the directory if it already exists. Otherwise, it will
    throw an exception.
"""
function serialize_simulation(sim::Simulation; path=nothing, force=false)
    if path === nothing
        directory = get_simulation_files_dir(sim)
    else
        directory = path
    end
    problems = get_model_names(get_models(sim))

    if !isempty(readdir(directory)) && !force
        throw(
            ArgumentError(
                "$directory has files already: $(readdir(directory)). Please delete them or pass force = true.",
            ),
        )
    end
    rm(directory, recursive=true, force=true)
    mkdir(directory)

    filename = joinpath(directory, SIMULATION_SERIALIZATION_FILENAME)
    obj = SimulationSerializationWrapper(
        get_steps(sim),
        problems,
        get_initial_time(sim),
        get_sequence(sim),
        get_simulation_dir(sim),
        get_name(sim),
    )
    Serialization.serialize(filename, obj)
    @info "Serialized simulation name = $(get_name(sim))" directory
    return directory
end

function deserialize_model(
    ::Type{Simulation},
    directory::AbstractString,
    problem_info::Dict,
)
    error("deserialization of a Simulation is not currently supported")
    orig = pwd()
    cd(directory)

    try
        filename = SIMULATION_SERIALIZATION_FILENAME
        if !ispath(filename)
            throw(ArgumentError("$filename does not exist"))
        end

        obj = Serialization.deserialize(filename)
        if !(obj isa SimulationSerializationWrapper)
            throw(
                IS.DataFormatError("deserialized object has incorrect type $(typeof(obj))"),
            )
        end

        models = Vector{DecisionModel{<:DecisionProblem}}()
        for name in obj.models
            model =
                deserialize_problem(DecisionProblem, joinpath("problems", "$(name).bin"))
            if !haskey(problem_info[key], "optimizer")
                throw(ArgumentError("problem_info must define 'optimizer'"))
            end
            push!(
                models,
                wrapper.problem_type(
                    name,
                    wrapper.template,
                    sys,
                    restore_from_copy(
                        wrapper.settings;
                        optimizer=problem_info[key]["optimizer"],
                    ),
                    get(problem_info[key], "jump_model", nothing),
                ),
            )
        end

        sim = Simulation(;
            name=obj.name,
            steps=obj.steps,
            models=SimulationModels(problems...),
            problems_sequence=obj.sequence,
            simulation_folder=obj.simulation_folder,
        )
        return sim
    finally
        cd(orig)
    end
end

function serialize_status(sim::Simulation)
    data = Dict("run_status" => string(get_simulation_status(sim)))
    filename = joinpath(get_results_dir(sim), "status.json")
    open(filename, "w") do io
        JSON3.write(io, data)
    end

    return
end

function deserialize_status(sim::Simulation)
    return deserialize_status(get_results_dir(sim))
end

function deserialize_status(results_path::AbstractString)
    filename = joinpath(results_path, "status.json")
    if !isfile(filename)
        error("run status file $filename does not exist")
    end

    data = open(filename, "r") do io
        JSON3.read(io, Dict)
    end

    return get_enum_value(RunStatus, data["run_status"])
end


text:
function is_hermitian(M; atol = ATOL)
	isapprox(M, M', atol = atol)
end

function is_positive_semidefinite(M; atol = ATOL)
	!any(eigvals(M) .< - atol)
end

function is_a_gram_matrix(M; atol = ATOL)
	is_hermitian(M, atol = atol) && is_positive_semidefinite(M, atol = atol) && all([isapprox(M[i,i], one(eltype(M)), atol = atol) for i in 1:size(M)[1]])
end

function is_unitary(U; atol = ATOL)
    isapprox(U' * U, Matrix(I, size(U)), atol = atol)
end

function is_orthonormal(U; atol = ATOL)
	for i = 1:size(U)[2]
		isapprox(abs(norm(U[:, i])), 1., atol = atol) ? nothing : return false
		for j = 1:size(U)[2]
			if j < i
				isapprox(abs(dot(U[:, i], U[:, j])), 0., atol = atol) ? nothing : return false
			end
		end
	end
end

function is_column_normalized(M; atol = ATOL)

	for i = 1:size(M)[2]
		@test  abs(norm(M[:, i])) ≈ 1. atol = atol
	end
end


text:
# Affine coupling layer from Dinh et al. (2017)
# Includes 1x1 convolution from in Putzky and Welling (2019)
# Author: Philipp Witte, pwitte3@gatech.edu
# Date: January 2020

export CouplingLayerBasic, CouplingLayerBasic3D

"""
    CL = CouplingLayerBasic(RB::ResidualBlock; logdet=false)

or

    CL = CouplingLayerBasic(n_in, n_hidden; k1=3, k2=3, p1=1, p2=1, s1=1, s2=1, logdet=false, ndims=2) (2D)

    CL = CouplingLayerBasic(n_in, n_hidden; k1=3, k2=3, p1=1, p2=1, s1=1, s2=1, logdet=false, ndims=3) (3D)

    CL = CouplingLayerBasic3D(n_in, n_hidden; k1=3, k2=3, p1=1, p2=1, s1=1, s2=1, logdet=false) (3D)

 Create a Real NVP-style invertible coupling layer with a residual block.

 *Input*:

 - `RB::ResidualBlock`: residual block layer consisting of 3 convolutional layers with ReLU activations.

 - `logdet`: bool to indicate whether to compte the logdet of the layer

 or

 - `n_in`, `n_hidden`: number of input and hidden channels

 - `k1`, `k2`: kernel size of convolutions in residual block. `k1` is the kernel of the first and third
    operator, `k2` is the kernel size of the second operator.

 - `p1`, `p2`: padding for the first and third convolution (`p1`) and the second convolution (`p2`)

 - `s1`, `s2`: stride for the first and third convolution (`s1`) and the second convolution (`s1`)

 - `ndims` : Number of dimensions

 *Output*:

 - `CL`: Invertible Real NVP coupling layer.

 *Usage:*

 - Forward mode: `Y1, Y2, logdet = CL.forward(X1, X2)`    (if constructed with `logdet=true`)

 - Inverse mode: `X1, X2 = CL.inverse(Y1, Y2)`

 - Backward mode: `ΔX1, ΔX2, X1, X2 = CL.backward(ΔY1, ΔY2, Y1, Y2)`

 *Trainable parameters:*

 - None in `CL` itself

 - Trainable parameters in residual block `CL.RB`

 See also: [`ResidualBlock`](@ref), [`get_params`](@ref), [`clear_grad!`](@ref)
"""
mutable struct CouplingLayerBasic <: NeuralNetLayer
    RB::Union{ResidualBlock, FluxBlock}
    logdet::Bool
    activation::ActivationFunction
    is_reversed::Bool
end

@Flux.functor CouplingLayerBasic

# Constructor from 1x1 convolution and residual block
function CouplingLayerBasic(RB::ResidualBlock; logdet=false, activation::ActivationFunction=SigmoidLayer())
    RB.fan == false && throw("Set ResidualBlock.fan == true")
    return CouplingLayerBasic(RB, logdet, activation, false)
end

CouplingLayerBasic(RB::FluxBlock; logdet=false, activation::ActivationFunction=SigmoidLayer()) = CouplingLayerBasic(RB, logdet, activation, false)

# 2D Constructor from input dimensions
function CouplingLayerBasic(n_in::Int64, n_hidden::Int64; k1=3, k2=3, p1=1, p2=1, s1=1, s2=1, logdet=false, activation::ActivationFunction=SigmoidLayer(), ndims=2)

    # 1x1 Convolution and residual block for invertible layer
    RB = ResidualBlock(n_in, n_hidden; k1=k1, k2=k2, p1=p1, p2=p2, s1=s1, s2=s2, fan=true, ndims=ndims)

    return CouplingLayerBasic(RB, logdet, activation, false)
end

CouplingLayerBasic3D(args...;kw...) = CouplingLayerBasic(args...; kw..., ndims=3)

# 2D/3D Forward pass: Input X, Output Y
function forward(X1::AbstractArray{Float32, N}, X2::AbstractArray{Float32, N}, L::CouplingLayerBasic; save::Bool=false, logdet=nothing) where N
    isnothing(logdet) ? logdet = (L.logdet && ~L.is_reversed) : logdet = logdet

    # Coupling layer
    logS_T1, logS_T2 = tensor_split(L.RB.forward(X1))
    S = L.activation.forward(logS_T1)
    Y2 = S.*X2 + logS_T2

    if logdet
        save ? (return X1, Y2, coupling_logdet_forward(S), S) : (return X1, Y2, coupling_logdet_forward(S))
    else
        save ? (return X1, Y2, S) : (return X1, Y2)
    end
end

# 2D/3D Inverse pass: Input Y, Output X
function inverse(Y1::AbstractArray{Float32, N}, Y2::AbstractArray{Float32, N}, L::CouplingLayerBasic; save::Bool=false, logdet=nothing) where N
    isnothing(logdet) ? logdet = (L.logdet && L.is_reversed) : logdet = logdet

    # Inverse layer
    logS_T1, logS_T2 = tensor_split(L.RB.forward(Y1))
    S = L.activation.forward(logS_T1)
    X2 = (Y2 - logS_T2) ./ (S .+ eps(1f0)) # add epsilon to avoid division by 0

    if logdet
        save == true ? (return Y1, X2, -coupling_logdet_forward(S), S) : (return Y1, X2, -coupling_logdet_forward(S))
    else
        save == true ? (return Y1, X2, S) : (return Y1, X2)
    end
end

# 2D/3D Backward pass: Input (ΔY, Y), Output (ΔX, X)
function backward(ΔY1, ΔY2, Y1, Y2, L::CouplingLayerBasic; set_grad::Bool=true)

    # Recompute forward state
    X1, X2, S = inverse(Y1, Y2, L; save=true, logdet=false)

    # Backpropagate residual
    ΔT = copy(ΔY2)
    ΔS = ΔY2 .* X2
    if L.logdet
        set_grad && (ΔS -= coupling_logdet_backward(S))
    end
    ΔX2 = ΔY2 .* S
    if set_grad
        ΔX1 = L.RB.backward(tensor_cat(L.activation.backward(ΔS, S), ΔT), X1) + ΔY1
    else
        ΔX1, Δθ = L.RB.backward(tensor_cat(L.activation.backward(ΔS, S), ΔT), X1; set_grad=set_grad)
        if L.logdet
            _, ∇logdet = L.RB.backward(tensor_cat(L.activation.backward(coupling_logdet_backward(S), S), 0f0.*ΔT), X1; set_grad=set_grad)
        end
        ΔX1 += ΔY1
    end

    if set_grad
        return ΔX1, ΔX2, X1, X2
    else
        L.logdet ? (return ΔX1, ΔX2, Δθ, X1, X2, ∇logdet) : (return ΔX1, ΔX2, Δθ, X1, X2)
    end
end

# 2D/3D Reverse backward pass: Input (ΔX, X), Output (ΔY, Y)
function backward_inv(ΔX1, ΔX2, X1, X2, L::CouplingLayerBasic; set_grad::Bool=true)

    # Recompute inverse state
    Y1, Y2, S = forward(X1, X2, L; save=true, logdet=false)

    # Backpropagate residual
    ΔT = -ΔX2 ./ S
    ΔS = X2 .* ΔT
    if L.logdet == true
        set_grad ? (ΔS += coupling_logdet_backward(S)) : (∇logdet = -coupling_logdet_backward(S))
    end
    if set_grad
        ΔY1 = L.RB.backward(tensor_cat(L.activation.backward(ΔS, S), ΔT), Y1) + ΔX1
    else
        ΔY1, Δθ = L.RB.backward(tensor_cat(L.activation.backward(ΔS, S), ΔT), Y1; set_grad=set_grad)
        ΔY1 += ΔX1
    end
    ΔY2 = - ΔT

    if set_grad
        return ΔY1, ΔY2, Y1, Y2
    else
        L.logdet ? (return ΔY1, ΔY2, Δθ, Y1, Y2, ∇logdet) : (return ΔY1, ΔY2, Δθ, Y1, Y2)
    end
end

## Jacobian-related functions

# 2D
function jacobian(ΔX1::AbstractArray{Float32, N}, ΔX2::AbstractArray{Float32, N}, Δθ::AbstractArray{Parameter, 1},
                  X1::AbstractArray{Float32, N}, X2::AbstractArray{Float32, N}, L::CouplingLayerBasic;
                  save=false, logdet=nothing) where N
    isnothing(logdet) ? logdet = (L.logdet && ~L.is_reversed) : logdet = logdet

    logS_T1, logS_T2 = tensor_split(L.RB.forward(X1))
    ΔlogS_T1, ΔlogS_T2 = tensor_split(jacobian(ΔX1, Δθ, X1, L.RB)[1])
    S = L.activation.forward(logS_T1)
    ΔS = L.activation.backward(ΔlogS_T1, S)
    Y2 = S.*X2 + logS_T2
    ΔY2 = ΔS.*X2 + S.*ΔX2 + ΔlogS_T2

    if logdet
        # Gauss-Newton approximation of logdet terms
        JΔθ = tensor_split(L.RB.jacobian(zeros(Float32, size(ΔX1)), Δθ, X1)[1])[1]
        GNΔθ = -L.RB.adjointJacobian(tensor_cat(L.activation.backward(JΔθ, S), zeros(Float32, size(S))), X1)[2]

        save ? (return ΔX1, ΔY2, X1, Y2, coupling_logdet_forward(S), GNΔθ, S) : (return ΔX1, ΔY2, X1, Y2, coupling_logdet_forward(S), GNΔθ)
    else
        save ? (return ΔX1, ΔY2, X1, Y2, S) : (return ΔX1, ΔY2, X1, Y2)
    end
end

# 2D/3D
function adjointJacobian(ΔY1, ΔY2, Y1, Y2, L::CouplingLayerBasic)
    return backward(ΔY1, ΔY2, Y1, Y2, L; set_grad=false)
end

## Logdet utils
coupling_logdet_forward(S) = sum(log.(abs.(S))) / size(S)[end]
coupling_logdet_backward(S) = 1f0./ S / size(S)[end]

## Other utils
# Clear gradients
clear_grad!(L::CouplingLayerBasic) = clear_grad!(L.RB)

# Get parameters
get_params(L::CouplingLayerBasic) = get_params(L.RB)

# Set is_reversed flag
function tag_as_reversed!(L::CouplingLayerBasic, tag::Bool)
    L.is_reversed = tag
    return L
end


text:
function find_direction(B, P, ∇f, scheme::QuasiNewton{<:Direct})
    return -(B \ ∇f)
end
function find_direction(B, P, ∇f, scheme::Newton{<:Direct})
    return scheme.linsolve(B, -∇f)
end
function find_direction(A, P, ∇f, scheme::QuasiNewton{<:Inverse})
    return -A * ∇f
end
function find_direction!(d, B, P, ∇f, scheme::QuasiNewton{<:Direct})
    d .= -(B \ ∇f)
    d
end
function find_direction!(d, B, P, ∇f, scheme::Newton{<:Direct})
    scheme.linsolve(d, B, -∇f)
end
function find_direction!(d, A, P, ∇f, scheme::QuasiNewton{<:Inverse})
    rmul!(mul!(d, A, ∇f), -1)
    d
end
function find_direction(B, P, ∇f, scheme::GradientDescent)
    -apply_preconditioner(OutOfPlace(), P, nothing, ∇f)
end
function find_direction!(d, B, P, ∇f, scheme::GradientDescent)
    d = apply_preconditioner(InPlace(), P, d, ∇f)
    d .= .-d
end

max_stars_repo_name: JuliaNLSolvers/NLSolvers

max_issues_repo_name: JuliaNLSolvers/NLSolvers

max_forks_repo_name: JuliaNLSolvers/NLSolvers


text:
################################################################################
### Measurement constructors
################################################################################

function checkflavors(model, N=2)
    if nflavors(model) != N
        @warn(
            "$N flavors are required, but $(nflavors(model)) have been found"
        )
    end
    nothing
end

# This has lattice_iteratorator = Nothing, because it straight up copies G
function greens_measurement(
        mc::DQMC, model::Model, greens_iterator = Greens();
        capacity = _default_capacity(mc), eltype = geltype(mc),
        obs = let
            N = length(lattice(model)) * nflavors(model)
            LogBinner(zeros(eltype, (N, N)), capacity=capacity)
        end, kwargs...
    )
    Measurement(
        mc, model, greens_iterator, Nothing, greens_kernel,
        obs = obs; kwargs...
    )
end

function occupation(mc::DQMC, model::Model; wrapper = nothing, kwargs...)
    li = wrapper === nothing ? EachSiteAndFlavor : wrapper{EachSiteAndFlavor}
    Measurement(mc, model, Greens(), li, occupation_kernel; kwargs...)
end

function charge_density(
        mc::DQMC, model::Model, greens_iterator;
        wrapper = nothing, lattice_iterator = EachSitePairByDistance, kwargs...
    )
    checkflavors(model)
    li = wrapper === nothing ? lattice_iterator : wrapper{lattice_iterator}
    Measurement(mc, model, greens_iterator, li, cdc_kernel; kwargs...)
end

charge_density_correlation(mc, m; kwargs...) = charge_density(mc, m, Greens(); kwargs...)
charge_density_susceptibility(mc, m; kwargs...) = charge_density(mc, m, CombinedGreensIterator(mc); kwargs...)

"""
    magnetization(mc, model, dir[; lattice_iterator, kwargs...])

Returns the x, y or z magnetization measurement given `dir = :x`, `:y` or `:z`
respectively.

NOTE:
We're skipping the multiplication by `-1im` during the measurement of the y
magnetization. To get the correct result, multiply the final result by `-1im`.
"""
function magnetization(
        mc::DQMC, model::Model, dir::Symbol;
        wrapper = nothing, lattice_iterator = EachSite, kwargs...
    )
    checkflavors(model)
    li = wrapper === nothing ? lattice_iterator : wrapper{lattice_iterator}
    if dir == :x;
        return Measurement(mc, model, Greens(), li, mx_kernel; kwargs...)
    elseif dir == :y;
        return Measurement(mc, model, Greens(), li, my_kernel; kwargs...)
    elseif dir == :z;
        return Measurement(mc, model, Greens(), li, mz_kernel; kwargs...)
    else throw(ArgumentError("`dir` must be :x, :y or :z, but is $dir"))
    end

end

function spin_density(
        dqmc, model, dir::Symbol, greens_iterator;
        wrapper = nothing, lattice_iterator = EachSitePairByDistance, kwargs...
    )
    checkflavors(model)
    li = wrapper === nothing ? lattice_iterator : wrapper{lattice_iterator}
    dir in (:x, :y, :z) || throw(ArgumentError("`dir` must be :x, :y or :z, but is $dir"))
    if     dir == :x
        return Measurement(dqmc, model, greens_iterator, li, sdc_x_kernel; kwargs...)
    elseif dir == :y
        return Measurement(dqmc, model, greens_iterator, li, sdc_y_kernel; kwargs...)
    elseif dir == :z
        return Measurement(dqmc, model, greens_iterator, li, sdc_z_kernel; kwargs...)
    else
        error("$greens_iterator not recongized")
    end
end
spin_density_correlation(args...; kwargs...) = spin_density(args..., Greens(); kwargs...)
spin_density_susceptibility(mc, args...; kwargs...) = spin_density(mc, args..., CombinedGreensIterator(mc); kwargs...)

function pairing(
        dqmc::DQMC, model::Model, greens_iterator;
        K = 1+length(neighbors(lattice(model), 1)), wrapper = nothing,
        lattice_iterator = EachLocalQuadByDistance{K}, kwargs...
    )
    li = wrapper === nothing ? lattice_iterator : wrapper{lattice_iterator}
    Measurement(dqmc, model, greens_iterator, li, pc_kernel; kwargs...)
end
pairing_correlation(mc, m; kwargs...) = pairing(mc, m, Greens(); kwargs...)
pairing_susceptibility(mc, m; kwargs...) = pairing(mc, m, CombinedGreensIterator(mc); kwargs...)

#=
Λ^L = Λxx(qx->0, qy=0, iω=0)
    = ∑_r ∫[0, ß] dτ ⟨j_x(r, τ) j_x(0, 0)⟩ exp(iqr) exp(iωτ)
j_x(r, τ) = it ∑_σ [c_{r+e_x, σ}^†(τ) c_{r, σ}(τ) - c_{r, σ}^†(τ) c_{r+e_x, σ}(τ)]

# reframe as:
0 -> r, r -> r + Δr, e_x -> r+e_x, r+e_x -> r+e_x+Δr
then use EachLocalQuadByDistance
Well actually no, because we want e_x to only pick one neighbor.

Λ_{trg-src}(Δr) = \sum_src1 cc_kernel(...)
∑_Δr exp(iq⋅Δr) real(Λ_{trg-src}(Δr)) # Why real? this is silly...
where q = Δq_i = k_1, k_2 (or k_1+k_2?)

                       v- (0, qy)
ρs = 1/8 * (real(Λxxq[1,2]) - real(Λxxq[2,1]) + real(Λyyq[1,2]) - real(Λyyq[2,1]))
                                (qx, 0) -^
This should correspond to
ρ = -K_s - Λ^T = Λxx^L - Λxx^T = Λxx((0, qy)) - Λxx((qx, 0))
  ~ real(Λxxq[1,2]) + real(Λyyq[1,2]) - real(Λxxq[2,1]) - real(Λyyq[2,1])
    ^-        sum over NNs?        -^

So I need to compute reciprocal lattice vectors I guess
So what is qx, qy then? I assume those reciprocal lattice vectors? Or literally
just a dx/dy?

Any real-space runtime computation could be EachLocalQuadByDistance, we can do
summing, direction picking and Fourier afterwards...

To be more efficient:
Should implement new lattice iterator that picks trg1, trg2 to be in the same
direction and skips trg == src
EachSyncedNNQuadByDistance{K}?
=#
function current_current_susceptibility(
        dqmc::DQMC, model::Model;
        K = 1+length(neighbors(lattice(model), 1)),
        greens_iterator = CombinedGreensIterator(dqmc), wrapper = nothing,
        lattice_iterator = EachLocalQuadBySyncedDistance{K}, kwargs...
    )
    li = wrapper === nothing ? lattice_iterator : wrapper{lattice_iterator}
    Measurement(dqmc, model, greens_iterator, li, cc_kernel; kwargs...)
end
function superfluid_density(
        dqmc::DQMC, model::Model, Ls = size(lattice(model));
        K = 1+length(neighbors(lattice(model), 1)),
        capacity = _default_capacity(dqmc),
        obs = LogBinner(ComplexF64(0), capacity=capacity),
        kwargs...
    )
    @assert K > 1
    dirs = directions(lattice(model))

    # Note: this only works with 2d vecs
    # Note: longs and trans are epsilon-vectors, i.e. they're the smallest
    # step we can take in discrete reciprocal space
    lvecs = lattice_vectors(lattice(model))
    uc_vecs = lvecs ./ Ls
    prefactor = 2pi / dot(lvecs...)
    rvecs = map(lv -> prefactor * cross([lv; 0], [0, 0, 1])[[1, 2]], uc_vecs)

    # Note: this only works for 2d lattices... maybe
    longs = []
    trans = []
    dir_idxs = []
    for i in 2:K
        if -uc_vecs[1] ≈ dirs[i] || uc_vecs[1] ≈ dirs[i]
            push!(longs, rvecs[1])
            push!(trans, rvecs[2])
            push!(dir_idxs, i)
        elseif -uc_vecs[2] ≈ dirs[i] || uc_vecs[2] ≈ dirs[i]
            push!(longs, rvecs[2])
            push!(trans, rvecs[1])
            push!(dir_idxs, i)
        else
            # We kinda need the i, j from R = i*a_1 + j*a_2 here so we can
            # construct the matching reciprocal vectors here
            @error("Skipping $(dirs[i]) - not a nearest neighbor")
        end
    end

    #=
    longs = normalize.(dirs[2:K]) * 1/L
    trans = map(dirs[2:K]) do v
        n = [normalize(v)..., 0]
        u = cross([0,0,1], n)
        u[1:2] / L
    end
    longs .*= 2pi
    trans .*= 2pi
    =#
    li = SuperfluidDensity{EachLocalQuadBySyncedDistance{K}}(
        dir_idxs, longs, trans
    )
    # Measurement(dqmc, model, CombinedGreensIterator, li, cc_kernel, obs=obs; kwargs...)
    current_current_susceptibility(dqmc, model, lattice_iterator = li)
end

function boson_energy_measurement(dqmc, model; kwargs...)
    Measurement(dqmc, model, Nothing, Nothing, energy_boson; kwargs...)
end

function noninteracting_energy(dqmc, model; kwargs...)
    Measurement(dqmc, model, Greens(), Nothing, nonintE_kernel; kwargs...)
end

# These require the model to implement intE_kernel
function interacting_energy(dqmc, model; kwargs...)
    Measurement(dqmc, model, Greens(), Nothing, intE_kernel; kwargs...)
end

function total_energy(dqmc, model; kwargs...)
    Measurement(dqmc, model, Greens(), Nothing, totalE_kernel; kwargs...)
end

################################################################################
### Wicks expanded Kernel functions
################################################################################

# Notes on math and stuff:
#=
- Use Wicks theorem to go from one expectation value with 2N operators to
  N expectation values with 2 operators each
- use G_{ij} = ⟨c_i c_j^†⟩
- use ⟨c_i^† c_j⟩ = δ_{ij} - ⟨c_j c_i^†⟩ = (I - G)_{ji} = dagger(::GreensMatrix)[i, j]
- we define G[i,j] as spin up, G[i+N,j+N] as spin down (N = number of sites)
- we used (I - G)[i, j+N] = G[i, j+N] etc when possible
- Gl0_{i,j} = ⟨c_i(l) c_j(0)^†⟩
- ⟨c_i(0) c_j(l)^†⟩ = ⟨(c_j(l) c_i(0)^†)^†⟩ = (⟨c_j(l) c_i(0)^†⟩)^†
                    = Gl0_{j,i}^† = Gl0_{j,i}^*
=#

"""
    greens_kernel(mc, model, G::GreensMatrix)

Returns the unprocessed Greens function `greens(mc) = {⟨cᵢcⱼ^†⟩}`.

* Lattice Iterators: `nothing` (zero index)
* Greens Iterators: `Greens` or `GreensAt`
"""
greens_kernel(mc, model, G::GreensMatrix) = G.val

"""
    occupation_kernel(mc, model, i::Integer, G::GreensMatrix)

Returns the per index occupation `⟨nᵢ⟩`.

* Lattice Iterators: `EachSiteAndFlavor`, `EachSite`
* Greens Iterators: `Greens` or `GreensAt`
"""
occupation_kernel(mc, model, i::Integer, G::GreensMatrix) = 1 - G[i, i]

"""
    cdc_kernel(mc, model, ij::NTuple{2, Integer}, G::GreensMatrix)
    cdc_kernel(mc, model, ij::NTuple{2, Integer}, Gs::NTuple{4, GreensMatrix})

Returns the per-site-pair charge density `⟨nᵢ(τ) nⱼ(0)⟩ - ⟨nᵢ(τ)⟩⟨nⱼ(0)⟩` where
`τ = 0` for the first signature.

* Lattice Iterators: `OnSite`, `EachSitePair` or `EachSitePairByDistance`
* Greens Iterators: `Greens`, `GreensAt` or `CombinedGreensIterator`
"""
function cdc_kernel(mc, model, ij::NTuple{2}, G::GreensMatrix)
    i, j = ij
    N = length(lattice(mc))
    # ⟨n↑n↑⟩
    dagger(G)[i, i] * dagger(G)[j, j] +
    dagger(G)[i, j] * G[i, j] +
    # ⟨n↑n↓⟩
    dagger(G)[i, i] * dagger(G)[j+N, j+N] +
    dagger(G)[i, j+N] * G[i, j + N] +
    # ⟨n↓n↑⟩
    dagger(G)[i+N, i+N] * dagger(G)[j, j] +
    dagger(G)[i+N, j] * G[i+N, j] +
    # ⟨n↓n↓⟩
    dagger(G)[i+N, i+N]  * dagger(G)[j+N, j+N] +
    dagger(G)[i+N, j+N] * G[i+N, j+N]
end

function cdc_kernel(mc, model, ij::NTuple{2}, packed_greens::NTuple{4})
    i, j = ij
	G00, G0l, Gl0, Gll = packed_greens
    N = length(lattice(mc))
    # ⟨n↑(l)n↑(0)⟩
    dagger(Gll)[i, i] * dagger(G00)[j, j] +
    dagger(G0l)[i, j] * Gl0[i, j] +
    # ⟨n↑(l)n↓(0)⟩
    dagger(Gll)[i, i] * dagger(G00)[j+N, j+N] +
    dagger(G0l)[i, j+N] * Gl0[i, j + N] +
    # ⟨n↓(l)n↑(0)⟩
    dagger(Gll)[i+N, i+N] * dagger(G00)[j, j] +
    dagger(G0l)[i+N, j] * Gl0[i+N, j] +
    # ⟨n↓(l)n↓(0)⟩
    dagger(Gll)[i+N, i+N] * dagger(G00)[j+N, j+N] +
    dagger(G0l)[i+N, j+N] * Gl0[i+N, j+N]
end

"""
    mx_kernel(mc, model, i::Integer, G::GreensMatrix)

Returns the per-site x-magnetization `⟨cᵢ↑^† cᵢ↓ + cᵢ↓^† cᵢ↑⟩`.

* Lattice Iterators: `EachSite`
* Greens Iterators: `Greens` or `GreensAt`
"""
function mx_kernel(mc, model, i, G::GreensMatrix)
    N = length(lattice(model))
    -G[i+N, i] - G[i, i+N]
end

"""
    my_kernel(mc, model, i::Integer, G::GreensMatrix)

Returns the per-site y-magnetization `-⟨cᵢ↑^† cᵢ↓ - cᵢ↓^† cᵢ↑⟩` without the
imaginary prefactor.

* Lattice Iterators: `EachSite`
* Greens Iterators: `Greens` or `GreensAt`
"""
function my_kernel(mc, model, i, G::GreensMatrix)
    N = length(lattice(model))
    G[i+N, i] - G[i, i+N]
end

"""
    mz_kernel(mc, model, i::Integer, G::GreensMatrix)

Returns the per-site z-magnetization `⟨nᵢ↑ - nᵢ↓⟩`.

* Lattice Iterators: `EachSite`
* Greens Iterators: `Greens` or `GreensAt`
"""
function mz_kernel(mc, model, i, G::GreensMatrix)
    N = length(lattice(model))
    G[i+N, i+N] - G[i, i]
end

"""
    sdc_x_kernel(mc, model, ij::NTuple{2, Integer}, G::GreensMatrix)
    sdc_x_kernel(mc, model, ij::NTuple{2, Integer}, Gs::NTuple{4, GreensMatrix})

Returns the per-site-pair x-spin density `⟨mxᵢ(τ) mxⱼ(0)⟩ - ⟨mxᵢ(τ)⟩⟨mxⱼ(0)⟩`
where `τ = 0` for the first signature.

* Lattice Iterators: `OnSite`, `EachSitePair` or `EachSitePairByDistance`
* Greens Iterators: `Greens`, `GreensAt` or `CombinedGreensIterator`
"""
function sdc_x_kernel(mc, model, ij::NTuple{2}, G::GreensMatrix)
    i, j = ij
    N = length(lattice(model))
    G[i+N, i] * G[j+N, j] + G[i+N, i] * G[j, j+N] +
    G[i, i+N] * G[j+N, j] + G[i, i+N] * G[j, j+N] +
    dagger(G)[i, j+N] * G[i+N, j] + dagger(G)[i, j] * G[i+N, j+N] +
    dagger(G)[i+N, j+N] * G[i, j] + dagger(G)[i+N, j] * G[i, j+N]
end
function sdc_x_kernel(mc, model, ij::NTuple{2}, packed_greens::NTuple{4})
    i, j = ij
	G00, G0l, Gl0, Gll = packed_greens
    N = length(lattice(model))
    Gll[i+N, i] * G00[j+N, j] + Gll[i+N, i] * G00[j, j+N] +
    Gll[i, i+N] * G00[j+N, j] + Gll[i, i+N] * G00[j, j+N] +
    dagger(G0l)[i, j+N] * Gl0[i+N, j] + dagger(G0l)[i, j] * Gl0[i+N, j+N] +
    dagger(G0l)[i+N, j+N] * Gl0[i, j] + dagger(G0l)[i+N, j] * Gl0[i, j+N]
end

"""
    sdc_y_kernel(mc, model, ij::NTuple{2, Integer}, G::GreensMatrix)
    sdc_y_kernel(mc, model, ij::NTuple{2, Integer}, Gs::NTuple{4, GreensMatrix})

Returns the per-site-pair x-spin density `⟨myᵢ(τ) myⱼ(0)⟩ - ⟨myᵢ(τ)⟩⟨myⱼ(0)⟩`
where `τ = 0` for the first signature.

* Lattice Iterators: `OnSite`, `EachSitePair` or `EachSitePairByDistance`
* Greens Iterators: `Greens`, `GreensAt` or `CombinedGreensIterator`
"""
function sdc_y_kernel(mc, model, ij::NTuple{2}, G::GreensMatrix)
    i, j = ij
    N = length(lattice(model))
    - G[i+N, i] * G[j+N, j] + G[i+N, i] * G[j, j+N] +
      G[i, i+N] * G[j+N, j] - G[i, i+N] * G[j, j+N] +
    - dagger(G)[i, j+N] * G[i+N, j] + dagger(G)[i, j] * G[i+N, j+N] +
      dagger(G)[i+N, j+N] * G[i, j] - dagger(G)[i+N, j] * G[i, j+N]
end
function sdc_y_kernel(mc, model, ij::NTuple{2}, packed_greens::NTuple{4})
    i, j = ij
	G00, G0l, Gl0, Gll = packed_greens
    N = length(lattice(model))
    - Gll[i+N, i] * G00[j+N, j] + Gll[i+N, i] * G00[j, j+N] +
      Gll[i, i+N] * G00[j+N, j] - Gll[i, i+N] * G00[j, j+N] +
    - dagger(G0l)[i, j+N] * Gl0[i+N, j] + dagger(G0l)[i, j] * Gl0[i+N, j+N] +
      dagger(G0l)[i+N, j+N] * Gl0[i, j] - dagger(G0l)[i+N, j] * Gl0[i, j+N]
end

"""
    sdc_z_kernel(mc, model, ij::NTuple{2, Integer}, G::GreensMatrix)
    sdc_z_kernel(mc, model, ij::NTuple{2, Integer}, Gs::NTuple{4, GreensMatrix})

Returns the per-site-pair x-spin density `⟨mzᵢ(τ) mzⱼ(0)⟩ - ⟨mzᵢ(τ)⟩⟨mzⱼ(0)⟩`
where `τ = 0` for the first signature.

* Lattice Iterators: `OnSite`, `EachSitePair` or `EachSitePairByDistance`
* Greens Iterators: `Greens`, `GreensAt` or `CombinedGreensIterator`
"""
function sdc_z_kernel(mc, model, ij::NTuple{2}, G::GreensMatrix)
    i, j = ij
    N = length(lattice(model))
    dagger(G)[i, i]     * dagger(G)[j, j] -
    dagger(G)[i, i]     * dagger(G)[j+N, j+N] -
    dagger(G)[i+N, i+N] * dagger(G)[j, j] +
    dagger(G)[i+N, i+N] * dagger(G)[j+N, j+N] +
    dagger(G)[i, j] * G[i, j]     - dagger(G)[i, j+N] * G[i, j+N] -
    dagger(G)[i+N, j] * G[i+N, j] + dagger(G)[i+N, j+N] * G[i+N, j+N]
end
function sdc_z_kernel(mc, model, ij::NTuple{2}, packed_greens::NTuple{4})
    i, j = ij
	G00, G0l, Gl0, Gll = packed_greens
    N = length(lattice(model))
    dagger(Gll)[i, i] * dagger(G00)[j, j] -
    dagger(Gll)[i, i] * dagger(G00)[j+N, j+N] -
    dagger(Gll)[i+N, i+N] * dagger(G00)[j, j] +
    dagger(Gll)[i+N, i+N] * dagger(G00)[j+N, j+N] +
    dagger(G0l)[i, j] * Gl0[i, j]     - dagger(G0l)[i, j+N] * Gl0[i, j+N] -
    dagger(G0l)[i+N, j] * Gl0[i+N, j] + dagger(G0l)[i+N, j+N] * Gl0[i+N, j+N]
end

"""
    pc_kernel(mc, model, ij::NTuple{4, Integer}, G::GreensMatrix)
    pc_kernel(mc, model, ij::NTuple{4, Integer}, Gs::NTuple{4, GreensMatrix})

Returns the per-site-pair pairing `⟨Δᵢⱼ(τ) Δₖₗ^†(0)⟩` where `τ = 0` for the
first signature and `Δᵢⱼ(τ) = 0.5 cᵢ↑(τ) cⱼ↓(τ)`. The Delta operator is a
deferred version of the  pair field operator `Δᵢ = 0.5 ∑ₐ f(a) cᵢ↑ cᵢ₊ₐ↓` which
leaves the execution of the sum for after the simulation.

* Lattice Iterators: `EachLocalQuadByDistance` or `EachLocalQuadBySyncedDistance`
* Greens Iterators: `Greens`, `GreensAt` or `CombinedGreensIterator`
"""
function pc_kernel(mc, model, sites::NTuple{4}, G::GreensMatrix)
    pc_kernel(mc, model, sites, (G, G, G, G))
end
function pc_kernel(mc, model, sites::NTuple{4}, packed_greens::NTuple{4})
    src1, trg1, src2, trg2 = sites
	G00, G0l, Gl0, Gll = packed_greens
    N = length(lattice(model))
    # Δ_v(src1, trg1)(τ) Δ_v^†(src2, trg2)(0)
    # G_{i, j}^{↑, ↑}(τ, 0) G_{i+d, j+d'}^{↓, ↓}(τ, 0) -
    # G_{i, j+d'}^{↑, ↓}(τ, 0) G_{i+d, j}^{↓, ↑}(τ, 0)
    Gl0[src1, src2] * Gl0[trg1+N, trg2+N] - Gl0[src1, trg2+N] * Gl0[trg1+N, src2]
end

"""
    pc_kernel(mc, model, ij::NTuple{4, Integer}, G::GreensMatrix)
    pc_kernel(mc, model, ij::NTuple{4, Integer}, Gs::NTuple{4, GreensMatrix})

Returns the per-site-pair pairing `⟨Δᵢⱼ^†(τ) Δₖₗ(0)⟩` where `τ = 0` for the
first signature and `Δᵢⱼ(τ) = 0.5 cᵢ↑(τ) cⱼ↓(τ)`. The Delta operator is a
deferred version of the  pair field operator `Δᵢ = 0.5 ∑ₐ f(a) cᵢ↑ cᵢ₊ₐ↓` which
leaves the execution of the sum for after the simulation.

* Lattice Iterators: `EachLocalQuadByDistance` or `EachLocalQuadBySyncedDistance`
* Greens Iterators: `Greens`, `GreensAt` or `CombinedGreensIterator`
"""
function pc_alt_kernel(mc, model, sites::NTuple{4}, G::GreensMatrix)
    pc_alt_kernel(mc, model, sites, (G, G, G, G))
end
function pc_alt_kernel(mc, model, sites::NTuple{4}, packed_greens::NTuple{4})
    src1, trg1, src2, trg2 = sites
	G00, G0l, Gl0, Gll = packed_greens
    N = length(lattice(model))
    # Δ_v^†(src1, trg1)(τ) Δ_v(src2, trg2)(0)
    # (I-G)_{j, i}^{↑, ↑}(0, τ) (I-G)_{j+d', i+d}^{↓, ↓}(0, τ) -
    # (I-G)_{j, i+d}^{↑, ↓}(0, τ) G_{j+d', i}^{↓, ↑}(0, τ)
    dagger(G0l)[trg1+N, trg2+N] * dagger(G0l)[src1, src2] -
    G0l[src2, trg1+N] * G0l[trg2+N, src1]
end

function pc_combined_kernel(mc, model, sites::NTuple{4}, G)
    # Δ^† Δ + Δ Δ^†
    pc_kernel(mc, model, sites, G) + pc_alt_kernel(mc, model, sites, G)
end

function pc_ref_kernel(mc, model, sites::NTuple{4}, G::GreensMatrix)
    # Δ^† Δ + Δ Δ^† but ↑ and ↓ are swapped
    pc_ref_kernel(mc, model, sites, (G, G, G, G))
end
function pc_ref_kernel(mc, model, sites::NTuple{4}, packed_greens::NTuple{4})
    src1, trg1, src2, trg2 = sites
	G00, G0l, Gl0, Gll = packed_greens
    N = length(lattice(model))
    Gl0[src1+N, src2+N] * Gl0[trg1, trg2] -
    Gl0[src1+N, trg2] * Gl0[trg1, src2+N] +
    (I[trg2, trg1] - G0l[trg2, trg1]) * (I[src2, src1] - G0l[src2+N, src1+N]) -
    (I[src2, trg1] - G0l[src2+N, trg1]) * (I[trg2, src1] - G0l[trg2, src1+N])
end

function cc_kernel(mc, model, sites::NTuple{4}, packed_greens::NTuple{4})
    # Computes
    # ⟨j_{t2-s2}(s2, l) j_{t1-s1}(s1, 0)⟩
    # where t2-s2 (t1-s1) is usually a NN vector/jump, and
    # j_{t2-s2}(s2, l) = i \sum_σ [T_{ts} c_t^†(l) c_s(τ) - T_{st} c_s^†(τ) c_t(τ)]
    src1, trg1, src2, trg2 = sites
	G00, G0l, Gl0, Gll = packed_greens
    N = length(lattice(model))
    T = mc.stack.hopping_matrix
    output = zero(eltype(G00))

    # Iterate through (spin up, spin down)
    for σ1 in (0, N), σ2 in (0, N)
        s1 = src1 + σ1; t1 = trg1 + σ1
        s2 = src2 + σ2; t2 = trg2 + σ2
        # Note: if H is real and Hermitian, T can be pulled out and the I's cancel
        # Note: This matches crstnbr/dqmc if H real, Hermitian
        # Note: I for G0l and Gl0 auto-cancels
        output -= (
                T[t2, s2] * (I[s2, t2] - Gll[s2, t2]) -
                T[s2, t2] * (I[t2, s2] - Gll[t2, s2])
            ) * (
                T[t1, s1] * (I[s1, t1] - G00[s1, t1]) -
                T[s1, t1] * (I[t1, s1] - G00[t1, s1])
            ) +
            - T[t2, s2] * T[t1, s1] * G0l[s1, t2] * Gl0[s2, t1] +
            + T[t2, s2] * T[s1, t1] * G0l[t1, t2] * Gl0[s2, s1] +
            + T[s2, t2] * T[t1, s1] * G0l[s1, s2] * Gl0[t2, t1] +
            - T[s2, t2] * T[s1, t1] * G0l[t1, s2] * Gl0[t2, s1]
    end

    # OLD PARTIALLY OUTDATED

    # This should compute
    # ⟨j_{trg1-src1}(src1, τ) j_{trg2-src2}(src2, 0)⟩
    # where (trg-src) picks a direction (e.g. NN directions)
    # and (src1-src2) is the distance vector that the Fourier transform applies to
    # From dos Santos: Introduction to Quantum Monte Carlo Simulations
    # j_{trg-src}(src, τ) = it \sum\sigma (c^\dagger(trg,\sigma, \tau) c(src, \sigma, \tau) - c^\dagger(src, \sigma, \tau) c(trg, \sigma \tau))
    # where i -> src, i+x -> trg as a generalization
    # and t is assumed to be hopping matrix element, generalizing to
    # = i \sum\sigma (T[trg, src] c^\dagger(trg,\sigma, \tau) c(src, \sigma, \tau) - T[src, trg] c^\dagger(src, \sigma, \tau) c(trg, \sigma \tau))

            # Why no I? - delta_0l = 0
            # T[t1, s1] * T[t2, s2] * (I[s2, t1] - G0l[s2, t1]) * Gl0[s1, t2] -
            # T[s1, t1] * T[t2, s2] * (I[s2, s1] - G0l[s2, s1]) * Gl0[t1, t2] -
            # T[t1, s1] * T[s2, t2] * (I[t2, t1] - G0l[t2, t1]) * Gl0[s1, s2] +
            # T[s1, t1] * T[s2, t2] * (I[t2, s1] - G0l[t2, s1]) * Gl0[t1, s2]

        # Uncompressed Wicks expansion
        # output += T[t1, s1] * T[t2, s2] *
        #     ((I[s1, t1] - Gll[s1, t1]) * (I[s2, t2] - G00[s2, t2]) +
        #     (I[s2, t1] - G0l[s2, t1]) * Gl0[s1, t2])
        # output -= T[s1, t1] * T[t2, s2] *
        #     ((I[t1, s1] - Gll[t1, s1]) * (I[s2, t2] - G00[s2, t2]) +
        #     (I[s2, s1] - G0l[s2, s1]) * Gl0[t1, t2])
        # output -= T[t1, s1] * T[s2, t2] *
        #     ((I[s1, t1] - Gll[s1, t1]) * (I[t2, s2] - G00[t2, s2]) +
        #     (I[t2, t1] - G0l[t2, t1]) * Gl0[s1, s2])
        # output += T[s1, t1] * T[s2, t2] *
        #     ((I[t1, s1] - Gll[t1, s1]) * (I[t2, s2] - G00[t2, s2]) +
        #     (I[t2, s1] - G0l[t2, s1]) * Gl0[t1, s2])
    output
end

@inline function nonintE_kernel(mc, model, G::GreensMatrix)
    # <T> = \sum Tji * (Iij - Gij) = - \sum Tji * (Gij - Iij)
    T = mc.stack.hopping_matrix
    nonintE(T, G.val)
end

# TODO should this be moved/work differently?
nonintE(T::AbstractArray, G::GreensMatrix) = nonintE(T, G.val)
function nonintE(T::AbstractArray, G::AbstractArray)
    output = zero(eltype(G))
    for i in axes(G, 1), j in axes(G, 2)
        output += T[j, i] * (I[i, j] - G[i, j])
    end
    # 2 because we're using spin up/down symmetry
    2.0 * output
end
function nonintE(T::BlockDiagonal{X, N}, G::BlockDiagonal{X, N}) where {X, N}
    output = zero(eltype(G))
    @inbounds n = size(T.blocks[1], 1)
    @inbounds for i in 1:N
        t = T.blocks[i]
        g = G.blocks[i]
        @turbo for k in 1:n, l in 1:n
            output += t[k,l] * (ifelse(k==l, 1.0, 0.0) - g[k,l])
        end
    end
    output
end

function totalE_kernel(mc, model, G::GreensMatrix)
    nonintE_kernel(mc, model, G) + intE_kernel(mc, model, G)
end


text:
f(z::Number)=z^2
g(x) = x+1

max_stars_repo_name: rlerichev/ProjectGitExercise

max_issues_repo_name: rlerichev/ProjectGitExercise

max_forks_repo_name: rlerichev/ProjectGitExercise