import FITSIO
import JLD
import Optim

import ..Log
using ..Model
import ..SDSSIO
import ..Infer
import ..SDSSIO: RunCamcolField
import ..PSF

using ..DeterministicVI

function union_find!(i, components_tree)
    root = i
    while components_tree[i] != i
        i = components_tree[i]
    end

    while root != i
        root2 = components_tree[root]
        components_tree[root] = i
        root = root2
    end
    return i
end

"""
Computes connected components given a conflict graph, and the range of sources
to process. Writes to the components argument.

- source_start - start source
- end_source - end source (inclusively processed)
- sources - the actual source ids to find connected components for (indexed by source_start and source_end)
- neighbor_map - conflict graph of sources
- components_tree - the components array to do union find on.
- components - Dict{Int64, Vector{Int64}} dictionary from cc_id to target_source index in that component
- src_to_target_index - a mapping from light source id -> index within the target_sources array. This is required
                        since we need to write the index within target sources to the output, rather than the source id itself.

Returns:
- Nothing
"""
function compute_connected_components(source_start, source_end, sources, neighbor_map,
                                      components_tree, components, src_to_target_index)
    # Construct reverse map from source_id -> index in source. Only do this for
    # the specified portion we are computing connected components.
    source_to_index = Dict{Int64, Int64}()
    index_to_source = Dict{Int64, Int64}()
    for i = source_start:source_end
        source_to_index[sources[i]] = i-source_start+1
        index_to_source[i-source_start+1] = sources[i]
    end

    for i = source_start:source_end
        components_tree[i-source_start+1] = i-source_start+1
    end

    # Run union find algorithm to find connected components.
    for i = source_start:source_end
        target = union_find!(i-source_start+1, components_tree)
        for neighbor in neighbor_map[index_to_source[i-source_start+1]]
            if haskey(source_to_index, neighbor)
                neighbor_idx = source_to_index[neighbor]
                component_source = union_find!(neighbor_idx, components_tree)
                components_tree[component_source] = target
            end

        end
    end

    # Write to components dictionary.
    for i = source_start:source_end
        index = union_find!(i-source_start+1, components_tree)
        if !haskey(components, index)
            components[index] = Vector{Int64}()
        end
        push!(components[index], src_to_target_index[index_to_source[i-source_start+1]])
    end
end

"""
Partitions sources via the cyclades algorithm.

- nprocthreads - number of threads to which to distribute sources
- target_sources - array of target sources. Elements should match keys of neighbor_map.
- neighbor_map - graph of connections of sources

Returns:
- An array of vectors representing the workload of each thread ([thread][batch][sources(indices)])
"""
function partition_cyclades(nprocthreads, target_sources, neighbor_map; batch_size=60)
    Log.info("Starting Cyclades partitioning...")
    tic()

    n_sources = length(target_sources)
    n_total_batches = convert(Int64, ceil(n_sources / batch_size))

    # Construct src_to_indx map
    src_to_indx = Dict{Int64, Int64}()
    for i=1:n_sources
        src_to_indx[target_sources[i]] = i
    end

    # The final workload distribution.
    thread_sources_assignment = Array(Vector{Vector{Int64}}, nprocthreads)
    for thread = 1:nprocthreads
        thread_sources_assignment[thread] = Vector{Vector{Int64}}(n_total_batches)
        for batch = 1:n_total_batches
            thread_sources_assignment[thread][batch] = Vector{Int64}()
        end
    end

    # First shuffle the sources. Note Cyclades is serially equivalent
    # to this permutation of sources.
    sources = [x for x in keys(neighbor_map)]
    shuffle!(sources)

    # Process batch_size number of sources at a time.
    # TODO max - parallelize everything below (particularly CC computation)

    # The components tree is for union-find.
    components_tree = [Array(Int64, batch_size) for i=1:nprocthreads]

    # We have n_total_batches components, where each component is a dictionary
    # from the component id, to a list of sources in that component
    components = [Dict{Int64, Vector{Int64}}() for i=1:n_total_batches]
    sources_in_components = 0
    for (cur_batch, source_index) in enumerate(collect(1:batch_size:n_sources))
        source_start = source_index
        source_end = min(source_index+batch_size-1, n_sources)
        @assert source_end - source_start + 1 <= batch_size

        # TODO max - parallelize this
        compute_connected_components(source_start, source_end, sources, neighbor_map,
                                     components_tree[1],
                                     components[cur_batch],
                                     src_to_indx)
    end

    assigned_sources = 0
    # Load balance the connected components within each batch into thread_sources_assignment.
    for (cur_batch, cur_batch_component) in enumerate(components)
        # Priority queue for load balancing.
        pqueue = Base.Collections.PriorityQueue([i for i=1:nprocthreads],
                                                [0 for i=1:nprocthreads],
                                                Base.Order.Forward)

        # Assign non-conflicting group of sources to different threads
        for (component_group_id, sources_of_component) in cur_batch_component
            least_loaded_thread = Base.Collections.peek(pqueue)[1]
            for source_to_assign in sources_of_component
                push!(thread_sources_assignment[least_loaded_thread][cur_batch], source_to_assign)
                assigned_sources += 1
            end
            pqueue[least_loaded_thread] += length(sources_of_component)
        end
    end

    Log.info("Cyclades - Assigned sources: $(assigned_sources) vs correct number of sources: $(n_sources)")
    @assert assigned_sources == n_sources
    Log.info("Cyclades - Number of batches: $(n_total_batches)")
    Log.info("Finished Cyclades partitioning.  Elapsed time: $(toq()) seconds")
    thread_sources_assignment
end

"""
Partitions sources across threads equally.

- nprocthreads - number of threads to which to distribute processing the sources.
- n_sources - the number of total sources to process.

Returns:
- An array of vectors representing the workload of each thread ([thread][batch][sources]).
  Note for partition equally, there is only 1 batch.

"""
function partition_equally(nprocthreads, n_sources)
    Log.info("Starting basic source partitioning...")
    tic()
    n_sources_per_thread = floor(Int64, n_sources / nprocthreads)
    thread_sources_assignment = Array(Vector{Vector{Int64}}, nprocthreads)
    n_sources_assigned = 0
    for thread = 1:nprocthreads
        start_source = (thread-1) * n_sources_per_thread
        end_source = thread * n_sources_per_thread
        if thread == nprocthreads
            end_source = n_sources
        end
        # Only 1 batch for partitioning equally
        thread_sources_assignment[thread] = Vector{Vector{Int64}}(1)
        thread_sources_assignment[thread][1] = Vector{Int64}(end_source-start_source)
        for source_assignment = start_source:end_source-1
            thread_sources_assignment[thread][1][source_assignment-start_source+1] = source_assignment+1
            n_sources_assigned += 1
        end
    end
    @assert n_sources_assigned == n_sources
    Log.info("Finished basic source partitioning. Elapsed time: $(toq()) seconds")
    thread_sources_assignment
end

"""
Like one_node_infer, uses multiple threads on one node to fit the Celeste
model over numerous iterations.

catalog - the catalog of light sources
target_sources - light sources to optimize
neighbor_map - ligh_source index -> neighbor light_source id

cyclades_partition - use the cyclades algorithm to partition into non conflicting batches for updates.
batch_size - size of a single batch of sources for updates
within_batch_shuffling - whether or not to process sources within a batch randomly
joint_inference_terminate - whether to terminate once sources seem to be stable
termination_percent - stop optimization once a certain percentage of sources have been optimized.

Returns:

- Vector of OptimizedSource results
"""
function one_node_joint_infer(catalog, target_sources, neighbor_map, images;
                              cyclades_partition=true,
                              batch_size=60,
                              within_batch_shuffling=true,
                              termination_percent=.95,
                              n_iters=10)
    # Seed random number generator to ensure the same results per run.
    srand(42)

    nprocthreads = nthreads()

    # Partition the sources
    n_sources = length(target_sources)
    Log.info("Optimizing $(n_sources) sources")
    if cyclades_partition
        # Convert neighbormap to cyclades map (map from source_id -> [neighbor source_ids])
        cyclades_neighbor_map = Dict{Int64, Vector{Int64}}()
        for (index, neighbors) in enumerate(neighbor_map)
            source_id = target_sources[index]
            cyclades_neighbor_map[source_id] = neighbors
        end
        thread_sources_assignment = partition_cyclades(nprocthreads, target_sources, cyclades_neighbor_map, batch_size=batch_size)
    else
        thread_sources_assignment = partition_equally(nprocthreads, n_sources)
    end

    Log.info("Done assigning sources to threads for processing")

    # Pre allocate elbo args variational params
    target_source_variational_params = Dict{Int64, Array{Float64}}()
    for target_source in target_sources
        cat = catalog[target_source]
        target_source_variational_params[target_source] = init_source(cat)
    end

    # Pre-allocate dictionary of elboargs, call it ea_vec.
    ea_vec = Array{ElboArgs}(n_sources)
    function initialize_elboargs_sources(sources)
#        nputs(dt_nodeid, "Thread $(Threads.threadid()) allocating mem for $(length(sources)) sources")
        for cur_source_index in sources
            entry_id = target_sources[cur_source_index]
            entry = catalog[target_sources[cur_source_index]]
            neighbor_ids = neighbor_map[cur_source_index]
            neighbors = catalog[neighbor_map[cur_source_index]]

            # TODO max: refactor this portion? It's reused in infer_source.
#            nputs(dt_nodeid, "Thread $(Threads.threadid()) allocating mem for source $(target_sources[cur_source_index]): objid=$(entry.objid)")
            cat_local = vcat([entry], neighbors)
            ids_local = vcat([entry_id], neighbor_ids)

            #vp = Vector{Float64}[init_source(ce) for ce in cat_local]
            vp = Vector{Float64}[haskey(target_source_variational_params, x) ?
                        target_source_variational_params[x] :
                        init_source(catalog[x]) for x in ids_local]
            patches = Infer.get_sky_patches(images, cat_local)
            ea = ElboArgs(images, vp, patches, [1])
            Infer.load_active_pixels!(ea.images, ea.patches)
            ea_vec[cur_source_index] = ea
        end
    end

    # Initialize elboargs in parallel
    tic()
    thread_initialize_sources_assignment = partition_equally(nprocthreads, n_sources)

    Threads.@threads for i=1:nprocthreads
        for batch = 1:length(thread_initialize_sources_assignment[i])
            initialize_elboargs_sources(thread_initialize_sources_assignment[i][batch])
        end
    end

    Log.info("Done preallocating array of elboargs. Elapsed time: $(toq())")

    # Keep track of which sources have converged.
    sources_converged = Dict{Int64, Bool}()
    for source in target_sources
        sources_converged[source] = false
    end
    n_sources_converged = 0
    n_sources_converged_lock = SpinLock()

    function should_optimize_source(src_indx)
        src_has_converged = sources_converged[target_sources[src_indx]]
        neighbors_have_converged = true
        for neighbor in neighbor_map[src_indx]
            if haskey(sources_converged, neighbor)
                neighbors_have_converged = neighbors_have_converged && sources_converged[neighbor]
            end
        end
        return !src_has_converged || !neighbors_have_converged
    end

    # Process partition of sources. Multiple threads call this function in parallel.
    function process_sources(source_assignment::Vector{Int64}, iter)
        try

            # Use a constant number of newton steps
            n_newton_steps = 5

            # Shuffle the source assignments within each batch of each process.
            # This is disabled by default because it ruins the deterministic outcome
            # required by the test cases.
            if within_batch_shuffling
                shuffle!(source_assignment)
            end
            for cur_source_indx in source_assignment
                # Optimize only if source has not converged or at least
                # one of its neighbors has not converged
                if should_optimize_source(cur_source_indx)
                    cur_entry = catalog[target_sources[cur_source_indx]]
                    iter_count, obj_value, max_x, r = DeterministicVI.maximize_f(
                        DeterministicVI.elbo,
                        ea_vec[cur_source_indx],
                        max_iters=n_newton_steps,
                        use_default_optim_params=true)
                    sources_converged[target_sources[cur_source_indx]] = Optim.converged(r)
                end

                # Maintain count of sources that have converged
                if sources_converged[target_sources[cur_source_indx]]
                    lock(n_sources_converged_lock)
                    n_sources_converged += 1
                    unlock(n_sources_converged_lock)
                end
            end
        catch ex
            if is_production_run || nthreads() > 1
                Log.error(string(ex))
            else
                rethrow(ex)
            end
        end
    end

    # Process sources in parallel using nprocthreads.
    tic()
    n_batches = length(thread_sources_assignment[1])
    for iter = 1:n_iters
        # Reset number of sources converged
        n_sources_converged = 0

        # Process every batch of every iteration. We do the batches on the outside
        # Since there is an implicit barrier after the inner threaded for loop below.
        # We want this barrier because there may be conflict _between_ Cyclades batches.
        for batch = 1:n_batches
            # Process every batch of every iteration with nprocthreads
            Threads.@threads for i = 1:nprocthreads
                process_sources(thread_sources_assignment[i][batch], iter)
            end
        end

        if n_sources_converged >= termination_percent * n_sources
            break
        end
    end
    Log.info("$(n_sources_converged) / $(n_sources) converged")
    Log.info("Done fitting elboargs. Elapsed time: $(toq())")

    # Return add results to vector
    results = OptimizedSource[]

    for i = 1:n_sources
        entry = catalog[target_sources[i]]
        result = OptimizedSource(entry.thing_id,
                                 entry.objid,
                                 entry.pos[1],
                                 entry.pos[2],
                                 ea_vec[i].vp[1])
        push!(results, result)
    end

    results
end


@Base.deprecate predict(mm::LinearModel, newx::AbstractMatrix, interval::Symbol, level::Real = 0.95) predict(mm, newx; interval=interval, level=level)


@testset "construct by periods" begin
    @test NanoDate(yr,mn,dy,hr,mi,sc,ms,cs,ns) == nd
    @test NanoDate(Yr,Mn,Dy,Hr,Mi,Sc,Ms,Cs,Ns) == nd
    @test NanoDate(yr,mn,dy,hr,mi,sc,ms) == nd0
    @test NanoDate(Yr,Mn,Dy,Hr,Mi,Sc,Ms) == nd0
    @test NanoDate(yr,mn,dy) == nd00
    @test NanoDate(Yr,Mn,Dy) == nd00
end

@testset "construct by Dates types" begin
    @test NanoDate(dayt, tyme) == nd
    @test NanoDate(daytime, ananosecs) == nd
    @test NanoDate(daytime) == nd0
    @test NanoDate(dayt) == nd00
end


using TerminalUserInterfaces
const TUI = TerminalUserInterfaces
using Random

function main()
    TUI.initialize()
    y, x = 1, 1

    count = 1
    t = TUI.Terminal()

    # TUI.enableRawMode()
    TUI.clear_screen()
    TUI.hide_cursor()

    while true

        w, h = TUI.terminal_size()

        r = TUI.Rect(x, y, w ÷ 2, h ÷ 2)

        b = TUI.Block(title = "Bar Chart")
        p = TUI.BarChart(
            block = b,
            data = [
                ("B1", 1 * count)
                ("B2", 2)
                ("B3", 3 * count)
                ("B4", 4)
                ("B1", 5 * count)
                ("B2", 6)
                ("B3", 7 * count)
                ("B4", 8)
            ],
        )

        TUI.draw(t, p, r)

        TUI.flush(t)

        c = TUI.get_event(t)

        if c == 'j'
            count += 1
        elseif c == 'k'
            count -= 1
        elseif c == 'q'
            break
        end

        if count < 1
            count = 1
        end

    end

    TUI.cleanup()

end

main()


import XLSX

function gen_excel_charstring(L)
    charstring = ""
    total_length = Int('A') + L - 1
    if total_length < 91
        return string(Char(total_length))
    end
    div_count = 0
    tmp = total_length - 91
    first_int = tmp
    while first_int > 90
        div_count += 1
        first_int -= 90
    end
    first_char = Char(Int('A') + div_count)
    second_int = total_length - 91 -(div_count*90) - 1
    second_char = Char(Int('A') + second_int)
    return string(first_char, second_char)
end

function write_xlsx(path, providers, clients, D_MN, D_NN)
    XLSX.openxlsx(path, mode="w") do xf
        sheet = xf[1]
        XLSX.rename!(sheet, "Providers")
        XLSX.addsheet!(xf, "Clients")
        XLSX.addsheet!(xf, "D_MN")
        XLSX.addsheet!(xf, "D_NN")

        sheet = xf["Providers"]
        m = length(providers)
        sheet["A1", dim=2] = ["id", "cx", "cy", "type"]
        sheet["A2", dim=1] = [provider.id for provider in providers]
        sheet["B2", dim=1] = [provider.cx for provider in providers]
        sheet["C2", dim=1] = [provider.cy for provider in providers]
        sheet["D2", dim=1] = [provider.type for provider in providers]

        sheet = xf["Clients"]
        n = length(clients)
        sheet["A1", dim=2] = ["id", "cx", "cy", "type"]
        sheet["A2", dim=1] = [client.id for client in clients]
        sheet["B2", dim=1] = [client.cx for client in clients]
        sheet["C2", dim=1] = [client.cy for client in clients]
        sheet["D2", dim=1] = [client.type for client in clients]

        sheet = xf["D_MN"]
        col_index = gen_excel_charstring(n)
        row_index = m
        sheet[string("A1:", col_index, row_index)] = D_MN

        sheet = xf["D_NN"]
        col_index = gen_excel_charstring(n)
        row_index = n
        sheet[string("A1:", col_index, row_index)] = D_NN
    end
end


const use_return  = haskey(ENV, "USE_RETURN" ) ? parse(Bool, ENV["USE_RETURN"] ) : false
const USE_GPU     = haskey(ENV, "USE_GPU"    ) ? parse(Bool, ENV["USE_GPU"]    ) : false
const do_viz      = haskey(ENV, "DO_VIZ"     ) ? parse(Bool, ENV["DO_VIZ"]     ) : false
const do_save     = haskey(ENV, "DO_SAVE"    ) ? parse(Bool, ENV["DO_SAVE"]    ) : false
const do_save_viz = haskey(ENV, "DO_SAVE_VIZ") ? parse(Bool, ENV["DO_SAVE_VIZ"]) : false
const nx          = haskey(ENV, "NX"         ) ? parse(Int , ENV["NX"]         ) : 64 - 1
const ny          = haskey(ENV, "NY"         ) ? parse(Int , ENV["NY"]         ) : 64 - 1
const nz          = haskey(ENV, "NZ"         ) ? parse(Int , ENV["NZ"]         ) : 64 - 1
###
using ParallelStencil
using ParallelStencil.FiniteDifferences3D
@static if USE_GPU
    @init_parallel_stencil(CUDA, Float64, 3)
else
    @init_parallel_stencil(Threads, Float64, 3)
end
using ImplicitGlobalGrid, Plots, Printf, Statistics, LinearAlgebra, MAT
import MPI
# Global reductions
mean_g(A)    = (mean_l = mean(A); MPI.Allreduce(mean_l, MPI.SUM, MPI.COMM_WORLD)/MPI.Comm_size(MPI.COMM_WORLD))
norm_g(A)    = (sum2_l = sum(A.^2); sqrt(MPI.Allreduce(sum2_l, MPI.SUM, MPI.COMM_WORLD)))
maximum_g(A) = (max_l  = maximum(A); MPI.Allreduce(max_l,  MPI.MAX, MPI.COMM_WORLD))
minimum_g(A) = (min_l  = minimum(A); MPI.Allreduce(min_l,  MPI.MIN, MPI.COMM_WORLD))
# CPU functions
@views av_xza(A) = (A[1:end-1,:,1:end-1] .+ A[1:end-1,:,2:end] .+ A[2:end,:,1:end-1] .+ A[2:end,:,2:end]).*0.25
@views av_zi(A)  = (A[2:end-1,2:end-1,2:end-2] .+ A[2:end-1,2:end-1,3:end-1]).*0.5
@views av_za(A)  = (A[:,:,1:end-1] .+ A[:,:,2:end]).*0.5
@views inn(A)    =  A[2:end-1,2:end-1,2:end-1]

@parallel function smooth!(A2::Data.Array, A::Data.Array, fact::Data.Number)
    @inn(A2)   = @inn(A) + 1.0/6.1/fact*(@d2_xi(A) + @d2_yi(A) + @d2_zi(A))
    return
end

@parallel function compute_maxloc!(Musτ2::Data.Array, Musτ::Data.Array)
    @inn(Musτ2) = @maxloc(Musτ)
    return
end

macro Mu_eff() esc(:(1.0/(1.0/@all(Musτ) + 1.0/(G*dt)))) end
@parallel function compute_iter_params!(dτ_Rho::Data.Array, Gdτ::Data.Array, Musτ::Data.Array, Vpdτ::Data.Number, G::Data.Number, dt::Data.Number, Re::Data.Number, r::Data.Number, max_lxyz::Data.Number)
    @all(dτ_Rho) = Vpdτ*max_lxyz/Re/@Mu_eff()
    @all(Gdτ)    = Vpdτ^2/@all(dτ_Rho)/(r+2)
    return
end

@parallel_indices (ix,iy,iz) function assign_τ!(τxx::Data.Array, τyy::Data.Array, τzz::Data.Array, τxy::Data.Array,  τxz::Data.Array,  τyz::Data.Array, τxx_o::Data.Array, τyy_o::Data.Array, τzz_o::Data.Array, τxy_o::Data.Array, τxz_o::Data.Array, τyz_o::Data.Array)
    if (ix<=size(τxx,1) && iy<=size(τxx,2) && iz<=size(τxx,3))  τxx_o[ix,iy,iz] = τxx[ix,iy,iz]  end
    if (ix<=size(τyy,1) && iy<=size(τyy,2) && iz<=size(τyy,3))  τyy_o[ix,iy,iz] = τyy[ix,iy,iz]  end
    if (ix<=size(τzz,1) && iy<=size(τzz,2) && iz<=size(τzz,3))  τzz_o[ix,iy,iz] = τzz[ix,iy,iz]  end
    if (ix<=size(τxy,1) && iy<=size(τxy,2) && iz<=size(τxy,3))  τxy_o[ix,iy,iz] = τxy[ix,iy,iz]  end
    if (ix<=size(τxz,1) && iy<=size(τxz,2) && iz<=size(τxz,3))  τxz_o[ix,iy,iz] = τxz[ix,iy,iz]  end
    if (ix<=size(τyz,1) && iy<=size(τyz,2) && iz<=size(τyz,3))  τyz_o[ix,iy,iz] = τyz[ix,iy,iz]  end
    return
end

macro inn_yz_Gdτ(ix,iy,iz) esc(:( Gdτ[$ix  ,$iy+1,$iz+1] )) end
macro inn_xz_Gdτ(ix,iy,iz) esc(:( Gdτ[$ix+1,$iy  ,$iz+1] )) end
macro inn_xy_Gdτ(ix,iy,iz) esc(:( Gdτ[$ix+1,$iy+1,$iz  ] )) end
macro inn_yz_Mus(ix,iy,iz) esc(:( Mus[$ix  ,$iy+1,$iz+1] )) end
macro inn_xz_Mus(ix,iy,iz) esc(:( Mus[$ix+1,$iy  ,$iz+1] )) end
macro inn_xy_Mus(ix,iy,iz) esc(:( Mus[$ix+1,$iy+1,$iz  ] )) end
macro av_xyi_Gdτ(ix,iy,iz) esc(:( (Gdτ[$ix  ,$iy  ,$iz+1] + Gdτ[$ix+1,$iy  ,$iz+1] + Gdτ[$ix  ,$iy+1,$iz+1] + Gdτ[$ix+1,$iy+1,$iz+1])*0.25 )) end
macro av_xzi_Gdτ(ix,iy,iz) esc(:( (Gdτ[$ix  ,$iy+1,$iz  ] + Gdτ[$ix+1,$iy+1,$iz  ] + Gdτ[$ix  ,$iy+1,$iz+1] + Gdτ[$ix+1,$iy+1,$iz+1])*0.25 )) end
macro av_yzi_Gdτ(ix,iy,iz) esc(:( (Gdτ[$ix+1,$iy  ,$iz  ] + Gdτ[$ix+1,$iy+1,$iz  ] + Gdτ[$ix+1,$iy  ,$iz+1] + Gdτ[$ix+1,$iy+1,$iz+1])*0.25 )) end
macro av_xyi_Mus(ix,iy,iz) esc(:( (Mus[$ix  ,$iy  ,$iz+1] + Mus[$ix+1,$iy  ,$iz+1] + Mus[$ix  ,$iy+1,$iz+1] + Mus[$ix+1,$iy+1,$iz+1])*0.25 )) end
macro av_xzi_Mus(ix,iy,iz) esc(:( (Mus[$ix  ,$iy+1,$iz  ] + Mus[$ix+1,$iy+1,$iz  ] + Mus[$ix  ,$iy+1,$iz+1] + Mus[$ix+1,$iy+1,$iz+1])*0.25 )) end
macro av_yzi_Mus(ix,iy,iz) esc(:( (Mus[$ix+1,$iy  ,$iz  ] + Mus[$ix+1,$iy+1,$iz  ] + Mus[$ix+1,$iy  ,$iz+1] + Mus[$ix+1,$iy+1,$iz+1])*0.25 )) end
macro inn_yz_Gr(ix,iy,iz) esc(:( @inn_yz_Gdτ($ix,$iy,$iz)/(G*dt) )) end
macro inn_xz_Gr(ix,iy,iz) esc(:( @inn_xz_Gdτ($ix,$iy,$iz)/(G*dt) )) end
macro inn_xy_Gr(ix,iy,iz) esc(:( @inn_xy_Gdτ($ix,$iy,$iz)/(G*dt) )) end
macro av_xyi_Gr(ix,iy,iz) esc(:( @av_xyi_Gdτ($ix,$iy,$iz)/(G*dt) )) end
macro av_xzi_Gr(ix,iy,iz) esc(:( @av_xzi_Gdτ($ix,$iy,$iz)/(G*dt) )) end
macro av_yzi_Gr(ix,iy,iz) esc(:( @av_yzi_Gdτ($ix,$iy,$iz)/(G*dt) )) end
@parallel_indices (ix,iy,iz) function compute_Pt_τ!(Pt::Data.Array, τxx::Data.Array, τyy::Data.Array, τzz::Data.Array, τxy::Data.Array, τxz::Data.Array, τyz::Data.Array, τxx_o::Data.Array, τyy_o::Data.Array, τzz_o::Data.Array, τxy_o::Data.Array, τxz_o::Data.Array, τyz_o::Data.Array, Vx::Data.Array, Vy::Data.Array, Vz::Data.Array, Mus::Data.Array, Gdτ::Data.Array, r::Data.Number, G::Data.Number, dt::Data.Number, _dx::Data.Number, _dy::Data.Number, _dz::Data.Number)
    if (ix<=size(Pt,1)  && iy<=size(Pt,2)  && iz<=size(Pt,3))    Pt[ix,iy,iz] = Pt[ix,iy,iz] - r*Gdτ[ix,iy,iz]*(_dx*(Vx[ix+1,iy,iz] - Vx[ix,iy,iz]) + _dy*(Vy[ix,iy+1,iz] - Vy[ix,iy,iz]) + _dz*(Vz[ix,iy,iz+1] - Vz[ix,iy,iz]))  end
    if (ix<=size(τxx,1) && iy<=size(τxx,2) && iz<=size(τxx,3))  τxx[ix,iy,iz] = ( τxx[ix,iy,iz] + τxx_o[ix,iy,iz]*@inn_yz_Gr(ix,iy,iz) + 2.0*@inn_yz_Gdτ(ix,iy,iz)*(_dx*(Vx[ix+1,iy+1,iz+1] - Vx[ix,iy+1,iz+1])))/(1.0 + @inn_yz_Gdτ(ix,iy,iz)/@inn_yz_Mus(ix,iy,iz) + @inn_yz_Gr(ix,iy,iz))  end
    if (ix<=size(τyy,1) && iy<=size(τyy,2) && iz<=size(τyy,3))  τyy[ix,iy,iz] = ( τyy[ix,iy,iz] + τyy_o[ix,iy,iz]*@inn_xz_Gr(ix,iy,iz) + 2.0*@inn_xz_Gdτ(ix,iy,iz)*(_dy*(Vy[ix+1,iy+1,iz+1] - Vy[ix+1,iy,iz+1])))/(1.0 + @inn_xz_Gdτ(ix,iy,iz)/@inn_xz_Mus(ix,iy,iz) + @inn_xz_Gr(ix,iy,iz))  end
    if (ix<=size(τzz,1) && iy<=size(τzz,2) && iz<=size(τzz,3))  τzz[ix,iy,iz] = ( τzz[ix,iy,iz] + τzz_o[ix,iy,iz]*@inn_xy_Gr(ix,iy,iz) + 2.0*@inn_xy_Gdτ(ix,iy,iz)*(_dz*(Vz[ix+1,iy+1,iz+1] - Vz[ix+1,iy+1,iz])))/(1.0 + @inn_xy_Gdτ(ix,iy,iz)/@inn_xy_Mus(ix,iy,iz) + @inn_xy_Gr(ix,iy,iz))  end
    if (ix<=size(τxy,1) && iy<=size(τxy,2) && iz<=size(τxy,3))  τxy[ix,iy,iz] = ( τxy[ix,iy,iz] + τxy_o[ix,iy,iz]*@av_xyi_Gr(ix,iy,iz) + 2.0*@av_xyi_Gdτ(ix,iy,iz)*(0.5*(_dy*(Vx[ix+1,iy+1,iz+1] - Vx[ix+1,iy,iz+1]) + _dx*(Vy[ix+1,iy+1,iz+1] - Vy[ix,iy+1,iz+1])))) / (1.0 + @av_xyi_Gdτ(ix,iy,iz)/@av_xyi_Mus(ix,iy,iz) + @av_xyi_Gr(ix,iy,iz))  end
    if (ix<=size(τxz,1) && iy<=size(τxz,2) && iz<=size(τxz,3))  τxz[ix,iy,iz] = ( τxz[ix,iy,iz] + τxz_o[ix,iy,iz]*@av_xzi_Gr(ix,iy,iz) + 2.0*@av_xzi_Gdτ(ix,iy,iz)*(0.5*(_dz*(Vx[ix+1,iy+1,iz+1] - Vx[ix+1,iy+1,iz]) + _dx*(Vz[ix+1,iy+1,iz+1] - Vz[ix,iy+1,iz+1])))) / (1.0 + @av_xzi_Gdτ(ix,iy,iz)/@av_xzi_Mus(ix,iy,iz) + @av_xzi_Gr(ix,iy,iz))  end
    if (ix<=size(τyz,1) && iy<=size(τyz,2) && iz<=size(τyz,3))  τyz[ix,iy,iz] = ( τyz[ix,iy,iz] + τyz_o[ix,iy,iz]*@av_yzi_Gr(ix,iy,iz) + 2.0*@av_yzi_Gdτ(ix,iy,iz)*(0.5*(_dz*(Vy[ix+1,iy+1,iz+1] - Vy[ix+1,iy+1,iz]) + _dy*(Vz[ix+1,iy+1,iz+1] - Vz[ix+1,iy,iz+1])))) / (1.0 + @av_yzi_Gdτ(ix,iy,iz)/@av_yzi_Mus(ix,iy,iz) + @av_yzi_Gr(ix,iy,iz))  end
    return
end

macro av_xi_dτ_Rho(ix,iy,iz)  esc(:( (dτ_Rho[$ix  ,$iy+1,$iz+1] + dτ_Rho[$ix+1,$iy+1,$iz+1])*0.5 )) end
macro av_yi_dτ_Rho(ix,iy,iz)  esc(:( (dτ_Rho[$ix+1,$iy  ,$iz+1] + dτ_Rho[$ix+1,$iy+1,$iz+1])*0.5 )) end
macro av_zi_dτ_Rho(ix,iy,iz)  esc(:( (dτ_Rho[$ix+1,$iy+1,$iz  ] + dτ_Rho[$ix+1,$iy+1,$iz+1])*0.5 )) end
@parallel_indices (ix,iy,iz) function compute_V!(Vx::Data.Array, Vy::Data.Array, Vz::Data.Array, Pt::Data.Array, τxx::Data.Array, τyy::Data.Array, τzz::Data.Array, τxy::Data.Array, τxz::Data.Array, τyz::Data.Array, dτ_Rho::Data.Array, _dx::Data.Number, _dy::Data.Number, _dz::Data.Number, nx_1, nx_2, ny_1, ny_2, nz_1, nz_2)
    if (ix<=nx_1 && iy<=ny_2 && iz<=nz_2)  Vx[ix+1,iy+1,iz+1] = Vx[ix+1,iy+1,iz+1] + (_dx*(τxx[ix+1,iy,iz] - τxx[ix,iy,iz]) + _dy*(τxy[ix,iy+1,iz] - τxy[ix,iy,iz]) + _dz*(τxz[ix,iy,iz+1] - τxz[ix,iy,iz]) - _dx*(Pt[ix+1,iy+1,iz+1] - Pt[ix,iy+1,iz+1]))*@av_xi_dτ_Rho(ix,iy,iz)  end
    if (ix<=nx_2 && iy<=ny_1 && iz<=nz_2)  Vy[ix+1,iy+1,iz+1] = Vy[ix+1,iy+1,iz+1] + (_dy*(τyy[ix,iy+1,iz] - τyy[ix,iy,iz]) + _dx*(τxy[ix+1,iy,iz] - τxy[ix,iy,iz]) + _dz*(τyz[ix,iy,iz+1] - τyz[ix,iy,iz]) - _dy*(Pt[ix+1,iy+1,iz+1] - Pt[ix+1,iy,iz+1]))*@av_yi_dτ_Rho(ix,iy,iz)  end
    if (ix<=nx_2 && iy<=ny_2 && iz<=nz_1)  Vz[ix+1,iy+1,iz+1] = Vz[ix+1,iy+1,iz+1] + (_dz*(τzz[ix,iy,iz+1] - τzz[ix,iy,iz]) + _dx*(τxz[ix+1,iy,iz] - τxz[ix,iy,iz]) + _dy*(τyz[ix,iy+1,iz] - τyz[ix,iy,iz]) - _dz*(Pt[ix+1,iy+1,iz+1] - Pt[ix+1,iy+1,iz]))*@av_zi_dτ_Rho(ix,iy,iz)  end
    return
end

@parallel_indices (ix,iy,iz) function compute_Res!(∇V::Data.Array, Rx::Data.Array, Ry::Data.Array, Rz::Data.Array, Vx::Data.Array, Vy::Data.Array, Vz::Data.Array, Pt::Data.Array, τxx::Data.Array, τyy::Data.Array, τzz::Data.Array, τxy::Data.Array, τxz::Data.Array, τyz::Data.Array, _dx::Data.Number, _dy::Data.Number, _dz::Data.Number)
    if (ix<=size(∇V,1) && iy<=size(∇V,2) && iz<=size(∇V,3))  ∇V[ix,iy,iz] = _dx*(Vx[ix+1,iy,iz] - Vx[ix,iy,iz]) + _dy*(Vy[ix,iy+1,iz] - Vy[ix,iy,iz]) + _dz*(Vz[ix,iy,iz+1] - Vz[ix,iy,iz]) end
    if (ix<=size(Rx,1) && iy<=size(Rx,2) && iz<=size(Rx,3))  Rx[ix,iy,iz] = _dx*(τxx[ix+1,iy,iz] - τxx[ix,iy,iz]) + _dy*(τxy[ix,iy+1,iz] - τxy[ix,iy,iz]) + _dz*(τxz[ix,iy,iz+1] - τxz[ix,iy,iz]) - _dx*(Pt[ix+1,iy+1,iz+1] - Pt[ix,iy+1,iz+1]) end
    if (ix<=size(Ry,1) && iy<=size(Ry,2) && iz<=size(Ry,3))  Ry[ix,iy,iz] = _dy*(τyy[ix,iy+1,iz] - τyy[ix,iy,iz]) + _dx*(τxy[ix+1,iy,iz] - τxy[ix,iy,iz]) + _dz*(τyz[ix,iy,iz+1] - τyz[ix,iy,iz]) - _dy*(Pt[ix+1,iy+1,iz+1] - Pt[ix+1,iy,iz+1]) end
    if (ix<=size(Rz,1) && iy<=size(Rz,2) && iz<=size(Rz,3))  Rz[ix,iy,iz] = _dz*(τzz[ix,iy,iz+1] - τzz[ix,iy,iz]) + _dx*(τxz[ix+1,iy,iz] - τxz[ix,iy,iz]) + _dy*(τyz[ix,iy+1,iz] - τyz[ix,iy,iz]) - _dz*(Pt[ix+1,iy+1,iz+1] - Pt[ix+1,iy+1,iz]) end
    return
end

@parallel_indices (iy,iz) function bc_x!(A::Data.Array)
    A[  1, iy,  iz] = A[    2,   iy,   iz]
    A[end, iy,  iz] = A[end-1,   iy,   iz]
    return
end

@parallel_indices (ix,iz) function bc_y!(A::Data.Array)
    A[ ix,  1,  iz] = A[   ix,    2,   iz]
    A[ ix,end,  iz] = A[   ix,end-1,   iz]
    return
end

@parallel_indices (ix,iy) function bc_z!(A::Data.Array)
    A[ ix,  iy,  1] = A[   ix,   iy,    2]
    A[ ix,  iy,end] = A[   ix,   iy,end-1]
    return
end

@views function Stokes3D_()
    # Physics
    lx, ly, lz = 10.0, 10.0, 10.0  # domain extends
    ξ          = 1.0               # Maxwell relaxation time
    μs0        = 1.0               # matrix viscosity
    μsi        = 1e-3              # inclusion viscosity
    G          = 1.0               # elastic shear modulus
    εbg        = 1.0               # background strain-rate
    dt         = μs0/(G*ξ)
    # Numerics
    nt         = 1#5                 # number of time steps
    iterMax    = 200#2e5               # maximum number of pseudo-transient iterations
    nout       = 2000               # error checking frequency
    Re         = 6.0π              # Reynolds number
    r          = 1.0               # Bulk to shear elastic modulus ratio
    CFL        = 0.8/sqrt(3)       # CFL number
    ε          = 1e-8              # nonlinear absolute tolerence
    # nx, ny, nz = 127, 127, 127     # numerical grid resolution; should be a mulitple of 32-1 for optimal GPU perf
    b_width    = (16, 8, 4)         # boundary width for comm/comp overlap
    # Derived numerics
    me, dims, nprocs = init_global_grid(nx, ny, nz) # MPI initialisation
    @static if USE_GPU select_device() end    # select one GPU per MPI local rank (if >1 GPU per node)
    dx, dy, dz = lx/nx_g(), ly/ny_g(), lz/nz_g()      # cell sizes
    max_lxyz   = max(lx,ly,lz)
    Vpdτ       = min(dx,dy,dz)*CFL
    _dx, _dy, _dz = 1.0/dx, 1.0/dy, 1.0/dz
    xc, yc, zc = LinRange(dx/2, lx-dx/2, nx), LinRange(dy/2, ly-dy/2, ny), LinRange(dz/2, lz-dz/2, nz)
    # Array allocations
    Pt         = @zeros(nx  ,ny  ,nz  )
    dτ_Rho     = @zeros(nx  ,ny  ,nz  )
    Gdτ        = @zeros(nx  ,ny  ,nz  )
    ∇V         = @zeros(nx  ,ny  ,nz  )
    Vy         = @zeros(nx  ,ny+1,nz  )
    τxx        = @zeros(nx  ,ny-2,nz-2)
    τyy        = @zeros(nx-2,ny  ,nz-2)
    τzz        = @zeros(nx-2,ny-2,nz  )
    τxy        = @zeros(nx-1,ny-1,nz-2)
    τxz        = @zeros(nx-1,ny-2,nz-1)
    τyz        = @zeros(nx-2,ny-1,nz-1)
    τxx_o      = @zeros(nx  ,ny-2,nz-2)
    τyy_o      = @zeros(nx-2,ny  ,nz-2)
    τzz_o      = @zeros(nx-2,ny-2,nz  )
    τxy_o      = @zeros(nx-1,ny-1,nz-2)
    τxz_o      = @zeros(nx-1,ny-2,nz-1)
    τyz_o      = @zeros(nx-2,ny-1,nz-1)
    Rx         = @zeros(nx-1,ny-2,nz-2)
    Ry         = @zeros(nx-2,ny-1,nz-2)
    Rz         = @zeros(nx-2,ny-2,nz-1)
    Mus2       = @zeros(nx  ,ny  ,nz  )
    Musτ       = @zeros(nx  ,ny  ,nz  )
    # Initial conditions
    Rad2       =  zeros(nx  ,ny  ,nz  )
    Vx         =  zeros(nx+1,ny  ,nz  )
    Vz         =  zeros(nx  ,ny  ,nz+1)
    Rad2      .= [(x_g(ix,dx,Rad2) +0.5*dx -0.5*lx)^2 + (y_g(iy,dy,Rad2) +0.5*dy -0.5*ly)^2 + (z_g(iz,dz,Rad2) +0.5*dz -0.5*lz)^2 for ix=1:size(Rad2,1), iy=1:size(Rad2,2), iz=1:size(Rad2,3)]
    Vx         = Data.Array( -εbg.*[(x_g(ix,dx,Vx) +0.5*dx -0.5*lx) for ix=1:size(Vx,1), iy=1:size(Vx,2), iz=1:size(Vx,3)] )
    Vz         = Data.Array(  εbg.*[(z_g(iz,dz,Vz) +0.5*dz -0.5*lz) for ix=1:size(Vz,1), iy=1:size(Vz,2), iz=1:size(Vz,3)] )
    Mus        = μs0*ones(nx,ny,nz)
    Mus[Rad2.<1.0] .= μsi
    Mus        = Data.Array(Mus)
    Mus2      .= Mus
    for ism=1:10#15
        @hide_communication b_width begin # communication/computation overlap
            @parallel smooth!(Mus2, Mus, 1.0)
            Mus, Mus2 = Mus2, Mus
            update_halo!(Mus)
        end
    end
    Musτ      .= Mus
    @hide_communication b_width begin # communication/computation overlap
        @parallel compute_maxloc!(Musτ, Mus)
        @parallel (1:size(Musτ,2), 1:size(Musτ,3)) bc_x!(Musτ)
        @parallel (1:size(Musτ,1), 1:size(Musτ,3)) bc_y!(Musτ)
        @parallel (1:size(Musτ,1), 1:size(Musτ,2)) bc_z!(Musτ)
        update_halo!(Musτ)
    end
    nx_1, nx_2, ny_1, ny_2, nz_1, nz_2 = nx-1, nx-2, ny-1, ny-2, nz-1, nz-2
    len_Rx_g   = ((nx-2-1)*dims[1]+2)*((ny-2-2)*dims[2]+2)*((nz-2-2)*dims[3]+2)
    len_Ry_g   = ((nx-2-2)*dims[1]+2)*((ny-2-1)*dims[2]+2)*((nz-2-2)*dims[3]+2)
    len_Rz_g   = ((nx-2-2)*dims[1]+2)*((ny-2-2)*dims[2]+2)*((nz-2-1)*dims[3]+2)
    len_∇V_g   = ((nx-2  )*dims[1]+2)*((ny-2  )*dims[2]+2)*((nz-2  )*dims[3]+2)
    # Preparation of visualisation
    if do_viz || do_save_viz
        if (me==0) ENV["GKSwstype"]="nul"; if do_viz !ispath("../../figures") && mkdir("../../figures") end; end
        nx_v, ny_v, nz_v = (nx-2)*dims[1], (ny-2)*dims[2], (nz-2)*dims[3]
        if (nx_v*ny_v*nz_v*sizeof(Data.Number) > 0.8*Sys.free_memory()) error("Not enough memory for visualization.") end
        Pt_v   = zeros(nx_v, ny_v, nz_v) # global array for visu
        Vz_v   = zeros(nx_v, ny_v, nz_v)
        Rz_v   = zeros(nx_v, ny_v, nz_v)
        Mus_v  = zeros(nx_v, ny_v, nz_v)
        τxz_v  = zeros(nx_v, ny_v, nz_v)
        Pt_inn = zeros(nx-2, ny-2, nz-2) # no halo local array for visu
        Vz_inn = zeros(nx-2, ny-2, nz-2)
        Rz_inn = zeros(nx-2, ny-2, nz-2)
        Mus_inn= zeros(nx-2, ny-2, nz-2)
        τxz_inn= zeros(nx-2, ny-2, nz-2)
        y_sl2, y_sl = Int(ceil((ny_g()-2)/2)), Int(ceil(ny_g()/2))
        Xi_g, Zi_g  = dx+dx/2:dx:(lx-dx-dx/2), dz+dz/2:dz:(lz-dz-dz/2) # inner points only
    end
    # Time loop
    @parallel compute_iter_params!(dτ_Rho, Gdτ, Musτ, Vpdτ, G, dt, Re, r, max_lxyz)
    t=0.0; ittot=0; evo_t=[]; evo_τzz=[]
    for it = 1:nt
        err=2*ε; iter=0; err_evo1=[]; err_evo2=[]
        @parallel assign_τ!(τxx, τyy, τzz, τxy, τxz, τyz, τxx_o, τyy_o, τzz_o, τxy_o, τxz_o, τyz_o)
        # Pseudo-transient iteration
        while err > ε && iter <= iterMax
            if (it==1 && iter==11)  tic()  end
            @parallel compute_Pt_τ!(Pt, τxx, τyy, τzz, τxy, τxz, τyz, τxx_o, τyy_o, τzz_o, τxy_o, τxz_o, τyz_o, Vx, Vy, Vz, Mus, Gdτ, r, G, dt, _dx, _dy, _dz)
            @hide_communication b_width begin # communication/computation overlap
                @parallel compute_V!(Vx, Vy, Vz, Pt, τxx, τyy, τzz, τxy, τxz, τyz, dτ_Rho, _dx, _dy, _dz, nx_1, nx_2, ny_1, ny_2, nz_1, nz_2)
                @parallel (1:size(Vy,2), 1:size(Vy,3)) bc_x!(Vy)
                @parallel (1:size(Vz,2), 1:size(Vz,3)) bc_x!(Vz)
                @parallel (1:size(Vx,1), 1:size(Vx,3)) bc_y!(Vx)
                @parallel (1:size(Vz,1), 1:size(Vz,3)) bc_y!(Vz)
                @parallel (1:size(Vx,1), 1:size(Vx,2)) bc_z!(Vx)
                @parallel (1:size(Vy,1), 1:size(Vy,2)) bc_z!(Vy)
                update_halo!(Vx, Vy, Vz)
            end
            iter += 1
            if iter % nout == 0
                @parallel compute_Res!(∇V, Rx, Ry, Rz, Vx, Vy, Vz, Pt, τxx, τyy, τzz, τxy, τxz, τyz, _dx, _dy, _dz)
                Vmin, Vmax = minimum_g(Vx), maximum_g(Vx)
                Pmin, Pmax = minimum_g(Pt), maximum_g(Pt)
                norm_Rx    = norm_g(Rx)/(Pmax-Pmin)*lx/sqrt(len_Rx_g)
                norm_Ry    = norm_g(Ry)/(Pmax-Pmin)*lx/sqrt(len_Ry_g)
                norm_Rz    = norm_g(Rz)/(Pmax-Pmin)*lx/sqrt(len_Rz_g)
                norm_∇V    = norm_g(∇V)/(Vmax-Vmin)*lx/sqrt(len_∇V_g)
                # norm_Rx = norm_g(Rx)/len_Rx_g; norm_Ry = norm_g(Ry)/len_Ry_g; norm_Rz = norm_g(Rz)/len_Rz_g; norm_∇V = norm_g(∇V)/len_∇V_g
                err = maximum([norm_Rx, norm_Ry, norm_Rz, norm_∇V])
                if isnan(err) error("NaN") end
                push!(err_evo1,maximum([norm_Rx, norm_Ry, norm_Rz, norm_∇V])); push!(err_evo2,iter)
                if (me==0) @printf("Step = %d, iter = %d, err = %1.3e [norm_Rx=%1.3e, norm_Ry=%1.3e, norm_Rz=%1.3e, norm_∇V=%1.3e] \n", it, iter, err, norm_Rx, norm_Ry, norm_Rz, norm_∇V) end
            end
        end
        ittot += iter; t += dt
        # push!(evo_t, t); push!(evo_τzz, maximum(τzz))
    end
    # Performance
    wtime    = toc()
    A_eff    = (10*2 + 6*1 + 3)/1e9*nx*ny*nz*sizeof(Data.Number) # Effective main memory access per iteration [GB] (Lower bound of required memory access: Te has to be read and written: 2 whole-array memaccess; Ci has to be read: : 1 whole-array memaccess)
    wtime_it = wtime/(ittot-10)                           # Execution time per iteration [s]
    T_eff    = A_eff/wtime_it                             # Effective memory throughput [GB/s]
    if (me==0) @printf("Total iters = %d (%d steps), time = %1.3e sec (@ T_eff = %1.2f GB/s) \n", ittot, nt, wtime, round(T_eff, sigdigits=3)) end
    # Visualisation
    if do_viz || do_save_viz
        Pt_inn .= inn(Pt);   gather!(Pt_inn, Pt_v)
        Vz_inn .= av_zi(Vz); gather!(Vz_inn, Vz_v)
        Rz_inn .= av_za(Rz); gather!(Rz_inn, Rz_v)
        Mus_inn.= inn(Mus);  gather!(Mus_inn, Mus_v)
        τxz_inn.= av_xza(τxz); gather!(τxz_inn, τxz_v)
        if me==0 && do_viz
            p1 = heatmap(Xi_g, Zi_g, Pt_v[:,y_sl,:]', aspect_ratio=1, xlims=(Xi_g[1],Xi_g[end]), zlims=(Zi_g[1],Zi_g[end]), c=:viridis, title="Pressure")
            p2 = heatmap(Xi_g, Zi_g, Vz_v[:,y_sl,:]', aspect_ratio=1, xlims=(Xi_g[1],Xi_g[end]), zlims=(Zi_g[1],Zi_g[end]), c=:viridis, title="Vz")
            p4 = heatmap(Xi_g, Zi_g, log10.(abs.(Rz_v[:,y_sl2,:]')), aspect_ratio=1,  xlims=(Xi_g[1],Xi_g[end]), zlims=(Zi_g[1],Zi_g[end]), c=:viridis, title="log10(Rz)")
            #p5 = plot(err_evo2,err_evo1, legend=false, xlabel="# iterations", ylabel="log10(error)", linewidth=2, markershape=:circle, markersize=3, labels="max(error)", yaxis=:log10)
            p3 = plot(evo_t, evo_τzz, legend=false, xlabel="time", ylabel="max(τzz)", linewidth=0, markershape=:circle, framestyle=:box, markersize=3)
               #plot!(evo_t, 2.0.*εbg.*μs0.*(1.0.-exp.(.-evo_t.*G./μs0)), linewidth=2.0) # analytical solution
            plot(p1, p2, p4, p3)
            savefig("../../figures/Stokes_3D_ve_perf_xpu_$(nx_g()).png")
        end
    end
    if me==0 && do_save
        !ispath("../../output") && mkdir("../../output")
        open("../../output/out_Stokes3D_ve_xpu_perf.txt","a") do io
            println(io, "$(nprocs) $(nx_g()) $(ny_g()) $(nz_g()) $(ittot) $(wtime) $(A_eff) $(wtime_it) $(T_eff)")
        end
    end
    if me==0 && do_save_viz
        !ispath("../../out_visu") && mkdir("../../out_visu")
        matwrite("../../out_visu/Stokes_3D_ve_perf_xpu.mat", Dict("Pt_3D"=> Pt_v, "Mus_3D"=> Mus_v, "Txz_3D"=> τxz_v, "Vz_3D"=> Vz_v, "dx_3D"=> dx, "dy_3D"=> dy, "dz_3D"=> dz); compress = true)
    end
    finalize_global_grid()
    return xc, yc, zc, Pt
end

if use_return
    xc, yc, zc, P = Stokes3D_();
else
    Stokes3D = begin Stokes3D_(); return; end
end


module RefractionCorrection

include("utils.jl")
include("refract.jl")
include("io.jl")

export cart2pol
export pol2cart
export correct!
export correct
export write!

end


using Test
using Yota
using Yota: Tape, Input, Call, Constant, trace, itrace, play!, transform, binarize_ops
using Yota: mean_grad, setfield_nested!, copy_with, simplegrad, remove_unused
using Yota: find_field_source_var, unwind_iterate, eliminate_common
using Yota: unvcat, unhcat

include("test_trace_cassette.jl")
include("test_trace_interp.jl")
include("test_helpers.jl")
include("gradcheck.jl")
include("test_simple.jl")
include("test_grad.jl")
include("test_dynamic.jl")
include("test_simplegrad.jl")
include("test_update.jl")
include("test_transform.jl")
include("test_examples.jl")


"""
    write(C::ClimGrid, filename::String)

Write to disk ClimGrid C to netCDF file. Still experimental, some attributes are hardcoded.
"""
function write(C::ClimGrid, filename::String)

    # Test extension
    if extension(filename) .!= ".nc"
        filename = string(filename, ".nc")
    end

    # This creates a new NetCDF file
    # The mode "c" stands for creating a new file (clobber)
    ds = Dataset(filename, "c")

    latsymbol, lonsymbol = ClimateTools.getsymbols(C)
    x, y, timevec = ClimateTools.getdims(C)
    longrid, latgrid = ClimateTools.getgrids(C)

    # DIMENSIONS
    defDim(ds, string(lonsymbol),length(x))
    defDim(ds, string(latsymbol), length(y))
    defDim(ds, "time", length(timevec))

    nclon = defVar(ds, string(lonsymbol), Float64, (string(lonsymbol),))
    nclon.attrib["units"] = C.lonunits#"degrees_east"
    nclon.attrib["long_name"] = "longitude"
    nclon.attrib["standard_name"] = "longitude"
    nclon.attrib["actual_range"] = [minimum(x), maximum(x)]

    nclat = defVar(ds, string(latsymbol), Float64, (string(latsymbol),))
    nclat.attrib["units"] = C.latunits#"degrees_north"
    nclat.attrib["long_name"] = "latitude"
    nclat.attrib["standard_name"] = "latitude"
    nclat.attrib["actual_range"] = [minimum(y), maximum(y)]

    if latsymbol != :lat # a lon-lat grid is needed
        ncrlat = defVar(ds, "lat", Float64, (string(lonsymbol), string(latsymbol)))
        ncrlat.attrib["long_name"] = "latitude"
        ncrlat.attrib["units"] = "degrees"
        ncrlat.attrib["standard_name"] = "grid_latitude"
        ncrlat.attrib["axis"] = "Y"
        ncrlat.attrib["coordinate_defines"] = "point"
        ncrlat.attrib["actual_range"] = [minimum(latgrid), maximum(latgrid)]
    end

    if lonsymbol != :lon
        ncrlon = defVar(ds,"lon", Float64, (string(lonsymbol), string(latsymbol)))
        ncrlon.attrib["long_name"] = "longitude"
        ncrlon.attrib["units"] = "degrees"
        ncrlon.attrib["standard_name"] = "grid_longitude"
        ncrlon.attrib["axis"] = "X"
        ncrlon.attrib["coordinate_defines"] = "point"
        ncrlon.attrib["actual_range"] = [minimum(longrid), maximum(longrid)]
    end

    if ClimateTools.@isdefined C.grid_mapping["grid_mapping"]
        ncmapping = defVar(ds, C.grid_mapping["grid_mapping"], Char, ())
    elseif ClimateTools.@isdefined C.grid_mapping["grid_mapping_name"]
        ncmapping = defVar(ds, C.grid_mapping["grid_mapping_name"], Char, ())
    end
    for iattr in keys(C.grid_mapping)
        ncmapping.attrib[iattr] = C.grid_mapping[iattr]
    end

    # time
    timeout = Array{Float64, 1}(undef, length(timevec))
    for itime = 1:length(timevec)
        timeout[itime] = NCDatasets.timeencode([timevec[itime]], C.timeattrib["units"], C.timeattrib["calendar"])[1]
    end

    # daysfrom = string("days since ", string(timevec[1] - Dates.Day(1))[1:10], " ", string(timevec[1])[12:end])

    nctime = defVar(ds,"time", Float64, ("time",))
    nctime.attrib["long_name"] = "time"
    nctime.attrib["standard_name"] = "time"
    nctime.attrib["axis"] = "T"
    nctime.attrib["calendar"] = C.typeofcal#"365_day"
    nctime.attrib["units"] = C.timeattrib["units"]#daysfrom
    # nctime.attrib["bounds"] = "time_bnds"
    nctime.attrib["coordinate_defines"] = "point"

    # Define the variables contained in ClimGrid C
    v = defVar(ds, C.variable, Float32, (string(lonsymbol),string(latsymbol), "time"))

    # write attributes
    for iattr in keys(C.varattribs)
        v.attrib[iattr] = C.varattribs[iattr]
    end

    # Define global attributes
    for iattr in keys(C.globalattribs)
        ds.attrib[iattr] = C.globalattribs[iattr]
    end

    v[:] = C[1].data

    # Time vector
    nctime[:] = timeout

    # Dimensions
    nclon[:] = x
    nclat[:] = y

    # Dimensions
    if latsymbol != :lat
        ncrlon[:] = longrid
    end
    if latsymbol != :lat
        ncrlat[:] = latgrid
    end

    # Close file
    close(ds)

    return true

end


create_ksp() = KSPCreate()

function create_ksp(A::PetscMat)
    ksp = KSPCreate()
    KSPSetOperators(ksp, A, A)
    return ksp
end

function create_ksp(Amat::PetscMat, Pmat::PetscMat)
    ksp = KSPCreate()
    KSPSetOperators(ksp, Amat, Pmat)
    return ksp
end

set_operators!(ksp::PetscKSP, Amat::PetscMat) = KSPSetOperators(ksp, Amat, Amat)
set_operators!(ksp::PetscKSP, Amat::PetscMat, Pmat::PetscMat) = KSPSetOperators(ksp, Amat, Pmat)

solve!(ksp::PetscKSP, b::PetscVec, x::PetscVec) = KSPSolve(ksp, b, x)

function solve(ksp::PetscKSP, b::PetscVec)
    x = VecDuplicate(b)
    KSPSolve(ksp, b, x)
    return x
end

set_up!(ksp::PetscKSP) = KSPSetUp(ksp)

set_from_options!(ksp::PetscKSP) = KSPSetFromOptions(ksp)

Base.show(::IO, ksp::PetscKSP) = KSPView(ksp)

destroy!(ksp::PetscKSP) = KSPDestroy(ksp)


function init_constraint_struct(com, set::ComplementSet{F,S}, internals) where {F,S}
    inner_set = set.set
    fct = internals.fct
    if !(internals.fct isa SAF) && F <: SAF
        fct = get_saf(internals.fct)
    end

    constraint = get_constraint(com, fct, inner_set)
    complement_constraint = get_complement_constraint(com, constraint)
    return complement_constraint
end


module GradientsTestsuite

using ..KernelAbstractions
using ..KernelGradients
using ..Enzyme
using ..Test

include("matmul.jl")

function testsuite(backend, backend_str, backend_mod, AT, DAT)
    @testset "Matmul" begin
        matmul_testsuite(backend, AT)
    end
end

end


module HydrophoneCalibrations

using Interpolations
using Unitful

export volt_to_pressure, volt_to_pressure_and_phase, calibration
export parse_onda

const CALIBRATION_LOADED = Ref{Bool}() # initializes to false

include("configuration_loading.jl")
include("volt_to_pressure.jl")
include("onda_file_parsing.jl")

_load_config()

end


@testset "Plots recipe" begin
    mtg = read_opf("files/simple_plant.opf", Dict)

    recipe = RecipesBase.apply_recipe(Dict{Symbol,Any}(), mtg)

    @test length(recipe) == length(mtg)
    #? NB: both should be equal because 6 lines (all except first point have an edge),
    #? and 1 scatter.

    @test recipe[1].plotattributes == Dict{Symbol,Any}(:label => "", :seriescolor => :black, :seriestype => :line)
    @test recipe[1].args == ([0.0, 0.0], [0.0, 0.2]) # coordinates of the two vertex for the edge
    @test recipe[6].args == ([0.0, 0.7071067811865475], [1.6, 2.3071067811865476]) # same, last one

    df_coordinates = PlantGeom.mtg_coordinates_df(mtg, force = true)

    @test recipe[7].args[1] == df_coordinates.XX
    @test recipe[7].plotattributes == Dict{Symbol,Any}(
        :color => [:black, :black, :black, :black, :black, :black, :black],
        :palette => colorschemes[:viridis],
        :label => "",
        :hover => [
            "name: `node_1`, link: `/`, symbol: `Scene`, index: `1`",
            "name: `node_2`, link: `/`, symbol: `Individual`, index: `2`",
            "name: `node_3`, link: `/`, symbol: `Axis`, index: `3`",
            "name: `node_4`, link: `/`, symbol: `Internode`, index: `4`",
            "name: `node_5`, link: `+`, symbol: `Leaf`, index: `5`",
            "name: `node_6`, link: `<`, symbol: `Internode`, index: `6`",
            "name: `node_7`, link: `+`, symbol: `Leaf`, index: `7`",
        ],
        :seriestype => :scatter,
        :colorbar_entry => false
    )
end


@time @testset "Systems.NonDeterministicInput" begin include("unit_NonDeterministicInput.jl") end
@time @testset "Systems.System" begin include("unit_System.jl") end


# adapted from plots because its not in Recipes Base
function partialcircle(start_θ, end_θ, n = 20, r=1)
    [r*cos(u) for u in linspace(start_θ, end_θ, n)],
    [r*sin(u) for u in linspace(start_θ, end_θ, n)]
end

function make_circle(cx, cy, r)
    x, y = partialcircle(0, 2pi, 20, r)
    x .+= cx
    y .+= cy
    collect(zip(x, y))
end

@recipe function plot(env::CartpoleEnvironment; force=0.)
    sim = env.sim
    xlims --> (sim.xlimits[1], sim.xlimits[2])
    ylims --> sim.pole_length .* (-1.2,1.2)
    ratio --> true
    xlabel --> "x [meter]"
    ylabel --> "y [meter]"
    title --> @sprintf("Cartpole at t=%0.2f s", sim.time)
    min_force, max_force = extrema(env.actions)
    cart_pos = configuration(sim.state)[1]
    cart_offset = -0.04
    cart_height = 0.1
    cart_width = 0.6
    arrow_height = cart_height / 2
    arrow_width_plus = 0.15 / max_force * force
    arrow_width_minus = 0.15 / min_force * force
    pole_width = sim.pole_com == sim.pole_length ? 0.005 : 0.02
    wheel_radius = 0.07
    wheel_offset = 0.2
    wheel_turns = (cart_pos+10) / (2pi * wheel_radius)
    wheel_angle = -2pi*(wheel_turns - floor(wheel_turns))
    actual_xlims = plotattributes[:xlims]
    @series begin # floor
        seriescolor := :grey
        linewidth := 1
        label := ""
        [actual_xlims...], cart_offset .- cart_height .- wheel_radius .+ zeros(2)
    end
    if force > 0
        @series begin
            seriescolor := :red
            seriestype := :shape
            alpha := 0.2
            label := ""
            (cart_pos .+ cart_width./2 .+ [0, 0, arrow_width_plus, arrow_width_plus*1.2, arrow_width_plus],
             cart_offset .- cart_height./2 .+ arrow_height./2 .* [-1, 1, 1, 0, -1])
        end
    elseif force < 0
        @series begin
            seriescolor := :red
            seriestype := :shape
            alpha := 0.2
            label := ""
            (cart_pos .- cart_width./2 .- [0, 0, arrow_width_minus, arrow_width_minus*1.2, arrow_width_minus],
             cart_offset .- cart_height./2 .+ arrow_height./2 .* [-1, 1, 1, 0, -1])
        end
    end
    @series begin # triangle
        seriescolor := :lightgrey
        seriestype := :shape
        label := ""
        cart_pos .+ [-.1, 0., .1], cart_offset .+ [0., 0.1, 0.]
    end
    @series begin # big cart
        seriescolor := :lightgrey
        seriestype := :shape
        label := ""
        cart_pos .+ cart_width .* [-1, 1, 1, -1] ./ 2, cart_offset .+ [-cart_height, -cart_height, 0., 0.]
    end
    @series begin # wheel left
        seriestype := :shape
        seriescolor := :grey
        label := ""
        make_circle(cart_pos - wheel_offset, cart_offset - cart_height, wheel_radius)
    end
    @series begin # wheel right
        seriestype := :shape
        seriescolor := :grey
        label := ""
        make_circle(cart_pos + wheel_offset, cart_offset - cart_height, wheel_radius)
    end
    wheel_x = [0.7wheel_radius*cos(wheel_angle)]
    wheel_y = [0.7wheel_radius*sin(wheel_angle)]
    @series begin # wheel left scatter
        seriescolor := :black
        seriestype := :scatter
        markersize := 1
        label := ""
        wheel_x .+ cart_pos .- wheel_offset, wheel_y .+ cart_offset .- cart_height
    end
    @series begin # wheel right scatter
        seriescolor := :black
        seriestype := :scatter
        markersize := 1
        label := ""
        wheel_x .+ cart_pos .+ wheel_offset, wheel_y .+ cart_offset .- cart_height
    end
    p_bot_l = transform(sim.state,
            Point3D(frame_after(joints(sim.mechansim)[3]), SVector(-pole_width,0,0.)),
            default_frame(bodies(sim.mechansim)[1])).v
    p_bot_r = transform(sim.state,
            Point3D(frame_after(joints(sim.mechansim)[3]), SVector(pole_width,0,0.)),
            default_frame(bodies(sim.mechansim)[1])).v
    p_top_l = transform(sim.state,
            Point3D(frame_after(joints(sim.mechansim)[3]), SVector{3}(-pole_width,0,sim.pole_length)),
            default_frame(bodies(sim.mechansim)[1])).v
    p_top_r = transform(sim.state,
            Point3D(frame_after(joints(sim.mechansim)[3]), SVector{3}(pole_width,0,sim.pole_length)),
            default_frame(bodies(sim.mechansim)[1])).v
    @series begin # pole itself
        seriescolor := :grey
        seriestype := :shape
        label := ""
        [p_bot_l[1],p_bot_r[1],p_top_r[1],p_top_l[1]], [p_bot_l[3],p_bot_r[3],p_top_r[3],p_top_l[3]]
    end
    if sim.pole_com == sim.pole_length
        p_bot_c = transform(sim.state,
            Point3D(frame_after(joints(sim.mechansim)[3]), SVector(0,0,0.)),
            default_frame(bodies(sim.mechansim)[1])).v
        p_top_c = transform(sim.state,
            Point3D(frame_after(joints(sim.mechansim)[3]), SVector(0.,0,sim.pole_length)),
            default_frame(bodies(sim.mechansim)[1])).v
        @series begin # sphere on top
            seriestype := :shape
            seriescolor := :grey
            label := ""
            make_circle(p_top_c[1], p_top_c[3], 0.07)
        end
    end
    @series begin # sphere on bottom
        seriestype := :shape
        seriescolor := :lightgrey
        label := ""
        make_circle(cart_pos, 0., 0.03)
    end
    if get(plotattributes, :legend, :best) != :none
        @series begin # x line
            seriescolor := :grey
            linewidth := 1
            linestyle := :dash
            label := ""
            [cart_pos, cart_pos], [cart_offset .- cart_height .- wheel_radius .- 0.05, sim.pole_length]
        end
        @series begin # x text
            seriescolor := :grey
            label := ""
            series_annotations := [@sprintf("x = %0.2f", cart_pos), ""]
            [cart_pos], [cart_offset .- cart_height .- wheel_radius .- 0.15]
        end
    end
    nothing
end


  # Vinicius Piccoli  - PhD student at M3G - UNICAMP

  #= Informations:

  This simple code reads protein pdb files and writes inputs for Packmol.

  The code only works for a simulation of a protein solvated by water and ionic liquids (e.g EMIMDCA, EMIMBF4, BMIMNO3, BMIMCl, etc.)

  Example:
  gmx pdb2gmx -f protein.pdb (crystallographic structure) -o processed.pdb
  measure_prot(processed.pdb,10) - this will return the dimensions of the protein in x,y and z axis plus 10 Angstrons.

  The conc function write an “inp” file (which will be the input for Packmol and, therefore, contains the information for the simulation box)
  and the topology file for the simulation

  conc(MM::Float64,Nome::String,Nomean::String, conc::Float64)
    MM - Ionic Liquid molar mass
    Nome - Name of the cation
    Nomean - Name of the anion
    conc - Ionic liquid concentration
  =#

  function measure_prot(nome,dist)

  ## Variables - Measurement of the protein(X, Y and Z axis)
    pdbfile   = nome
    segment   = "all"
    restype   = "all"
    atomtype  = "all"
    firstatom = "all"
    lastatom  = "all"

  ## Open the file -
    file = open("$pdbfile","r+")
    natoms = 0
    var = 0

    for line in eachline(file)

      data = split(line)
      consider = false
      if data[1] == "ATOM" || data[1] == "HEATOM"
        natoms = natoms + 1
        ss = data[11]
        rt = data[4]
        at = data[3]
        consider = true
      else
        continue
      end

      if firstatom != "all"
        if natoms < firstatom
          consider = false
        end
      end

      if lastatom != "all"
        if natoms > lastatom
          consider = false
        end
      end

      if segment != "all"
        if ss != segment
          consider = false
        end
      end

      if restype != "all"
        if rt != restype
          consider = false
        end
      end

      if atomtype != "all"
        if at != atomtype
          consider = false
        end
      end

      if consider == true
        global x = parse(Float64,data[7])
        global y = parse(Float64,data[8])
        global z = parse(Float64,data[9])
      end

      if var == 1
        if x < xmin
            xmin = x
        end
        if y < ymin
            ymin = y
        end
        if z < zmin
            zmin = z
        end
        if x > xmax
            xmax = x
        end
        if y > ymax
            ymax = y
        end
        if z > zmax
            zmax = z
        end
      else

        global   xmin = x
        global   ymin = y
        global   zmin = z
        global   xmax = x
        global   ymax = y
        global   zmax = z
        var = 1

      end

    end

  ## Box size - add 30 A in each size of box
    bx = round(Int,((xmax - xmin) + dist ))
    by = round(Int,((ymax - ymin) + dist ))
    bz = round(Int,((zmax - zmin) + dist ))

    return bx , by, bz

  end

  # Function to calculate number of components to put inside a box given a specific concentration

  function conc(MM::Float64,Nome::String,Nomean::String, conc::Float64)

      # number of particles calculationa

      # volume of the box (L)
      vol_box(a,b,c) = a*b*c*1e-27;

      # protein volume (L)
      vol_prot(m) = (m/(6.02e23))*1e-3;

      # Water volume (L)
      vol_wat(n,mw) = n * (mw/(6.02e23))*1e-3;

      # Solution volume
      vol_sol(vc,vp) = vc - vp;

      # number of ionic liquids molecules
      num_il(vs,cil) = round(Int128,(vs*cil*6.02e23));

      # Volume of ionic liquids moleculres
      v_il(nil,mil) = (nil*mil*1e-3)/(6.02e23);

      # each ionic liquid will occupy the half of the volume

      # Number of water molecules
      num_wat(vs,vil) = round(Int128,((vs - vil) * 6.02e23) / (18*1e-3));

      # Box dimensions
      lx,ly,lz = measure_prot("ubq.pdb")   # Ubiquitin pdb

      # Total concentration of ions
      c = conc;

      vs   = vol_box(2*lx,2*ly,2*lz) - vol_prot(8560); # Solution volume / Ubiquitin molar mass = 8560
      nil  = round(Int,(num_il(vs,c)));                # Number of ionic liquid molecules for a concentration "c"

      nwat1 = num_wat(vs, v_il(nil,MM));               # Number of water molecules
      nwat = nwat1

      io = open("box.inp","w")

      println(io,"tolerance 2.0")
      println(io,"output system.pdb")
      println(io,"add_box_sides 1.0")
      println(io,"filetype pdb")
      println(io,"seed -1")
      println(io,"                  ")
      println(io,"structure ubq.pdb")
      println(io," number 1")
      println(io," center")
      println(io," fixed 0. 0. 0. 0. 0. 0.")
      println(io,"end structure")
      println(io,"                  ")
      println(io,"structure WATER.pdb")
      println(io," number $nwat")
      println(io," inside box -$(lx). -$(ly). -$(lz). $(lx). $(ly). $(lz).")
      println(io,"end structure")
      println(io,"                  ")
      println(io,"structure $(Nome)_VSIL.pdb")
      println(io," number $nil")
      println(io," inside box -$(lx). -$(ly). -$(lz). $(lx). $(ly). $(lz).")
      println(io,"end structure")
      println(io,"                  ")
      println(io,"structure $(Nomean)_VSIL.pdb")
      println(io," number $nil")
      println(io," inside box -$(lx). -$(ly). -$(lz). $(lx). $(ly). $(lz).")
      println(io,"end structure")

      close(io)

      topa = open("topola.top","r+")

      # number of lines !
      nlines = 0
      for line in eachline(topa)
        nlines += 1
      end

      close(topa)

      Nomean2="NOTHING"
      Nome2="NOTHING"

      topa = open("topola.top","r+")
      top = open("topol.top","w")
      ntline = 0

      for line in eachline(topa)

        ntline += 1

        if occursin("AN1",line) && ntline != nlines && ntline < nlines - 4
          println(top,replace(line,"AN1" => "$Nomean",count = 1))
        elseif occursin("CA1",line) && ntline < nlines - 4
          println(top,replace(line,"CA1" => "$Nome",count = 1))
        elseif occursin("AN2",line) && ntline < nlines - 4
          println(top,replace(line,"AN2" => "$Nomean2",count = 1))
        elseif occursin("CA2",line) && ntline < nlines - 4
          println(top,replace(line,"CA2" => "$Nome2",count = 1))
        elseif ntline == (nlines - 4)
          println(top,replace(line,line => "SOL               $(nwat+58)",count = 1)) # 58 cristalization water molecules

        # data for the ions

        elseif ntline == (nlines - 1)
          if Nomean == "DCA"
            println(top,replace(line,line => "NC                $nil",count = 1))
          elseif Nomean == "BF4"
            println(top,replace(line,line => "BF4                $nil",count = 1))
          elseif Nomean == "NO3"
            println(top,replace(line,line => "NO3                $nil",count = 1))
          elseif Nomean == "Cl"
            println(top,replace(line,line => "Cl                 $nil",count = 1))
          elseif Nomean == "SCN"
            println(top,replace(line,line => "SCN                $nil",count = 1))
          end
        elseif ntline == (nlines - 3)
          println(top,replace(line,line => "$(Nome[1:3])                $nil ",count = 1))
        else
          println(top,line)
        end
      end
      close(top)
      close(topa)

  end


module house
    include("./fixed_parameters.jl")
    using .fixed_parameters
    include("./human.jl")
    using .human

    mutable struct House
        quarantine::Bool
        warning::Bool
        walls::Array
        house_border::Tuple
        crew::Array

        in_house::Any
        house_time::Any

        function House(walls::Array, number_of_students::Int64, number_of_teachers::Int64, intern_random_infection_rate::Float64, student_break_rules::Float64, teacher_break_rules::Float64)
            house_border = ((walls[1][1][1],walls[1][2][2]), (walls[2][1][1],walls[2][2][1]))
            crew = [Array{Any}(undef, number_of_students), Array{Any}(undef, number_of_teachers)]
            for i in 1:number_of_students
                spawnpoint = (rand(house_border[1][1]+d:0.01:house_border[2][1]-d), rand(house_border[2][2]+d:0.01:house_border[1][2]-d))
                crew[1][i] = Human(spawnpoint, student_break_rules, true, intern_random_infection_rate)
            end
            for i in 1:number_of_teachers
                spawnpoint = (rand(house_border[1][1]+d:0.01:house_border[2][1]-d), rand(house_border[2][2]+d:0.01:house_border[1][2]-d))
                crew[2][i] = Human(spawnpoint, teacher_break_rules, true, intern_random_infection_rate)
            end

            new(false, false, walls, house_border, crew, in_house, house_time)
        end
    end

    function in_house(self::House, position::Tuple)
        if position[1]>=self.house_border[1][1] && position[1]<=self.house_border[2][1] && position[2]>=self.house_border[2][2] && position[2]<=self.house_border[1][2]
            return true
        else
            return false
        end
    end

    function house_time(self::House, fakewalls::Array)
        for i in self.crew
            for j in i
                if !self.in_house(self, j.position) && (j.destination==nothing || !self.in_house(self, j.destination))
                    j.destination = (rand(self.house_border[1][1]+d:0.01:self.house_border[2][1]-d), rand(self.house_border[2][2]+d:0.01:self.house_border[1][2]-d))
                    j.step = move_speed
                elseif self.in_house(self, j.position) && j.destination!=nothing && sqrt((j.destination[1]-j.position[1])^2+(j.destination[2]-j.position[2])^2)<j.step
                    j.destination = nothing
                    j.step = indoor_speed
                elseif self.in_house(self, j.position) && j.destination==nothing
                    j.step = indoor_speed
                    if self.quarantine
                        j.step = quarantine_speed
                    end
                end
                loc_1 = j.position[2]>lund_size[2]-hwh || j.position[2]<cwh
                j.move(j, fakewalls)
                loc_2 = j.position[2]>lund_size[2]-hwh || j.position[2]<cwh
                if loc_1!=loc_2
                    if !loc_2 && (self.warning || rand()<=1-j.break_rules)
                        j.radius = radius_with_masks
                    else
                        j.radius = radius_without_masks
                    end
                end
            end
        end
    end

    export House

end


const LAYERSKELETON = 64

checklayer(J::Int) = J÷LAYERSKELETON == J/LAYERSKELETON

struct ThinSphericalHarmonicPlan{T} <: SphericalHarmonicPlan{T}
    RP::RotationPlan{T}
    BF::Vector{Butterfly{T}}
    p1::NormalizedLegendreToChebyshevPlan{T}
    p2::NormalizedLegendre1ToChebyshev2Plan{T}
    p1inv::ChebyshevToNormalizedLegendrePlan{T}
    p2inv::Chebyshev2ToNormalizedLegendre1Plan{T}
    B::Matrix{T}
end

function ThinSphericalHarmonicPlan(A::Matrix{T}, L::Int; opts...) where T
    M, N = size(A)
    n = (N+1)÷2
    RP = RotationPlan(T, M-1)
    p1 = plan_normleg2cheb(A)
    p2 = plan_normleg12cheb2(A)
    p1inv = plan_cheb2normleg(A)
    p2inv = plan_cheb22normleg1(A)
    B = zero(A)
    Ce = Matrix{T}(I, M, M)
    Co = Matrix{T}(I, M, M)
    BF = Vector{Butterfly{T}}(undef, n-2)
    P = Progress(n-2, 0.1, "Pre-computing...", 43)
    for J = 1:2:n-2
        mul!(Ce, RP.layers[J])
        checklayer(J+1) && (BF[J] = Butterfly(Ce, L; isorthogonal = true, opts...))
        next!(P)
    end
    for J = 2:2:n-2
        mul!(Co, RP.layers[J])
        checklayer(J) && (BF[J] = Butterfly(Co, L; isorthogonal = true, opts...))
        next!(P)
    end
    ThinSphericalHarmonicPlan(RP, BF, p1, p2, p1inv, p2inv, B)
end

ThinSphericalHarmonicPlan(A::Matrix; opts...) = ThinSphericalHarmonicPlan(A, floor(Int, log2(size(A, 1)+1)-6); opts...)

if VERSION ≥ v"0.7-"
    adjoint(P::ThinSphericalHarmonicPlan) = Adjoint(P)
    transpose(P::ThinSphericalHarmonicPlan) = Transpose(P)
end

function LAmul!(Y::Matrix, TP::ThinSphericalHarmonicPlan, X::Matrix)
    RP, BF, p1, p2, B = TP.RP, TP.BF, TP.p1, TP.p2, TP.B
    copyto!(B, X)
    M, N = size(X)

    @stepthreads for J = 3:2:N÷2
        if checklayer(J-1)
            mul_col_J!(Y, BF[J-1], B, 2J)
            2J < N && mul_col_J!(Y, BF[J-1], B, 2J+1)
        else
            ℓ = round(Int, (J-1)÷LAYERSKELETON)*LAYERSKELETON
            mul_col_J!(RP, B, 2J, ℓ+1, J-1)
            2J < N && mul_col_J!(RP, B, 2J+1, ℓ+1, J-1)
            if ℓ > LAYERSKELETON-2
                mul_col_J!(Y, BF[ℓ], B, 2J)
                2J < N && mul_col_J!(Y, BF[ℓ], B, 2J+1)
            else
                copyto!(Y, 1+M*(2J-1), B, 1+M*(2J-1), 2M)
            end
        end
    end

    @stepthreads for J = 2:2:N÷2
        if checklayer(J)
            mul_col_J!(Y, BF[J-1], B, 2J)
            2J < N && mul_col_J!(Y, BF[J-1], B, 2J+1)
        else
            ℓ = round(Int, J÷LAYERSKELETON)*LAYERSKELETON
            mul_col_J!(RP, B, 2J, ℓ, J-1)
            2J < N && mul_col_J!(RP, B, 2J+1, ℓ, J-1)
            if ℓ > LAYERSKELETON-2
                mul_col_J!(Y, BF[ℓ-1], B, 2J)
                2J < N && mul_col_J!(Y, BF[ℓ-1], B, 2J+1)
            else
                copyto!(Y, 1+M*(2J-1), B, 1+M*(2J-1), 2M)
            end
        end
    end

    copyto!(Y, 1, X, 1, 3M)
    copyto!(B, Y)
    fill!(Y, zero(eltype(Y)))

    mul_col_J!!(Y, p1, B, 1)
    @stepthreads for J = 2:4:N
        mul_col_J!!(Y, p2, B, J)
        J < N && mul_col_J!!(Y, p2, B, J+1)
    end
    @stepthreads for J = 4:4:N
        mul_col_J!!(Y, p1, B, J)
        J < N && mul_col_J!!(Y, p1, B, J+1)
    end
    Y
end

if VERSION < v"0.7-"
    function Base.At_mul_B!(Y::Matrix, TP::ThinSphericalHarmonicPlan, X::Matrix)
        RP, BF, p1inv, p2inv, B = TP.RP, TP.BF, TP.p1inv, TP.p2inv, TP.B
        copyto!(B, X)
        M, N = size(X)
        mul_col_J!!(Y, p1inv, B, 1)
        @stepthreads for J = 2:4:N
            mul_col_J!!(Y, p2inv, B, J)
            J < N && mul_col_J!!(Y, p2inv, B, J+1)
        end
        @stepthreads for J = 4:4:N
            mul_col_J!!(Y, p1inv, B, J)
            J < N && mul_col_J!!(Y, p1inv, B, J+1)
        end

        copyto!(B, Y)
        fill!(Y, zero(eltype(Y)))
        copyto!(Y, 1, B, 1, 3M)

        @stepthreads for J = 3:2:N÷2
            if checklayer(J-1)
                At_mul_B_col_J!(Y, BF[J-1], B, 2J)
                2J < N && At_mul_B_col_J!(Y, BF[J-1], B, 2J+1)
            else
                ℓ = round(Int, (J-1)÷LAYERSKELETON)*LAYERSKELETON
                if ℓ > LAYERSKELETON-2
                    At_mul_B_col_J!(Y, BF[ℓ], B, 2J)
                    2J < N && At_mul_B_col_J!(Y, BF[ℓ], B, 2J+1)
                else
                    copyto!(Y, 1+M*(2J-1), B, 1+M*(2J-1), 2M)
                end
                At_mul_B_col_J!(RP, Y, 2J, ℓ+1, J-1)
                2J < N && At_mul_B_col_J!(RP, Y, 2J+1, ℓ+1, J-1)
            end
        end

        @stepthreads for J = 2:2:N÷2
            if checklayer(J)
                At_mul_B_col_J!(Y, BF[J-1], B, 2J)
                2J < N && At_mul_B_col_J!(Y, BF[J-1], B, 2J+1)
            else
                ℓ = round(Int, J÷LAYERSKELETON)*LAYERSKELETON
                if ℓ > LAYERSKELETON-2
                    At_mul_B_col_J!(Y, BF[ℓ-1], B, 2J)
                    2J < N && At_mul_B_col_J!(Y, BF[ℓ-1], B, 2J+1)
                else
                    copyto!(Y, 1+M*(2J-1), B, 1+M*(2J-1), 2M)
                end
                At_mul_B_col_J!(RP, Y, 2J, ℓ, J-1)
                2J < N && At_mul_B_col_J!(RP, Y, 2J+1, ℓ, J-1)
            end
        end

        sph_zero_spurious_modes!(Y)
    end

    Base.Ac_mul_B!(Y::Matrix, TP::ThinSphericalHarmonicPlan, X::Matrix) = At_mul_B!(Y, TP, X)
else
    function LinearAlgebra.mul!(Y::Matrix, TPt::Transpose{T,<:ThinSphericalHarmonicPlan}, X::Matrix) where T
        TP = parent(TPt)
        RP, BF, p1inv, p2inv, B = TP.RP, TP.BF, TP.p1inv, TP.p2inv, TP.B
        copyto!(B, X)
        M, N = size(X)
        mul_col_J!!(Y, p1inv, B, 1)
        @stepthreads for J = 2:4:N
            mul_col_J!!(Y, p2inv, B, J)
            J < N && mul_col_J!!(Y, p2inv, B, J+1)
        end
        @stepthreads for J = 4:4:N
            mul_col_J!!(Y, p1inv, B, J)
            J < N && mul_col_J!!(Y, p1inv, B, J+1)
        end

        copyto!(B, Y)
        fill!(Y, zero(eltype(Y)))
        copyto!(Y, 1, B, 1, 3M)

        @stepthreads for J = 3:2:N÷2
            if checklayer(J-1)
                At_mul_B_col_J!(Y, BF[J-1], B, 2J)
                2J < N && At_mul_B_col_J!(Y, BF[J-1], B, 2J+1)
            else
                ℓ = round(Int, (J-1)÷LAYERSKELETON)*LAYERSKELETON
                if ℓ > LAYERSKELETON-2
                    At_mul_B_col_J!(Y, BF[ℓ], B, 2J)
                    2J < N && At_mul_B_col_J!(Y, BF[ℓ], B, 2J+1)
                else
                    copyto!(Y, 1+M*(2J-1), B, 1+M*(2J-1), 2M)
                end
                At_mul_B_col_J!(RP, Y, 2J, ℓ+1, J-1)
                2J < N && At_mul_B_col_J!(RP, Y, 2J+1, ℓ+1, J-1)
            end
        end

        @stepthreads for J = 2:2:N÷2
            if checklayer(J)
                At_mul_B_col_J!(Y, BF[J-1], B, 2J)
                2J < N && At_mul_B_col_J!(Y, BF[J-1], B, 2J+1)
            else
                ℓ = round(Int, J÷LAYERSKELETON)*LAYERSKELETON
                if ℓ > LAYERSKELETON-2
                    At_mul_B_col_J!(Y, BF[ℓ-1], B, 2J)
                    2J < N && At_mul_B_col_J!(Y, BF[ℓ-1], B, 2J+1)
                else
                    copyto!(Y, 1+M*(2J-1), B, 1+M*(2J-1), 2M)
                end
                At_mul_B_col_J!(RP, Y, 2J, ℓ, J-1)
                2J < N && At_mul_B_col_J!(RP, Y, 2J+1, ℓ, J-1)
            end
        end

        sph_zero_spurious_modes!(Y)
    end

    LinearAlgebra.mul!(Y::Matrix, TPc::Adjoint{T,<:ThinSphericalHarmonicPlan}, X::Matrix) where T =
        mul!(Y, transpose(parent(TPc)), X)
end
allranks(TP::ThinSphericalHarmonicPlan) = mapreduce(i->allranks(TP.BF[i]),vcat,sort!([LAYERSKELETON-1:LAYERSKELETON:length(TP.BF);LAYERSKELETON:LAYERSKELETON:length(TP.BF)]))

function mul_col_J!(P::RotationPlan, A::AbstractMatrix, J::Int, L1::Int, L2::Int)
    M, N = size(A)
    @inbounds for m = L2-1:-2:L1
        layer = P.layers[m+1]
        @simd for i = 1:length(layer)
            G = layer[i]
            a1, a2 = A[G.i1,J], A[G.i2,J]
            A[G.i1,J] = G.c*a1 + G.s*a2
            A[G.i2,J] = G.c*a2 - G.s*a1
        end
    end
    A
end

function At_mul_B_col_J!(P::RotationPlan, A::AbstractMatrix, J::Int, L1::Int, L2::Int)
    M, N = size(A)
    @inbounds for m = L1:2:L2-1
        layer = P.layers[m+1]
        @simd for i = length(layer):-1:1
            G = layer[i]
            a1, a2 = A[G.i1,J], A[G.i2,J]
            A[G.i1,J] = G.c*a1 - G.s*a2
            A[G.i2,J] = G.c*a2 + G.s*a1
        end
    end
    A
end


## Si-Midorikawa 1999 GMPE tests
@testset "Si-Midorikawa 1999 GMPE PGA" begin
  # init model parameters
  include("../examples/si-midorikawa-1999.conf")
  ## test at epicenter on grid M7.0
  @test gmpe_simidorikawa1999(eq_7,config_simidorikawa1999_crustal_pga,grid_epicenter)[1].pga == 59.04
  ## run PGA modeling on grid withoit minpga Depth <= 30 M6.0
  S_c = gmpe_simidorikawa1999(eq_6,config_simidorikawa1999_crustal_pga,grid)
  @test length(S_c) == TEST_GRID_SIZE
  @test round(sum([S_c[i].pga for i=1:length(S_c)]),digits=2) == 6.98
  S_intp = gmpe_simidorikawa1999(eq_6,config_simidorikawa1999_interplate_pga,grid)
  @test length(S_intp) == TEST_GRID_SIZE
  @test round(sum([S_intp[i].pga for i=1:length(S_intp)]),digits=2) == 8.38
  S_intra = gmpe_simidorikawa1999(eq_6,config_simidorikawa1999_intraplate_pga,grid)
  @test length(S_intra) == TEST_GRID_SIZE
  @test round(sum([S_intra[i].pga for i=1:length(S_intra)]),digits=2) == 13.9
  ## run PGA modeling on grid with minpga Depth <= 30 M6.0
  S_c = gmpe_simidorikawa1999(eq_6,config_simidorikawa1999_crustal_pga,grid,min_val=0.34)
  @test length(S_c) == WITH_MINPGA
  @test round(sum([S_c[i].pga for i=1:length(S_c)]),digits=2) == 4.61
  ## run PGA modeling on grid with minpga Depth > 30 M6.0
  eq_30 = Earthquake(143.04,51.92,35,6.0)
  S_c = gmpe_simidorikawa1999(eq_30,config_simidorikawa1999_crustal_pga,grid,min_val=0.15)
  @test length(S_c) == WITH_MINPGA
  @test round(sum([S_c[i].pga for i=1:length(S_c)]),digits=2) == 2.04
  ## run PGA modeling for plotting Depth <= 30 M6.0
  S_c = gmpe_simidorikawa1999(eq_6,config_simidorikawa1999_crustal_pga)
  @test length(S_c) == SIMULATION_ARRAY_SIZE
  @test round(sum(S_c),digits=2) == 1096.79
  S_intp = gmpe_simidorikawa1999(eq_6,config_simidorikawa1999_interplate_pga)
  @test round(sum(S_intp),digits=2) == 1318.74
  S_intra = gmpe_simidorikawa1999(eq_6,config_simidorikawa1999_intraplate_pga)
  @test round(sum(S_intra),digits=2) == 2188.89
  ## run PGA modeling for plotting Depth > 30 M6.0
  S_c = gmpe_simidorikawa1999(eq_30,config_simidorikawa1999_crustal_pga)
  @test round(sum(S_c),digits=2) == 722.93
end


import Base: size, eltype
import RegularizedLeastSquares: initkaczmarz, dot_with_matrix_row, kaczmarz_update!

export reconstructionMultiPatch, MultiPatchOperator

# Necessary for Multi-System-Matrix FF reconstruction
voxelSize(bSF::MultiMPIFile) = voxelSize(bSF[1])
sfGradient(bSF::MultiMPIFile,dim) = sfGradient(bSF[1],dim)
generateHeaderDict(bSF::MultiMPIFile,bMeas::MPIFile) =
   generateHeaderDict(bSF[1],bMeas)

function reconstructionMultiPatch(bSF, bMeas::MPIFile;
  minFreq=0, maxFreq=1.25e6, SNRThresh=-1,maxMixingOrder=-1, numUsedFreqs=-1, sortBySNR=false, recChannels=1:numReceivers(bMeas), kargs...)

  freq = filterFrequencies(bSF,minFreq=minFreq, maxFreq=maxFreq,recChannels=recChannels, SNRThresh=SNRThresh, numUsedFreqs=numUsedFreqs, sortBySNR=sortBySNR)

  @debug "selecting $(length(freq)) frequencies"

  return reconstructionMultiPatch(bSF, bMeas, freq; kargs...)
end

function reconstructionMultiPatch(bSF, bMeas::MPIFile, freq;
            frames=nothing, bEmpty=nothing, nAverages=1, numAverages=nAverages,
	    spectralLeakageCorrection=true, voxelsize=voxelSize(bSF), kargs...)

  bgCorrection = (bEmpty != nothing)

  FFOp = MultiPatchOperatorHighLevel(bSF, bMeas, freq, bgCorrection;
                    kargs... )

  L = acqNumFGFrames(bMeas)
  (frames==nothing) && (frames=collect(1:L))
  nFrames=length(frames)

  uTotal_ = getMeasurementsFD(bMeas,frequencies=freq, frames=frames, numAverages=numAverages,
                             spectralLeakageCorrection=spectralLeakageCorrection)

  periodsSortedbyFFPos = unflattenOffsetFieldShift(ffPos(bMeas))
  uTotal = similar(uTotal_,size(uTotal_,1),length(periodsSortedbyFFPos),size(uTotal_,3))

  for k=1:length(periodsSortedbyFFPos)
      uTotal[:,k,:] = mean(uTotal_[:,periodsSortedbyFFPos[k],:], dims=2)
  end

  # Here we call a regular reconstruction function
  c_ = reconstruction(FFOp, uTotal; kargs...)
  c = reshape(c_, shape(FFOp.grid)..., :)

  # calculate axis
  shp = size(c)
  pixspacing = (voxelsize ./ sfGradient(bMeas,3) .* sfGradient(bSF,3)) * 1000u"mm"
  offset = (fieldOfViewCenter(FFOp.grid) .- 0.5.*fieldOfView(FFOp.grid) .+ 0.5.*spacing(FFOp.grid)) * 1000u"mm"

  # TODO does this provide the correct value in the multi-patch case?
  dtframes = acqNumAverages(bMeas)*dfCycle(bMeas)*numAverages*1u"s"
  # create image
  c = reshape(c,1,size(c)...)
  im = makeAxisArray(c, pixspacing, offset, dtframes)
  imMeta = ImageMeta(im,generateHeaderDict(bSF,bMeas))
  return imMeta
end

# MultiPatchOperator is a type that acts as the MPI system matrix but exploits
# its sparse structure.
# Its very important to keep this type typestable
mutable struct MultiPatchOperator{V<:AbstractMatrix, U<:Positions}
  S::Vector{V}
  grid::U
  N::Int
  M::Int
  RowToPatch::Vector{Int}
  xcc::Vector{Vector{Int}}
  xss::Vector{Vector{Int}}
  sign::Matrix{Int}
  nPatches::Int
  patchToSMIdx::Vector{Int}
end

eltype(FFOp::MultiPatchOperator) = eltype(FFOp.S[1])

function MultiPatchOperatorHighLevel(SF::MPIFile, bMeas, freq, bgCorrection::Bool; kargs...)
  return MultiPatchOperatorHighLevel(MultiMPIFile([SF]), bMeas, freq, bgCorrection; kargs...)
end

function MultiPatchOperatorHighLevel(bSF::MultiMPIFile, bMeas, freq, bgCorrection::Bool;
        FFPos = zeros(0,0), FFPosSF = zeros(0,0), kargs...)

  FFPos_ = ffPos(bMeas)

  periodsSortedbyFFPos = unflattenOffsetFieldShift(FFPos_)
  idxFirstPeriod = getindex.(periodsSortedbyFFPos,1)
  FFPos_ = FFPos_[:,idxFirstPeriod]

  if length(FFPos) > 0
    FFPos_[:] = FFPos
  end

  if length(FFPosSF) == 0
    L = length(ffPos(bSF[1]))
    FFPosSF_ = [vec(ffPos(SF))[l] for l=1:L, SF in bSF] #[vec(ffPos(SF)) for SF in bSF]
  else
    FFPosSF_ = FFPosSF #[vec(FFPosSF[:,l]) for l=1:size(FFPosSF,2)]
  end

  gradient = acqGradient(bMeas)[:,:,1,idxFirstPeriod]

  FFOp = MultiPatchOperator(bSF, bMeas, freq, bgCorrection,
                  FFPos = FFPos_,
                  gradient = gradient,
                  FFPosSF = FFPosSF_; kargs...)
  return FFOp
end

function MultiPatchOperator(SF::MPIFile, bMeas, freq, bgCorrection::Bool; kargs...)
  return MultiPatchOperator(MultiMPIFile([SF]), bMeas, freq, bgCorrection; kargs...)
end

function findNearestPatch(ffPosSF, FFPos, gradientSF, gradient)
  idx = -1
  minDist = 1e20
  for l = 1:size(ffPosSF,2)
    if gradientSF[l][:,:,1,1] == gradient
      dist = norm(ffPosSF[:,l].-FFPos)
      if dist < minDist
        minDist = dist
        idx = l
      end
    end
  end
  if idx < 0
    error("Something went wrong")
  end
  return idx
end

function MultiPatchOperator(SFs::MultiMPIFile, bMeas, freq, bgCorrection::Bool;
        mapping=zeros(0), kargs...)
  if length(mapping) > 0
    return MultiPatchOperatorExpliciteMapping(SFs,bMeas,freq,bgCorrection; mapping=mapping, kargs...)
  else
    return MultiPatchOperatorRegular(SFs,bMeas,freq,bgCorrection; kargs...)
  end
end

function MultiPatchOperatorExpliciteMapping(SFs::MultiMPIFile, bMeas, freq, bgCorrection::Bool;
                    denoiseWeight=0, FFPos=zeros(0,0), FFPosSF=zeros(0,0),
                    gradient=zeros(0,0,0),
                    roundPatches = false,
                    SFGridCenter = zeros(0,0),
                    systemMatrices = nothing,
                    mapping=zeros(0),
		    calibsize = nothing,
                    grid = nothing, kargs...)

  @debug "Loading System matrix"
  numPatches = size(FFPos,2)
  M = length(freq)
  RowToPatch = kron(collect(1:numPatches), ones(Int,M))

  if systemMatrices == nothing
    S = [getSF(SF,freq,nothing,"kaczmarz", bgCorrection=bgCorrection)[1] for SF in SFs]
  else
    S = systemMatrices
  end

  if length(SFGridCenter) == 0
    SFGridCenter = zeros(3,length(SFs))
    for l=1:length(SFs)
      SFGridCenter[:,l] = calibFovCenter(SFs[l])
    end
  end

  gradientSF = [acqGradient(SF) for SF in SFs]

  grids = RegularGridPositions[]
  patchToSMIdx = mapping

  sign = ones(Int, M, numPatches)

  # We first check which system matrix fits best to each patch. Here we use only
  # those system matrices where the gradient matches. If the gradient matches, we take
  # the system matrix with the closes focus field shift
  for k=1:numPatches
    idx = mapping[k]
    SF = SFs[idx]

    diffFFPos = FFPosSF[:,idx] .- FFPos[:,k]

    calibsizeTmp = (calibsize == nothing) ? calibSize(SF) : calibsize

    push!(grids, RegularGridPositions(calibsizeTmp,calibFov(SF),SFGridCenter[:,idx].-diffFFPos))
  end

  # We now know all the subgrids for each patch, if the corresponding system matrix would be taken as is
  # and if a possible focus field missmatch has been taken into account (by changing the center)
  @debug "Calculate Reconstruction Grid"
  if grid == nothing
    recoGrid = RegularGridPositions(grids)
  else
    recoGrid = grid
    # Test whether given grid covers all subgrids
    recoGridTest = RegularGridPositions(grids)
    testMin = round.(recoGridTest.center,digits=3) - round.(recoGridTest.fov./2,digits=3)
    testMax = round.(recoGridTest.center,digits=3) + round.(recoGridTest.fov./2,digits=3)
    gridMin = round.(recoGrid.center,digits=3) - round.(recoGrid.fov./2,digits=3)
    gridMax = round.(recoGrid.center,digits=3) + round.(recoGrid.fov./2,digits=3)
    if false in ((gridMin .<= testMin) .* (gridMax .>= testMax))
        @warn "A larger grid is used for reconstruction, since the given grid does not cover all subgrids: $recoGrid ⊅  $recoGridTest"
        recoGrid = recoGridTest
    end
  end

  # Within the next loop we will refine our grid since we now know our reconstruction grid
  for k=1:numPatches
    idx = mapping[k]
    SF = SFs[idx]

    issubgrid = isSubgrid(recoGrid,grids[k])
    if !issubgrid
      grids[k] = deriveSubgrid(recoGrid, grids[k])
    end
  end
  @debug "Use $(length(S)) patches"

  @debug "Calculate LUT"
  # now that we have all grids we can calculate the indices within the recoGrid
  xcc, xss = calculateLUT(grids, recoGrid)

  return MultiPatchOperator(S, recoGrid, length(recoGrid), M*numPatches,
             RowToPatch, xcc, xss, sign, numPatches, patchToSMIdx)
end

function MultiPatchOperatorRegular(SFs::MultiMPIFile, bMeas, freq, bgCorrection::Bool;
                    denoiseWeight=0, gradient=zeros(0,0,0),
                    roundPatches = false, FFPos=zeros(0,0), FFPosSF=zeros(0,0),
                    kargs...)

  @debug "Loading System matrix"
  numPatches = size(FFPos,2)
  M = length(freq)
  RowToPatch = kron(collect(1:numPatches), ones(Int,M))

  S = AbstractMatrix[]
  SOrigIdx = Int[]
  SIsPlain = Bool[]

  gradientSF = [acqGradient(SF) for SF in SFs]

  grids = RegularGridPositions[]
  matchingSMIdx = zeros(Int,numPatches)
  patchToSMIdx = zeros(Int,numPatches)

  # We first check which system matrix fits best to each patch. Here we use only
  # those system matrices where the gradient matches. If the gradient matches, we take
  # the system matrix with the closes focus field shift
  for k=1:numPatches
    idx = findNearestPatch(FFPosSF, FFPos[:,k], gradientSF, gradient[:,:,k])

    SF = SFs[idx]

    if isapprox(FFPosSF[:,idx],FFPos[:,k])
      diffFFPos = zeros(3)
    else
      diffFFPos = FFPosSF[:,idx] .- FFPos[:,k]
    end
    push!(grids, RegularGridPositions(calibSize(SF),calibFov(SF),calibFovCenter(SF).-diffFFPos))
    matchingSMIdx[k] = idx
  end

  # We now know all the subgrids for each patch, if the corresponding system matrix would be taken as is
  # and if a possible focus field missmatch has been taken into account (by changing the center)
  @debug "Calculate Reconstruction Grid"
  recoGrid = RegularGridPositions(grids)

  # Within the next loop we will refine our grid since we now know our reconstruction grid
  for k=1:numPatches
    idx = matchingSMIdx[k]
    SF = SFs[idx]

    issubgrid = isSubgrid(recoGrid,grids[k])
    if !issubgrid &&
       roundPatches &&
       spacing(recoGrid) == spacing(grids[k])

      issubgrid = true
      grids[k] = deriveSubgrid(recoGrid, grids[k])

    end

    # if the patch is a true subgrid we don't need to apply interpolation and can load the
    # matrix as is.
    if issubgrid
      # we first check if the matrix is already in memory
      u = -1
      for l=1:length(SOrigIdx)
        if SOrigIdx[l] == idx && SIsPlain[l]
          u = l
          break
        end
      end
      if u > 0 # its already in memory
        patchToSMIdx[k] = u
      else     # not yet in memory  -> load it
        S_, grid = getSF(SF,freq,nothing,"kaczmarz", bgCorrection=bgCorrection)
        push!(S,S_)
        push!(SOrigIdx,idx)
        push!(SIsPlain,true) # mark this as a plain system matrix (without interpolation)
        patchToSMIdx[k] = length(S)
      end
    else
      # in this case the patch grid does not fit onto the reco grid. Lets derive a subgrid
      # that is very similar to grids[k]
      newGrid = deriveSubgrid(recoGrid, grids[k])

      # load the matrix on the new subgrid
      S_, grid = getSF(SF,freq,nothing,"kaczmarz", bgCorrection=bgCorrection,
                   gridsize=shape(newGrid),
                   fov=fieldOfView(newGrid),
                   center=fieldOfViewCenter(newGrid).-fieldOfViewCenter(grids[k]))
                   # @TODO: I don't know the sign of aboves statement

      grids[k] = newGrid # we need to change the stored Grid since we now have a true subgrid
      push!(S,S_)
      push!(SOrigIdx,idx)
      push!(SIsPlain,false)
      patchToSMIdx[k] = length(S)
    end
  end
  @debug "Use $(length(S)) patches"

  @debug "Calculate LUT"
  # now that we have all grids we can calculate the indices within the recoGrid
  xcc, xss = calculateLUT(grids, recoGrid)

  sign = ones(Int, M, numPatches)

  return MultiPatchOperator(S, recoGrid, length(recoGrid), M*numPatches,
             RowToPatch, xcc, xss, sign, numPatches, patchToSMIdx)
end

function calculateLUT(grids, recoGrid)
  xss = Vector{Int}[]
  xcc = Vector{Int}[]
  for k=1:length(grids)
    N = length(grids[k])
    push!(xss, collect(1:N))
    xc = zeros(Int64,N)
    for n=1:N
      xc[n] = posToLinIdx(recoGrid,grids[k][n])
    end
    push!(xcc, xc)
  end
  return xcc, xss
end

function size(FFOp::MultiPatchOperator,i::Int)
  if i==2
    return FFOp.N
  elseif i==1
    return FFOp.M
  else
    error("bounds error")
  end
end

size(FFOp::MultiPatchOperator) = (FFOp.M,FFOp.N)

length(FFOp::MultiPatchOperator) = size(FFOp,1)*size(FFOp,2)

### The following is intended to use the standard kaczmarz method ###

function calculateTraceOfNormalMatrix(Op::MultiPatchOperator, weights)
  if length(Op.S) == 1
    trace = calculateTraceOfNormalMatrix(Op.S[1],weights)
    trace *= Op.nPatches #*prod(Op.PixelSizeSF)/prod(Op.PixelSizeC)
  else
    trace = sum([calculateTraceOfNormalMatrix(S,weights) for S in Op.S])
    #trace *= prod(Op.PixelSizeSF)/prod(Op.PixelSizeC)
  end
  return trace
end

setlambda(::MultiPatchOperator, ::Any) = nothing

function dot_with_matrix_row(Op::MultiPatchOperator, x::AbstractArray{T}, k::Integer) where T
  p = Op.RowToPatch[k]
  xs = Op.xss[p]
  xc = Op.xcc[p]

  j = mod1(k,div(Op.M,Op.nPatches))
  A = Op.S[Op.patchToSMIdx[p]]
  sign = Op.sign[j,Op.patchToSMIdx[p]]

  return dot_with_matrix_row_(A,x,xs,xc,j,sign)
end

function dot_with_matrix_row_(A::AbstractArray{T},x,xs,xc,j,sign) where T
  tmp = zero(T)
  @simd  for i = 1:length(xs)
     @inbounds tmp += sign*A[j,xs[i]]*x[xc[i]]
  end
  tmp
end

function kaczmarz_update!(Op::MultiPatchOperator, x::AbstractArray, k::Integer, beta)
  p = Op.RowToPatch[k]
  xs = Op.xss[p]
  xc = Op.xcc[p]

  j = mod1(k,div(Op.M,Op.nPatches))
  A = Op.S[Op.patchToSMIdx[p]]
  sign = Op.sign[j,Op.patchToSMIdx[p]]

  kaczmarz_update_!(A,x,beta,xs,xc,j,sign)
end

function kaczmarz_update_!(A,x,beta,xs,xc,j,sign)
  @simd for i = 1:length(xs)
    @inbounds x[xc[i]] += beta* conj(sign*A[j,xs[i]])
  end
end

function initkaczmarz(Op::MultiPatchOperator,λ,weights::Vector)
  T = typeof(real(Op.S[1][1]))
  denom = T[] #zeros(T,Op.M)
  rowindex = Int64[] #zeros(Int64,Op.M)

  MSub = div(Op.M,Op.nPatches)

  if length(Op.S) == 1
    for i=1:MSub
      s² = rownorm²(Op.S[1],i)*weights[i]^2
      if s²>0
        for l=1:Op.nPatches
          k = i+MSub*(l-1)
          push!(denom,weights[i]^2/(s²+λ)) #denom[k] = weights[i]^2/(s²+λ)
          push!(rowindex,k) #rowindex[k] = k
        end
      end
    end
  else
    for l=1:Op.nPatches
      for i=1:MSub
        s² = rownorm²(Op.S[Op.patchToSMIdx[l]],i)*weights[i]^2
        if s²>0
          k = i+MSub*(l-1)
          push!(denom,weights[i]^2/(s²+λ)) #denom[k] = weights[i]^2/(s²+λ)
          push!(rowindex,k) #rowindex[k] = k
        end
      end
    end
  end

  denom, rowindex
end


export
    BP,
    infer

"""
    abstract type BP <: InstantAlgorithm

Belief Propagation algorithm
"""
abstract type BP <: InstantAlgorithm end

function answer(::Marginal, ::BP, runtime::Runtime, instance::VariableInstance, bounds = false)
    if bounds
        error("BP cannot provide bounded answers")
    end
    return get_belief(runtime, instance)
end

function infer(algorithm::BP, runtime::InstantRuntime,
    evidence::Dict{Symbol, Score} = Dict{Symbol, Score}(),
    interventions::Dict{Symbol, Dist} = Dict{Symbol, Dist}(),
    placeholder_beliefs = get_placeholder_beliefs(runtime))
    net = get_network(runtime)
    ensure_all!(runtime)
    order = topsort(get_initial_graph(net))
    set_ranges!(runtime, evidence, algorithm.default_range_size, 1, order, placeholder_beliefs)
    for ph in get_placeholders(net)
        if !(ph.name in keys(placeholder_beliefs))
            error("Placeholder ", ph.name, " does not have a belief")
        end
        pi = placeholder_beliefs[ph.name]
        for ch in get_children(net, ph)
            set_message!(runtime, ph, ch, :pi_message, pi)
        end
    end
    for (n,e) in evidence
        v = get_node(net, n)
        inst = current_instance(runtime, v)
        post_evidence!(runtime, inst, e)
    end
    for (n,i) in interventions
        v = get_node(net, n)
        inst = current_instance(runtime, v)
        post_intervention!(runtime, inst, i)
    end
    run_bp(algorithm, runtime)
end

# Run a backward step of BP for a single variable.
function _backstep(runtime, var, ranges)
    network = get_network(runtime)
    inst = current_instance(runtime, var)
    evidence = get_evidence(runtime, inst)
    sf = get_sfunc(inst)
    chs = get_children(get_network(runtime), var)
    msgs = collect_messages(runtime, chs, var, :lambda_message)
    O = output_type(sf)
    range::Vector{<:O} = [convert(output_type(sf), r) for r in ranges[var.name]]

    if isempty(msgs)
        ilams = Score{output_type(sf)}[]
        if !isnothing(evidence)
            push!(ilams, evidence)
        end
        lam = compute_lambda(sf, range, ilams)
    else
        incoming_lams = Score[]
        for m in msgs
            m1::Score = m
            push!(incoming_lams, m1)
        end
        if !isnothing(evidence)
            push!(incoming_lams, evidence)
        end
        lam = compute_lambda(sf, range, incoming_lams)
    end

    set_value!(runtime, inst, :lambda, lam)
    pars = get_parents(network, var)
    incoming_pis::Vector{Dist} = collect_messages(runtime, pars, var, :pi_message)
    parranges = [ranges[p.name] for p in pars]
    outgoing_lams = outgoing_lambdas(sf, lam, range, tuple(parranges...), tuple(incoming_pis...))
    distribute_messages!(runtime, var, pars, :lambda_message, outgoing_lams)
end

# Run a forward step of BP for a single variable
# The procedure is slightly different for the first pass, so use the firstpass flag
function _forwardstep(runtime, var, ranges, firstpass)
    network = get_network(runtime)
    inst = current_instance(runtime, var)
    sf = get_sfunc(inst)
    pars = get_parents(network, var)
    O = output_type(sf)
    range = convert(Vector{O}, ranges[var.name])
    parranges = [ranges[p.name] for p in pars]
    incoming_pis = tuple(collect_messages(runtime, pars, var, :pi_message)...)
    if has_intervention(runtime, inst)
        intervention = get_intervention(runtime, inst)
        pi = intervention
    elseif firstpass && has_evidence(runtime, inst)
        # In the first pass, evidence is treated like lambda messages from elsewhere
        evidence = get_evidence(runtime, inst)
        pi = Cat(range, [get_score(evidence, x) for x in range])
    else
        pi = compute_pi(sf, range, tuple(parranges...), incoming_pis)
    end
    set_value!(runtime, inst, :pi, pi)

    chs = get_children(network, var)
    if firstpass
        outpis = fill(pi, length(chs))
    else
        lam = get_value(runtime, inst, :lambda)
        bel = compute_bel(sf, range, pi, lam)
        post_belief!(runtime, inst, bel)
        incoming_lams::Vector{Score} = collect_messages(runtime, chs, var, :lambda_message)
        outpis = outgoing_pis(sf, range, bel, incoming_lams)
    end
    distribute_messages!(runtime, var, chs, :pi_message, outpis)
end


module Thyra
using CxxWrap, MPI
import ..libjltrilinos
import ..CxxUnion
import ..Teuchos
import ..Tpetra

export LinearOpWithSolve

@wrapmodule(libjltrilinos, :register_thyra)

struct LinearOpWithSolve{ST,LT,GT,NT}
  lows::CxxWrap.SmartPointer{Thyra.LinearOpWithSolveBase{ST}}
  domainmap::CxxWrap.SmartPointer{Tpetra.Map{LT,GT,NT}}
  rangespace::CxxWrap.SmartPointer{Thyra.TpetraVectorSpace{ST,LT,GT,NT}}
  domainspace::CxxWrap.SmartPointer{Thyra.TpetraVectorSpace{ST,LT,GT,NT}}

  function LinearOpWithSolve{ST,LT,GT,NT}(A::CxxUnion{Tpetra.CrsMatrix{ST,LT,GT,NT}}, parameters::CxxUnion{Teuchos.ParameterList}, verbosity::Teuchos.EVerbosityLevel) where {ST,LT,GT,NT}
    lows_factory = Thyra.BelosLinearOpWithSolveFactory(ST)
    Thyra.setParameterList(lows_factory, parameters)
    Thyra.setVerbLevel(lows_factory, verbosity)

    lows = Thyra.createOp(lows_factory)
    domainmap = Tpetra.getDomainMap(A)
    rangespace = Thyra.tpetraVectorSpace(Tpetra.getRangeMap(A))
    domainspace = Thyra.tpetraVectorSpace(domainmap)
    A_thyra = Thyra.tpetraLinearOp(rangespace, domainspace, A)
    Thyra.initializeOp(lows_factory, A_thyra, lows)

    return new{ST,LT,GT,NT}(lows, domainmap, rangespace, domainspace)
  end
end

function LinearOpWithSolve(A::CxxUnion{Tpetra.CrsMatrix{ST,LT,GT,NT}}, parameters::CxxUnion{Teuchos.ParameterList}=Teuchos.ParameterList(), verbosity::Teuchos.EVerbosityLevel=Teuchos.VERB_NONE) where {ST,LT,GT,NT}
  return LinearOpWithSolve{ST,LT,GT,NT}(A,parameters,verbosity)
end

import Base: \

function \(A::LinearOpWithSolve{ST,LT,GT,NT}, b::CxxUnion{Tpetra.Vector{ST,LT,GT,NT}}) where {ST,LT,GT,NT}
  x = Tpetra.Vector(A.domainmap)
  x_th = Thyra.tpetraVector(A.domainspace, x)
  b_th = Thyra.tpetraVector(A.rangespace, b)
  status = Thyra.solve(A.lows, Thyra.NOTRANS, b_th, x_th)
  return x
end

end # module