import .Tracker
using .Tracker: TrackedReal, TrackedVector, TrackedMatrix, TrackedArray, @grad

_eps(::Type{<:TrackedReal{T}}) where {T} = eps(T)

# AD implementations
function jacobian(
    b::Union{<:ADBijector{<:TrackerAD}, Inversed{<:ADBijector{<:TrackerAD}}},
    x::Real
)
    return Tracker.data(Tracker.gradient(b, x)[1])
end
function jacobian(
    b::Union{<:ADBijector{<:TrackerAD}, Inversed{<:ADBijector{<:TrackerAD}}},
    x::AbstractVector{<:Real}
)
    # We extract `data` so that we don't return a `Tracked` type
    return Tracker.data(Tracker.jacobian(b, x))
end

# implementations for Shift bijector
function _logabsdetjac_shift(a::TrackedReal, x::Real, ::Val{0})
    return Tracker.param(_logabsdetjac_shift(Tracker.data(a), Tracker.data(x), Val(0)))
end
function _logabsdetjac_shift(a::TrackedReal, x::AbstractVector{<:Real}, ::Val{0})
    return Tracker.param(_logabsdetjac_shift(Tracker.data(a), Tracker.data(x), Val(0)))
end
function _logabsdetjac_shift(
    a::Union{TrackedReal, TrackedVector{<:Real}},
    x::AbstractVector{<:Real},
    ::Val{1}
)
    return Tracker.param(_logabsdetjac_shift(Tracker.data(a), Tracker.data(x), Val(1)))
end
function _logabsdetjac_shift(
    a::Union{TrackedReal, TrackedVector{<:Real}},
    x::AbstractMatrix{<:Real},
    ::Val{1}
)
    return Tracker.param(_logabsdetjac_shift(Tracker.data(a), Tracker.data(x), Val(1)))
end

# implementations for Scale bijector
# Adjoints for 0-dim and 1-dim `Scale` using `Real`
function _logabsdetjac_scale(a::TrackedReal, x::Real, ::Val{0})
    return Tracker.track(_logabsdetjac_scale, a, Tracker.data(x), Val(0))
end
@grad function _logabsdetjac_scale(a::Real, x::Real, ::Val{0})
    return _logabsdetjac_scale(Tracker.data(a), Tracker.data(x), Val(0)), Δ -> (inv(Tracker.data(a)) .* Δ, nothing, nothing)
end
# Need to treat `AbstractVector` and `AbstractMatrix` separately due to ambiguity errors
function _logabsdetjac_scale(a::TrackedReal, x::AbstractVector, ::Val{0})
    return Tracker.track(_logabsdetjac_scale, a, Tracker.data(x), Val(0))
end
@grad function _logabsdetjac_scale(a::Real, x::AbstractVector, ::Val{0})
    da = Tracker.data(a)
    J = fill(inv.(da), length(x))
    return _logabsdetjac_scale(da, Tracker.data(x), Val(0)), Δ -> (transpose(J) * Δ, nothing, nothing)
end
function _logabsdetjac_scale(a::TrackedReal, x::AbstractMatrix, ::Val{0})
    return Tracker.track(_logabsdetjac_scale, a, Tracker.data(x), Val(0))
end
@grad function _logabsdetjac_scale(a::Real, x::AbstractMatrix, ::Val{0})
    da = Tracker.data(a)
    J = fill(size(x, 1) / da, size(x, 2))
    return _logabsdetjac_scale(da, Tracker.data(x), Val(0)), Δ -> (transpose(J) * Δ, nothing, nothing)
end
# adjoints for 1-dim and 2-dim `Scale` using `AbstractVector`
function _logabsdetjac_scale(a::TrackedVector, x::AbstractVector, ::Val{1})
    return Tracker.track(_logabsdetjac_scale, a, Tracker.data(x), Val(1))
end
@grad function _logabsdetjac_scale(a::TrackedVector, x::AbstractVector, ::Val{1})
    # ∂ᵢ (∑ⱼ log|aⱼ|) = ∑ⱼ δᵢⱼ ∂ᵢ log|aⱼ|
    #                 = ∂ᵢ log |aᵢ|
    #                 = (1 / aᵢ) ∂ᵢ aᵢ
    #                 = (1 / aᵢ)
    da = Tracker.data(a)
    J = inv.(da)
    return _logabsdetjac_scale(da, Tracker.data(x), Val(1)), Δ -> (J .* Δ, nothing, nothing)
end
function _logabsdetjac_scale(a::TrackedVector, x::AbstractMatrix, ::Val{1})
    return Tracker.track(_logabsdetjac_scale, a, Tracker.data(x), Val(1))
end
@grad function _logabsdetjac_scale(a::TrackedVector, x::AbstractMatrix, ::Val{1})
    da = Tracker.data(a)
    Jᵀ = repeat(inv.(da), 1, size(x, 2))
    return _logabsdetjac_scale(da, Tracker.data(x), Val(1)), Δ -> (Jᵀ * Δ, nothing, nothing)
end
# TODO: implement analytical gradient for scaling a vector using a matrix
# function _logabsdetjac_scale(a::TrackedMatrix, x::AbstractVector, ::Val{1})
#     Tracker.track(_logabsdetjac_scale, a, Tracker.data(x), Val{1})
# end
# @grad function _logabsdetjac_scale(a::TrackedMatrix, x::AbstractVector, ::Val{1})
#     throw
# end

# implementations for Stacked bijector
function logabsdetjac(b::Stacked, x::TrackedMatrix{<:Real})
    return Tracker.collect([logabsdetjac(b, x[:, i]) for i = 1:size(x, 2)])
end
# TODO: implement custom adjoint since we can exploit block-diagonal nature of `Stacked`
function (sb::Stacked)(x::TrackedMatrix{<:Real})
    init = reshape(sb(x[:, 1]), :, 1)
    return Tracker.collect(mapreduce(i -> sb(x[:, i]), hcat, 2:size(x, 2); init = init))
end


#!/bin/bash
#=
JULIA="${JULIA:-julia}"
JULIA_CMD="${JULIA_CMD:-$JULIA --color=yes --startup-file=no}"
# below gets the directory name of the script, even if there is a symlink involved
# from https://stackoverflow.com/questions/59895/get-the-source-directory-of-a-bash-script-from-within-the-script-itself
SOURCE="${BASH_SOURCE[0]}"
while [ -h "$SOURCE" ]; do # resolve $SOURCE until the file is no longer a symlink
  DIR="$( cd -P "$( dirname "$SOURCE" )" >/dev/null 2>&1 && pwd )"
  SOURCE="$(readlink "$SOURCE")"
  [[ $SOURCE != /* ]] && SOURCE="$DIR/$SOURCE" # if $SOURCE was a relative symlink, we need to resolve it relative to the path where the symlink file was located
done
DIR="$( cd -P "$( dirname "$SOURCE" )" >/dev/null 2>&1 && pwd )"
export JULIA_PROJECT=$DIR/..
export JULIA_LOAD_PATH=@:@stdlib  # exclude default environment
exec $JULIA_CMD -e 'include(popfirst!(ARGS))' "$SOURCE" "$@"
=#
using ArgParse
s = ArgParseSettings()
@add_arg_table s begin
    "basisfile"
        help = "name of the HDF5 containing basis and creation info"
        required = true
    "--outputfile","-o"
        help = "the output pdf file path, otherwise it will make up a name"

end

delete!(ENV, "PYTHONPATH")
using PyCall
using PyPlot
using HDF5
using Pope
using LinearAlgebra
using Statistics
parsed_args = parse_args(ARGS, s)
if parsed_args["outputfile"]==nothing
    parsed_args["outputfile"] = Pope.outputname(parsed_args["basisfile"],"plots","pdf")
end
display(parsed_args);println()
pdf = pyimport("matplotlib.backends.backend_pdf")

function titleme(basisinfo)
    title("Chan $(basisinfo.channel_number)
    Noise file: $(basename(basisinfo.noise_model_file))
    Pulse file: $(basename(basisinfo.pulse_file))")
end

function plot_model_and_residuals(basisinfo)
    figure(figsize=(10,10))
    subplot(311)
    Nsamp, Nb = size(basisinfo.svdbasis.basis)
    artists=plot(basisinfo.svdbasis.basis)
    xlabel("Sample number")
    titleme(basisinfo)
    legend(artists, basisinfo.singular_values,loc="right",title="Singular values")

    subplot(312)
    for i = 1:Nb
        p = basisinfo.svdbasis.projectors[i,:]
        plot(p ./ norm(p))
    end
    xlabel("Sample number")
    ylabel("Projectors (scaled to norm=1)")

    subplot(313)
    noise_std = √(basisinfo.svdbasis.noise_result.autocorr[1])
    normalized_sorted_residuals = sort(basisinfo.std_residuals)/noise_std
    Ntrain = length(normalized_sorted_residuals)
    a=normalized_sorted_residuals[1]
    b=max(1.5*a, normalized_sorted_residuals[ceil(Int, Ntrain*0.9)])
    plot(normalized_sorted_residuals, range(0, stop=1, length=length(normalized_sorted_residuals)))
    xlim(a,b)
    ylim(0,1)
    xlabel("(Residual std)/(Noise std)")
    ylabel("Frac. training pulses w/ lower residual")
    grid(true)
end

function plot_example_pulses(basisinfo)
    for i = 1:ceil(Int,length(basisinfo.percentiles_of_sample_pulses)/9)
        r = (1:9) .+ (i-1)*9
        figure(figsize=(10,10))
        subplot(311)
        artists=plot(basisinfo.example_pulses[:,r])
        legend(artists,basisinfo.percentiles_of_sample_pulses[r],loc="right",title="Percentile of residual std")
        xlabel("Sample number")
        ylabel("Measured pulses")
        titleme(basisinfo)

        subplot(312)
        model_pulses = basisinfo.svdbasis.basis*basisinfo.svdbasis.projectors*basisinfo.example_pulses[:,r]
        residuals = basisinfo.example_pulses[:,r]-model_pulses
        noise_std = √(basisinfo.svdbasis.noise_result.autocorr[1])
        artists=plot(residuals)
        std_residual_normalized = Float32.(std(residuals, dims=1)'[:]/noise_std)
        std_residual_normalized_from_training = basisinfo.std_residuals_of_example_pulses/noise_std
        legend(artists, std_residual_normalized, loc="right", title="std dev/noise std dev")
        xlabel("Sample number")
        ylabel("Residuals (measured-model)")
        ylim(-7*median(std_residual_normalized)*noise_std,7*median(std_residual_normalized)*noise_std)

        subplot(313)
        artists = plot(model_pulses)
        xlabel("Sample number")
        ylabel("Model pulses")
    end
end

function write_all_plots_to_pdf(pdffile)
    for i in plt.get_fignums()
        figure(i)
        pdffile.savefig(i)
        close(i)
    end
end

function main(parsed_args)
    h5 = h5open(parsed_args["basisfile"],"r")
    pdffile = pdf.PdfPages(parsed_args["outputfile"])
    channelnames = string.(sort(parse.(Int,names(h5))))
    for name in channelnames
        println("plotting channel $name")
        channel = parse(Int,name)
        basisinfo = Pope.hdf5load(Pope.SVDBasisWithCreationInfo,h5["$name"])
        plot_model_and_residuals(basisinfo)
        plot_example_pulses(basisinfo)
        write_all_plots_to_pdf(pdffile)
    end
    close(h5)
    pdffile.close()
    plt.close("all")
end

main(parsed_args)


# © Przemyslaw Szufel, 2018, run under Julia 1.0.0

struct ModelConfig
    # computational grid size
    GRID_SIZE :: Tuple{Int,Int}
    #data matrix size on a single process (all matrices are squares)
    dim_i::Int
    dim_j::Int
    #probabilities for empty, red and blue cells
    probs_w::Vector{Float64}
    #minimal required number of neighbours
    min_neighbours::Int
    slaves::Vector{Int64}
    master_id::Int64
end

mutable struct NodeData
    grid_i :: Int
    grid_j :: Int
    field :: Matrix{Int8}
    free_cells::Vector{Tuple{Int,Int}}
    rng::MersenneTwister
end

struct BorderData
    left :: Vector{Int8}
    right :: Vector{Int8}
    top :: Vector{Int8}
    bottom :: Vector{Int8}
end

struct CornerData
    topleft :: Int8
    topright :: Int8
    bottomleft :: Int8
    bottomright :: Int8
end

#The below const statements could be an alternative to barrier functions used
#const modelC  = ModelConfig((0,0),0,0,Vector{Float64}(),0)
#const node = NodeData(0,0,Matrix{Int8}(0,0),Vector{Tuple{Int,Int}}(),MersenneTwister(0),(0,0))
#const worker_grid = Matrix{Int}(0,0)
function get_field_interior():: Matrix{Int8}
    dim_i = modelC.dim_i
    dim_j = modelC.dim_j
    field = zeros(Int8,(dim_i,dim_j))
    for i = (1+1):(dim_i+1), j = (1+1):(dim_j+1)
        field[i-1,j-1]=node.field[i,j]
    end
    return field
end

function sample(rng::AbstractRNG, weights::Vector{Float64})::Int
    t = rand(rng) * sum(weights)
    n = length(weights)
    i = 1
    cw = weights[1]
    while cw < t && i < n
        i += 1
        @inbounds cw += weights[i]
    end
    return i
end

function init_node()::Bool
    dim_i = modelC.dim_i
    dim_j = modelC.dim_j
    values = Int8(0):Int8(length(modelC.probs_w)-1)
    node.field = zeros(Int8,(1+dim_i+1,1+dim_j+1))
    node.free_cells=Vector{Tuple{Int,Int}}()
    for i = (1+1):(dim_i+1), j = (1+1):(dim_j+1)
       val = values[sample(node.rng,modelC.probs_w)]
       if val == 0
           push!(node.free_cells,(i,j))
       end
       node.field[i,j] = val
    end
    return true
end

function get_border()::BorderData
    return get_border(modelC, node)
end
function get_border(modelC :: ModelConfig, node ::NodeData)::BorderData
    dim_i = modelC.dim_i
    dim_j = modelC.dim_j
    #Please note that 4 bytes are redundant, however this increases code readibility
    return BorderData(
        node.field[(1+1):(1+dim_i),1+1],
        node.field[(1+1):(1+dim_i),1+dim_j],
        node.field[1+1,(1+1):(1+dim_j)],
        node.field[1+dim_i,(1+1):(1+dim_j)] )
end

function get_border_top()::Vector{Int8}
    return get_border_top(modelC, node)
end
function get_border_bottom()::Vector{Int8}
    return get_border_bottom(modelC, node)
end
function get_border_left()::Vector{Int8}
    return get_border_left(modelC, node)
end
function get_border_right()::Vector{Int8}
    return get_border_right(modelC, node)
end

function get_border_top(modelC :: ModelConfig, node ::NodeData)::Vector{Int8}
    return node.field[(1+1),(1+1):(1+modelC.dim_j)]
end
function get_border_bottom(modelC :: ModelConfig, node ::NodeData)::Vector{Int8}
    return node.field[(1+modelC.dim_i),(1+1):(1+modelC.dim_j)]
end
function get_border_left(modelC :: ModelConfig, node ::NodeData)::Vector{Int8}
    return node.field[(1+1):(1+modelC.dim_i),(1+1)]
end
function get_border_right(modelC :: ModelConfig, node ::NodeData)::Vector{Int8}
    return node.field[(1+1):(1+modelC.dim_i),(1+modelC.dim_j)]
end

function get_corner_top_left()::Int8
    return get_corner_top_left(modelC, node)
end
function get_corner_top_right()::Int8
    return get_corner_top_right(modelC, node)
end
function get_corner_bottom_left()::Int8
    return get_corner_bottom_left(modelC, node)
end
function get_corner_bottom_right()::Int8
    return get_corner_bottom_right(modelC, node)
end

function get_corner_top_left(modelC :: ModelConfig, node ::NodeData)::Int8
    return node.field[1+1,1+1]
end
function get_corner_top_right(modelC :: ModelConfig, node ::NodeData)::Int8
    return node.field[1+1,1+modelC.dim_j]
end
function get_corner_bottom_left(modelC :: ModelConfig, node ::NodeData)::Int8
    return node.field[1+modelC.dim_i,1+1]
end
function get_corner_bottom_right(modelC :: ModelConfig, node ::NodeData)::Int8
    return node.field[1+modelC.dim_i,1+modelC.dim_j]
end

function step_setborder_nodes!()::Float64
    return step_setborder_nodes!(modelC, node)
end
function step_setborder_nodes!(modelC :: ModelConfig, node ::NodeData)::Bool
    start_t = time()
    dim_i = modelC.dim_i
    dim_j = modelC.dim_j
    #borders
    for j = 1:dim_j
        node.field[1,j+1]=border_data.top[j]
        node.field[1+dim_i+1,j+1]=border_data.bottom[j]
    end
    for i = 1:dim_i
        node.field[i+1,1]=border_data.left[i]
        node.field[i+1,1+dim_j+1]=border_data.right[i]
    end
    node.field[1,1] = corner_data.topleft
    node.field[1,1+dim_j+1] = corner_data.topright
    node.field[1+dim_i+1,1] = corner_data.bottomleft
    node.field[1+dim_i+1,1+dim_j+1] = corner_data.bottomright
    return time()-start_t
end

function step_setborder_nodes_p2p!()::Float64
    return step_setborder_nodes_p2p!(modelC, node, worker_grid)
end
function step_setborder_nodes_p2p!(modelC :: ModelConfig, node ::NodeData, worker_grid::Matrix{Int})::Float64
    start_t = time()
    j_left = node.grid_j-1
    if j_left < 1; j_left = modelC.GRID_SIZE[2]; end
    j_right = node.grid_j+1
    if j_right > modelC.GRID_SIZE[2]; j_right = 1; end
    i_top = node.grid_i-1
    if i_top < 1; i_top = modelC.GRID_SIZE[1]; end
    i_bottom = node.grid_i+1
    if i_bottom > modelC.GRID_SIZE[1]; i_bottom = 1; end
    dim_i = modelC.dim_i
    dim_j = modelC.dim_j
    #borders
    @sync begin
        #top
        @async node.field[1,(1+1):(dim_j+1)] = remotecall_fetch(get_border_bottom,worker_grid[i_top,node.grid_j])
        #bottom
        @async node.field[(1+dim_i+1),(1+1):(dim_j+1)] = remotecall_fetch(get_border_top,worker_grid[i_bottom,node.grid_j])
        #left
        @async node.field[(1+1):(dim_i+1),1] = remotecall_fetch(get_border_right,worker_grid[node.grid_i,j_left])
        #right
        @async node.field[(1+1):(dim_i+1),(1+dim_j+1)] = remotecall_fetch(get_border_left,worker_grid[node.grid_i,j_right])
        #top-left
        @async node.field[1,1] = remotecall_fetch(get_corner_bottom_right,worker_grid[i_top,j_left])
        #top-right
        @async node.field[1,(1+dim_j+1)] = remotecall_fetch(get_corner_bottom_left,worker_grid[i_top,j_right])
        #bottom-left
        @async node.field[(1+dim_i+1),1] = remotecall_fetch(get_corner_top_right,worker_grid[i_bottom,j_left])
        #bottom-right
        @async node.field[(1+dim_i+1),(1+dim_j+1)] = remotecall_fetch(get_corner_top_left,worker_grid[i_bottom,j_right])
    end
    return time()-start_t
end

struct LeaveLocData
    relocate_counts :: Vector{Int}
    free_cells_count :: Int
    elapsed_time :: Float64
end

function step_leave_location!()::LeaveLocData
    return step_leave_location!(modelC, node)::LeaveLocData
end
function step_leave_location!(modelC :: ModelConfig, node ::NodeData)::LeaveLocData
    start_t = time()
    field = node.field
    free_count = length(node.free_cells)
    dim_i = modelC.dim_i
    dim_j = modelC.dim_j
    ds = ((1,0),  (0,-1), (0,1), (-1,0), (1,1),  (1,-1), (-1,1), (-1,-1))
    for i = (1+1):(dim_i+1), j = (1+1):(dim_j+1)
        if field[i,j] > 0
            c = 0
            for (d1, d2) ∈ ds
                c += field[i,j] == field[i+d1,j+d2]
            end
            #print("c=",c," i=",i," j=",j," f=",field[i-1:i+1,j-1:j+1],"\r\n")
            if c < modelC.min_neighbours
                push!(node.free_cells,(i,j))
            end
        end
    end
    relocate_counts::Vector{Int} = zeros(Int,length(modelC.probs_w)-1)
    for (i, j) in node.free_cells[free_count+1:length(node.free_cells)]
        relocate_counts[field[i,j]] +=1
        field[i,j]=0
    end
    return LeaveLocData(relocate_counts,length(node.free_cells),time()-start_t)
end

function step_set_new_location!()::Float64
    return step_set_new_location!(modelC, node)
end
function step_set_new_location!(modelC :: ModelConfig, node ::NodeData)::Float64
    start_t = time()
    shuffle!(node.rng,node.free_cells)
    ix = length(node.free_cells)
    for z = 1:length(new_citizens)
        for a = 1:new_citizens[z]
            node.field[node.free_cells[ix]...]=z
            ix -=1
            pop!(node.free_cells)
        end
    end
    return time()-start_t
end


extended_source_modeling(model::Model, srcData, recGeometry, recData, weights, perturbation, srcnum::UnitRange{Int64}, op::Char, mode::Int64) =
    extended_source_modeling(model, srcData, recGeometry, recData, weights, perturbation, srcnum, op, mode, Options())

function extended_source_modeling(model::Model, srcData, recGeometry, recData, weights, dm, srcnum::UnitRange{Int64}, op::Char, mode::Int64, options)
# extended_source_modeling function for multiple sources. Depending on the operator and mode, this function distributes the sources
# and if applicable the input data amongst the available workers.

    results = judipmap(j ->extended_source_modeling(model, subsample(srcData, j), subsample(recGeometry, j),
                                                    subsample(recData, j), weights, dm, op, mode, subsample(options, j)), srcnum)

    if op=='F' || (op=='J' && mode==1)
        argout1 = vcat(results...)
    elseif op=='J' && mode==-1
        argout1 = sum(results)
    else
        error("operation no defined")
    end
    return argout1
end


@init @require FiniteDifferences="26cc04aa-876d-5657-8c51-4c34ba976000" begin
    using .FiniteDifferences

    function FiniteDifferences.to_vec(X::CollapsedDimArray)
        x_vec, back = to_vec(collapseddim(X))
        s = size(parent(X))
        si = X.si
        sj = X.sj
        ob = X.onebatch
        function CollapsedDimArray_from_vec(x_vec)
            return CollapsedDimArray(reshape(back(x_vec), s), si, sj, ob)
        end
        return x_vec, CollapsedDimArray_from_vec
    end
end

@inline function _sumbatch(ca::CollapsedDimArray)
     return CollapsedDimArray(sum(parent(ca), dims=ntuple(i->i-1+ca.sj, ndims(parent(ca))+1-ca.sj)), ca.si, ca.sj, static(true))
end

using ChainRulesCore
using ChainRulesCore: NoTangent, @thunk
import ChainRulesCore: ProjectTo
function ChainRulesCore.rrule(::Type{CollapsedDimArray}, x, dims, si, sj, onebatch)
    s = size(x)
    function CollapsedDimArray_pullback(Ȳ)
        ∂x = @thunk begin
            tmp = unwrap_collapse(unthunk(Ȳ))
            size(tmp) == s ? tmp : reshape(tmp, s)
        end
        return (NoTangent(), ∂x, NoTangent(), NoTangent(), NoTangent(), NoTangent())
    end
    return CollapsedDimArray(x, dims, si, sj, onebatch), CollapsedDimArray_pullback
end

function ChainRulesCore.rrule(::typeof(parent), x::CollapsedDimArray)
    si, sj, ob = x.si, x.sj, x.onebatch
    function collapsed_parent_pullback(Ȳ)
        ∂x = @thunk CollapsedDimArray(unthunk(Ȳ), si, sj, ob)
        return (NoTangent(), ∂x)
    end
    return parent(x), collapsed_parent_pullback
end

function ChainRulesCore.rrule(::typeof(matmul), A::AbstractVecOrMat, B::AbstractVecOrMat, s)
    Y = matmul(A, B, s)
    function matmul_pullback(Ȳ)
        Ȳ = unthunk(Ȳ)
        Athunk = @thunk matmul(Ȳ, B', s)
        Bthunk = @thunk matmul(A', Ȳ, s)
        sthunk = @thunk sum(reshape(Ȳ, :) .* reshape(Y, :)) * inv(s)
        return (NoTangent(), Athunk, Bthunk, sthunk)
    end
    return Y, matmul_pullback
end

function ProjectTo(ca::CollapsedDimArray)
    dims = size(parent(ca))
    return ProjectTo{CollapsedDimArray}(; dims=dims, si = ca.si, sj = ca.sj, onebatch = ca.onebatch)
end

function (project::ProjectTo{CollapsedDimArray})(dx::AbstractArray)
    dx = unwrap_collapse(dx)
    return CollapsedDimArray(reshape(dx, project.dims), project.si, project.sj, project.onebatch)
end

function (project::ProjectTo{AbstractArray})(dx::CollapsedDimArray)
    dx = reshape(unwrap_collapse(dx), project.axes)
    return project(dx)
end

function ChainRulesCore.rrule(::typeof(matmul), A::AbstractArray, B::AbstractArray, s)
    ProjA = ProjectTo(A)
    ProjB = ProjectTo(B)

    transA, pA = trans(A)
    transB, pB = trans(B)
    a = CollapsedDimArray(pA)
    b = CollapsedDimArray(pB)
    Y = matmul_wrapper(transA, transB, s, a, b)

    function matmul_pullback(Ȳ)
        Ȳ = unthunk(Ȳ)
        Â = trans(transA, a)
        B̂ = trans(transB, b)

        Athunk = @thunk begin
            tmp = matmul(Ȳ, batched_adjoint(B̂), s)
            ProjA(isonebatch(Â) ? _sumbatch(tmp) : tmp)
        end
        Bthunk = @thunk begin
            tmp = matmul(batched_adjoint(Â), Ȳ, s)
            ProjB(isonebatch(B̂) ? _sumbatch(tmp) : tmp)
        end
        sthunk = @thunk sum(reshape(Ȳ, :) .* reshape(unwrap_collapse(Y), :)) * inv(s)
        return (NoTangent(), Athunk, Bthunk, sthunk)
    end
    return Y, matmul_pullback
end


"""
This module provides concrete implementations of [`StochasticStyle`](@ref)s, which
specify the algorithm used by [`lomc!`](@ref) when
performing stochastic matrix-vector multiplication.

# Stochastic styles:

* [`StochasticStyle`](@ref)
* [`IsStochasticInteger`](@ref)
* [`IsDeterministic`](@ref)
* [`IsStochasticWithThreshold`](@ref)
* [`IsDynamicSemistochastic`](@ref)
* [`IsExplosive`](@ref)
* [`StyleUnknown`](@ref)

The offdiagonal spawning is defined in `spawning.jl` and is controlled by setting a
[`SpawningStrategy`](@ref).

The vector compression strategies are defined in `compression.jl` and are controlled by
setting a [`CompressionStrategy`](@ref).

"""
module StochasticStyles

using StaticArrays
using StatsBase
using ..ConsistentRNG
using ..Rimu: MultiScalar, localpart

using ..Interfaces
import ..Interfaces:
    deposit!, diagonal_element, offdiagonals, random_offdiagonal, default_style,
    fciqmc_col!, step_stats, update_dvec!, compress!
export
    StochasticStyle, IsStochasticInteger, IsDeterministic, IsStochasticWithThreshold,
    IsDynamicSemistochastic, StyleUnknown

include("spawning.jl")
include("compression.jl")
include("styles.jl")

end


"""
    energy_balance(object, meteo, constants = Constants())
    energy_balance!(object, meteo, constants = Constants())

Computes the energy balance of one or several components based on the type of the model it was
parameterized with in `object.energy`, and on one or several meteorology time-steps.

At the moment, two models are implemented in the package:

- [`Monteith`](@ref): the model found in Monteith and Unsworth (2013)
- `Missing`: if no computation of the energy balance is needed

# Arguments

- `object`: a [`Component`](@ref) struct ([`AbstractComponentModel`](@ref)), or a Dict/Array of.
- `meteo::Union{AbstractAtmosphere,Weather}`: meteorology structure, see [`Atmosphere`](@ref) or
[`Weather`](@ref)
- `constants = Constants()`: physical constants. See [`Constants`](@ref) for more details

# Note

Some models need input values for some variables. For example [`Monteith`](@ref) requires a
value for `Rₛ`, `d` and `skyFraction`. If you read the models from a file, you can
use [`init_status!`](@ref) (see examples).

# Examples

```julia
# ---Simple example---

meteo = Atmosphere(T = 20.0, Wind = 1.0, P = 101.3, Rh = 0.65)

# Using the model of Monteith and Unsworth (2013) for energy, Farquhar et al. (1980) for
# photosynthesis, and Medlyn et al. (2011) for stomatal conductance:
leaf = LeafModels(energy = Monteith(),
            photosynthesis = Fvcb(),
            stomatal_conductance = Medlyn(0.03, 12.0),
            Rₛ = 13.747, skyFraction = 1.0, PPFD = 1500.0, d = 0.03)

energy_balance(leaf,meteo)

# ---Using several components---

leaf2 = copy(leaf)
leaf2.status.PPFD = 800.0

energy_balance([leaf,leaf2],meteo)

# You can use a Dict if you'd like to keep track of the leaf in the returned DataFrame:
energy_balance(Dict(:leaf1 => leaf, :leaf2 => leaf2), meteo)

# ---Using several meteo time-steps---

w = Weather([Atmosphere(T = 20.0, Wind = 1.0, P = 101.3, Rh = 0.65),
             Atmosphere(T = 25.0, Wind = 1.5, P = 101.3, Rh = 0.55)],
             (site = "Test site",))

energy_balance(leaf, w)

# ---Using several meteo time-steps and several components---

energy_balance(Dict(:leaf1 => leaf, :leaf2 => leaf2), w)

# ---Using a model file---

model = read_model(joinpath(dirname(dirname(pathof(PlantBiophysics))),"test","inputs","models","plant_coffee.yml"))

# An example model file is available here:
# "https://raw.githubusercontent.com/VEZY/PlantBiophysics/main/test/inputs/models/plant_coffee.yml"

# Initialising the mandatory variables:
init_status!(model, Rₛ = 13.747, skyFraction = 1.0, PPFD = 1500.0, Tₗ = 25.0, d = 0.03)

# NB: To know which variables has to be initialised according to the models used, you can use
# `to_initialise(Component)`, *e.g.*:
to_initialise(model["Leaf"])

# Running a simulation for all component types in the same scene:
energy_balance!(model, meteo)

model["Leaf"].status.Rn
model["Leaf"].status.A
model["Leaf"].status.Cᵢ
```

# References

Duursma, R. A., et B. E. Medlyn. 2012. « MAESPA: a model to study interactions between water
limitation, environmental drivers and vegetation function at tree and stand levels, with an
example application to [CO2] × drought interactions ». Geoscientific Model Development 5 (4):
919‑40. https://doi.org/10.5194/gmd-5-919-2012.

Monteith, John L., et Mike H. Unsworth. 2013. « Chapter 13 - Steady-State Heat Balance: (i)
Water Surfaces, Soil, and Vegetation ». In Principles of Environmental Physics (Fourth Edition),
edited by John L. Monteith et Mike H. Unsworth, 217‑47. Boston: Academic Press.

Schymanski, Stanislaus J., et Dani Or. 2017. « Leaf-Scale Experiments Reveal an Important
Omission in the Penman–Monteith Equation ». Hydrology and Earth System Sciences 21 (2): 685‑706.
https://doi.org/10.5194/hess-21-685-2017.

Vezy, Rémi, Mathias Christina, Olivier Roupsard, Yann Nouvellon, Remko Duursma, Belinda Medlyn,
Maxime Soma, et al. 2018. « Measuring and modelling energy partitioning in canopies of varying
complexity using MAESPA model ». Agricultural and Forest Meteorology 253‑254 (printemps): 203‑17.
https://doi.org/10.1016/j.agrformet.2018.02.005.
"""
energy_balance!, energy_balance

function energy_balance!(object::AbstractComponentModel, meteo::AbstractAtmosphere, constants = Constants())
    is_init = is_initialised(object)
    !is_init && error("Some variables must be initialized before simulation")

    net_radiation!(object, meteo, constants)
    return nothing
end

# Same as above but non-mutating
function energy_balance(object::AbstractComponentModel, meteo::AbstractAtmosphere, constants = Constants())
    object_tmp = copy(object)
    energy_balance!(object_tmp, meteo, constants)

    return object_tmp
end

# energy_balance over several objects (e.g. all leaves of a plant) in an Array
function energy_balance!(object::O, meteo::AbstractAtmosphere, constants = Constants()) where O <: AbstractArray{<:AbstractComponentModel}

    for i in values(object)
        energy_balance!(i, meteo, constants)
    end

    return nothing
end

# energy_balance over several objects (e.g. all leaves of a plant) in a kind of Dict.
function energy_balance!(object::O, meteo::AbstractAtmosphere, constants = Constants()) where {O <: AbstractDict{N,<:AbstractComponentModel} where N}

    for (k, v) in object
        energy_balance!(v, meteo, constants)
    end

    return nothing
end

# same as the above but non-mutating
function energy_balance(
    object::O,
    meteo::AbstractAtmosphere,
    constants = Constants()
    ) where O <: Union{AbstractArray{<:AbstractComponentModel},AbstractDict{N,<:AbstractComponentModel} where N}

    # Copy the objects only once before the computation for performance reasons:
    object_tmp = copy(object)

    # Computation:
    energy_balance!(object_tmp, meteo, constants)

    return object_tmp
end

# energy_balance over several meteo time steps (called Weather) and possibly several components:
function energy_balance!(
    object::T,
    meteo::Weather,
    constants = Constants()
    ) where T <: Union{AbstractArray{<:AbstractComponentModel},AbstractDict{N,<:AbstractComponentModel} where N}

    # Check if the meteo data and the status have the same length (or length 1)
    check_status_wheather(object, meteo)

    # Computing for each time-steps:
    for (i, meteo_i) in enumerate(meteo.data)
        # Each object in a time-step:
        for obj in object
            energy_balance!(obj[i], meteo_i, constants)
        end
    end

end

# If we call weather with one component only, put it in an Array and call the function above
function energy_balance!(object::AbstractComponentModel, meteo::Weather, constants = Constants())
    energy_balance!([object], meteo, constants)
end

# energy_balance over several meteo time steps (same as above) but non-mutating
function energy_balance(
    object::T,
    meteo::Weather,
    constants = Constants()
    ) where T <: Union{AbstractComponentModel,AbstractDict{N,<:AbstractComponentModel} where N}

    object_tmp = copy(object)

    energy_balance!(object_tmp, meteo, constants)

    return object_tmp
end


# main file for computing dRdq transposed products

import PDESolver.evalResidual_revq

function evalResidual_revq(mesh::AbstractMesh, sbp::AbstractOperator, eqn::EulerData,
                     opts::Dict, t::Number=0.0)

  time = eqn.params.time
  eqn.params.t = t  # record t to params
  myrank = mesh.myrank

  # parallel data exchange must be started before this function is called
  # (both q and res_bar)

  # Forward sweep
  time.t_dataprep += @elapsed dataPrep_for_revq(mesh, sbp, eqn, opts)

  time.t_volume += @elapsed if opts["addVolumeIntegrals"]
    evalVolumeIntegrals_revq(mesh, sbp, eqn, opts)
  end

  if opts["use_GLS"]
    error("GLS not supported for revq product")
  end

  time.t_bndry += @elapsed if opts["addBoundaryIntegrals"]
    evalBoundaryIntegrals_revq(mesh, sbp, eqn, opts)
  end

  time.t_stab += @elapsed if opts["addStabilization"]
     addStabilization_revq(mesh, sbp, eqn, opts)
  end

  time.t_face += @elapsed if mesh.isDG && opts["addFaceIntegrals"]
    evalFaceIntegrals_revq(mesh, sbp, eqn, opts)
  end

  time.t_source += @elapsed evalSourceTerm_revq(mesh, sbp, eqn, opts)

  time.t_sharedface += @elapsed if mesh.commsize > 1
    evalSharedFaceIntegrals_revq(mesh, sbp, eqn, opts)
  end

  time.t_shock += @elapsed if opts["addShockCapturing"]
    evalShockCapturing_revq(mesh, sbp, eqn, opts)
  end

  # apply inverse mass matrix to eqn.res, necessary for CN
  if opts["use_Minv"]
    error("use_Minv not supported for revq product")
    applyMassMatrixInverse3D(mesh, sbp, eqn, opts, eqn.res)
  end

  time.t_dataprep += @elapsed dataPrep_revq(mesh, sbp, eqn, opts)

  return nothing

end  # end evalResidual

"""
  Dataprep-type function that should be called at the beginning of a reverse
  mode product, rather than the regular `dataPrep` function.
"""
function dataPrep_for_revq(mesh::AbstractMesh{Tmsh}, sbp::AbstractOperator,
                   eqn::AbstractEulerData{Tsol, Tres}, opts) where {Tmsh, Tsol, Tres}

  # apply filtering to input
  if eqn.params.use_filter
    applyFilter(mesh, sbp, eqn, eqn.q, opts)
  end

  getAuxVars(mesh, eqn)

  if opts["use_staggered_grid"]
    error("staggered grid not supported for revq product")
    aux_vars = zeros(Tres, 1, mesh.mesh2.numNodesPerElement)
    for i=1:mesh.numEl
      qs = ro_sview(eqn.q, :, :, i)
      qf = sview(eqn.q_flux, :, :, i)
      interpolateElementStaggered(eqn.params, mesh, qs, aux_vars, qf)
    end
  end

  # only eqn.q_bndry is required, all other methods are no-precompute
  if mesh.isDG
    interpolateBoundary(mesh, sbp, eqn, opts, eqn.q, eqn.q_bndry, eqn.aux_vars_bndry)
  end

  if eqn.params.use_edgestab
    stabscale(mesh, sbp, eqn)
  end

  return nothing
end

"""

Reverse mode of dataPrep w.r.t mesh metrics

* mesh  : a mesh object
* sbp   : SBP operator object
* eqn   : an EulerData object
* opts  : options dictionary

"""
function dataPrep_revq(mesh::AbstractMesh{Tmsh}, sbp::AbstractOperator,
                   eqn::AbstractEulerData{Tsol, Tres}, opts) where {Tmsh, Tsol, Tres}

  # nothing to do here

  return nothing
end

"""

Reverse mode of evalVolumeIntegrals with respect to the mesh metrics ∂ξ/∂x

**Arguments**

* mesh  : a mesh object
* sbp   : SBP operator object
* eqn   : an EulerData object
* opts  : options dictionary

"""
function evalVolumeIntegrals_revq(mesh::AbstractMesh{Tmsh},
                             sbp::AbstractOperator, eqn::EulerData{Tsol, Tres, Tdim}, opts) where {Tmsh,  Tsol, Tres, Tdim}

  integral_type = opts["volume_integral_type"]
  if integral_type == 1
    if opts["Q_transpose"] == true
      calcVolumeIntegrals_nopre_revq(mesh, sbp, eqn, opts)
    else
      error("Q_transpose = false not supported for revq product")
    end  # end if
  elseif integral_type == 2
    calcVolumeIntegralsSplitForm_revq(mesh, sbp, eqn, opts, eqn.volume_flux_func,
                                 eqn.volume_flux_func_revq)
  else
    throw(ErrorException("Unsupported volume integral type = $integral_type"))
  end

  return nothing
end  # end evalVolumeIntegrals

"""
### EulerEquationMod.evalBoundaryIntegrals_revq

Reverse mode of evalBoundaryIntegrals with respect to the mesh metrics ∂ξ/∂x

**Arguments**

* mesh  : a mesh object
* sbp   : SBP operator object
* eqn   : an EulerData object
* opts  : options dictionary

"""
function evalBoundaryIntegrals_revq(mesh::AbstractMesh{Tmsh},
                               sbp::AbstractOperator, eqn::EulerData{Tsol, Tres, Tdim}, opts) where {Tmsh, Tsol, Tres, Tdim}

  getBCFluxes_revq(mesh, sbp, eqn, opts)

  return nothing

end  # end evalBoundaryIntegrals

function addStabilization_revq(mesh::AbstractMesh, sbp::AbstractOperator,
                               eqn::EulerData, opts)

  if opts["use_lps"]
    applyLPStab_revq(mesh, sbp, eqn, opts)
  end

  return nothing
end

"""

Reverse mode of evalFaceIntegrals with respect to the mesh metrics ∂ξ/∂x

**Arguments**

* mesh  : a mesh object
* sbp   : SBP operator object
* eqn   : an EulerData object
* opts  : options dictionary

"""
function evalFaceIntegrals_revq(mesh::AbstractDGMesh{Tmsh},
                           sbp::AbstractOperator, eqn::EulerData{Tsol}, opts) where {Tmsh, Tsol}

  face_integral_type = opts["face_integral_type"]
  if face_integral_type == 1
    calcFaceIntegral_nopre_revq(mesh, sbp, eqn, opts, eqn.flux_func_revq,
                                mesh.interfaces)

  elseif face_integral_type == 2

    getFaceElementIntegral_revq(mesh, sbp, eqn, eqn.face_element_integral_func,
                           eqn.flux_func_revq, mesh.sbpface, mesh.interfaces)

  else
    throw(ErrorException("Unsupported face integral type = $face_integral_type"))
  end

  # do some output here?
  return nothing
end

"""

  Reverse mode evalSharedFaceIntegrals with respect to q.  Note that this
  uses the parallel approach of [`finishExchangeData_rev2`](@ref), so the
  functions called here are not precisely the reverse mode of the primal
  functions.

"""
function evalSharedFaceIntegrals_revq(mesh::AbstractDGMesh, sbp, eqn, opts)

  if mesh.npeers == 0
    return nothing
  end

  if getParallelData(eqn.shared_data) != PARALLEL_DATA_ELEMENT ||
     getParallelData(eqn.shared_data_bar) != PARALLEL_DATA_ELEMENT

    error("""parallel data setting must be PARALLEL_DATA_ELEMENT """)
  end

  face_integral_type = opts["face_integral_type"]
  if face_integral_type == 1
    finishExchangeData_rev2(mesh, sbp, eqn, opts, eqn.shared_data,
                      eqn.shared_data_bar, calcSharedFaceIntegrals_element_revq)
  elseif face_integral_type == 2
    finishExchangeData_rev2(mesh, sbp, eqn, opts, eqn.shared_data,
                      eqn.shared_data_bar, calcSharedFaceElementIntegrals_element_revq)
  else
    throw(ErrorException("unsupported face integral type = $face_integral_type"))
  end

  return nothing
end

"""

Reverse mode of evalSourceTerm w.r.t mesh metrics.

*This function has not been written properly. The developer must update this
documentation as and when this code is developed*

"""
function evalSourceTerm_revq(mesh::AbstractMesh{Tmsh},
                     sbp::AbstractOperator, eqn::EulerData{Tsol, Tres, Tdim},
                     opts) where {Tmsh, Tsol, Tres, Tdim}

  # placeholder for multiple source term functionality (similar to how
  # boundary conditions are done)
  if opts["use_src_term"]
    error("source terms not supported for revq product")
  end

  return nothing
end  # end function

"""

  Reverse mode of `evalShockCapturing`

  **Inputs**:

   * mesh : Abstract mesh type
   * sbp  : Summation-by-parts operator
   * eqn  : Euler equation object
   * opts : options dictonary
"""
function evalShockCapturing_revq(mesh::AbstractMesh{Tmsh},
                     sbp::AbstractOperator, eqn::EulerData{Tsol, Tres, Tdim},
                     opts) where {Tmsh, Tsol, Tres, Tdim}

  # currently there is only one choice for the shock capturing scheme.
  # If there are more, need something more sophisiticated to choose.
  # Perhaps add an abstract typed field to eqn containing the structs

  capture = eqn.shock_capturing
  applyShockCapturing_revq(mesh, sbp, eqn, opts, capture)

  return nothing
end


#===============================================================================

    amp_analysis.jl

    amplification analysis

    issue:
    the code assumes only 1 step/ramp because the current data format
    does not allow us to break the fluorescence data down by step_id/ramp_id

===============================================================================#

import DataStructures.OrderedDict
import StaticArrays: SVector, SMatrix
import Memento: debug, warn
using Ipopt

#===============================================================================
    function definitions >>
===============================================================================#

"Analyse amplification data via calls to `calibrate` and `get_fit_results`."
function amp_analysis(i ::AmpInput) # ; asrp ::AmpStepRampProperties)
    debug(logger, "at amp_analysis()")
    #
    ## deconvolute and normalize
    calibration_results =
        calibrate(i, i.calibration_data, i.calibration_args, i.raw, array)
    #
    ## initialize output
    o = AmpOutput(
        Val{i.amp_output},
        i,
        calibration_results...,
        # cq_method,
        default_ct_fluos(i))
    #
    ## fit amplification models and report results
    if i.num_cycles <= 2
		num_cycles = i.num_cycles
        warn(logger, "number of cycles $num_cycles <= 2: baseline subtraction " *
            "and Cq calculation will not be performed")
    else ## num_cycles > 2
        baseline_cyc_bounds = check_bl_cyc_bounds(i, DEFAULT_AMP_BL_CYC_BOUNDS)
        try
            set_ct_fluos!(o, i, baseline_cyc_bounds)
        catch e
            warn(logger, "set_ct_fluos catch error: " * sprint(showerror, e))
        end
        try
            set_output_fields!(o, i, get_fit_results(o, i, baseline_cyc_bounds))
        catch e
            warn(logger, "get_fit_results catch error: " * sprint(showerror, e))
        end
        set_qt_fluos!(o, i)
        set_report_cq!(o, i)
    end ## if
    return o
end ## amp_analysis()

#==============================================================================#

"Set the amplification output field `ct_fluos`."
function set_ct_fluos!(
    o                       ::AmpOutput,
    i                       ::AmpInput,
    baseline_cyc_bounds     ::AbstractArray,
)
    debug(logger, "at calc_ct_fluos()")
    (length(o.ct_fluos) > 0) && return nothing
    o.ct_fluos = default_ct_fluos(i)
    (i.cq_method != :ct)     && return nothing
    (i.amp_model != :SFC)    && return nothing
    ## else
    o.ct_fluos =
        map(1:i.num_channels) do channel_i
            fits =
                map(1:i.num_wells) do well_i
                    fluos = o.rbbs_ary3[:, well_i, channel_i]
                    amp_fit_model(
                        Val{SFCModel},
                        AmpCqFluoModelResults,
                        i,
                        fluos,
                        bl_cyc_bounds[well_i, channel_i],
                        DEFAULT_AMP_CT_FLUO_METHOD, ## cq_method
                        NaN_T) ## i.ct_fluo
                end ## do well_i
            fits |>
                their(:quant_status) |>
                find_idc_useful |>
                curry(getindex)(fits) |>
                their(:cq_fluo) |>
                median
        end ## do channel_i
    return nothing ## side effects only
end ## calc_ct_fluos()

## called by calc_ct_fluos() >>

@inline function find_idc_useful(postbl_stata ::AbstractVector)
    idc_useful = find(postbl_stata .== :Optimal)
    (length(idc_useful) > 0) && return idc_useful
    idc_useful = find(postbl_stata .== :UserLimit)
    (length(idc_useful) > 0) && return idc_useful
    return eachindex(postbl_stata)
end ## find_idc_useful()

default_ct_fluos(i ::AmpInput) =
    SVector{i.num_channels, Float_T}(fill(NaN_T, i.num_channels))

#==============================================================================#

"Fit amplification model to data for each well and channel."
function get_fit_results(
    o                       ::AmpOutput,
    i                       ::AmpInput,
    bl_cyc_bounds           ::AbstractArray,
)
    function do_model_fit(wi ::Int_T, ci ::Int_T)
        debug(logger, "at do_model_fit($wi, $ci)")
        if isa(solver, Ipopt.IpoptSolver) && length(prefix) > 0
            ipopt_file = string(join([prefix, ci, wi], '_')) * ".txt"
            push!(solver.options, (:output_file, ipopt_file))
        end
        fluos = o.rbbs_ary3[:, wi, ci]
        amp_fit_model(
            Val{i.amp_model},
            i.amp_model_results,
            i,
            fluos,
            bl_cyc_bounds[wi, ci],
            i.cq_method,
            o.ct_fluos[ci])
    end ## fit model()

    ## << end of function definition nested within set_fit_results!()

    debug(logger, "at set_fit_results!()")
    solver = i.solver
    prefix = i.ipopt_print_prefix
    fit_results =
        SMatrix{i.num_wells, i.num_channels, i.amp_model_results}([
            do_model_fit(wi, ci)
            for wi in 1:i.num_wells, ci in 1:i.num_channels])
end ## set_fit_results!()

#==============================================================================#

"Format the results of the amplification analyses."
@inline function set_output_fields!(
    o                       ::AmpOutput,
    i                       ::AmpInput,
    results                 ::AbstractArray,
)
    debug(logger, "at set_output_fields!()")
    foreach(fieldnames(first(results))) do fieldname
        output_field = getfield(o, fieldname)
        T = output_field |> eltype
        vector_output_field = ndims(output_field) == 3
        if vector_output_field
            setfield!(o, fieldname,
                results |>
                moose(bless(Vector{T}) ∘ field(fieldname), hcat) |>
                morph(:, i.num_wells, i.num_channels))
        else
            setfield!(o, fieldname,
                results |> mold(bless(T) ∘ field(fieldname)) |>
                bless(SMatrix{i.num_wells, i.num_channels, T}))
        end ## if
    end ## next fieldname
    return nothing ## side effects only
end ## set_output_fields!()

#==============================================================================#

"Calculate the amplification analysis fields `qt_fluos` and `max_qt_fluo`, when
the output format is `long`."
function set_qt_fluos!(
    o                       ::AmpLongOutput,
    i                       ::AmpInput,
)
    debug(logger, "at set_qt_fluos!()")
    o.qt_fluos =
        [   quantile(o.blsub_fluos[:, well_i, channel_i], i.qt_prob)
            for well_i in 1:i.num_wells, channel_i in 1:i.num_channels ]
    o.max_qt_fluo = maximum(o.qt_fluos)
    return nothing ## side effects only
end ## set_qt_fluos!()

"Do nothing when the output format is `short`."
set_qt_fluos!(
    o                       ::AmpShortOutput,
    i                       ::AmpInput,
) = nothing

#==============================================================================#

"Call `report_cq!` for each well and channel, when the output format is `long`."
function set_report_cq!(
    o                       ::AmpLongOutput,
    i                       ::AmpInput,
)
    debug(logger, "at set_report_cq!()")
    for well_i in 1:i.num_wells, channel_i in 1:i.num_channels
        report_cq!(i, o, well_i, channel_i)
    end
    return nothing ## side effects only
end ## set_report_cq!()

"Do nothing when the output format is `short`."
function set_report_cq!(
    o                       ::AmpShortOutput,
    i                       ::AmpInput,
)
    debug(logger, "at set_report_cq!() SHORT FORMAT")
    temp_cq=fill(NaN_T,size(o.cq)[1], size(o.cq)[2])

    max_dr1_lb, max_dr2_lb, max_bsf_lb = i.max_dr1_lb, i.max_dr2_lb, i.max_bsf_lb
    qt_fluos_ =
        [   quantile(o.blsub_fluos_flb[:, well_i, channel_i], i.qt_prob)
            for well_i in 1: i.num_wells, channel_i in 1: i.num_channels ]
    max_qt_fluo_ = maximum(qt_fluos_)
    for well_i in 1: i.num_wells, channel_i in 1: i.num_channels
        b_ = o.coefs[1, well_i, channel_i]
        postbl_status = o.quant_status[well_i, channel_i]
        max_dr1= maximum(o.dr1_pred[:, well_i, channel_i])
        max_dr2= maximum(o.dr2_pred[:, well_i, channel_i])
        max_bsf = maximum(o.blsub_fluos_flb[:, well_i, channel_i])
		max_bsf_ = maximum(o.blsub_fluos[(end - Int(floor((size(o.blsub_fluos)[1]) / 4))): end, well_i, channel_i])
        (scaled_max_dr1, scaled_max_dr2, scaled_max_bsf) =
        [max_dr1, max_dr2, max_bsf] ./ max_qt_fluo_

        min_reliable_cyc_ = length(i.baseline_cyc_bounds) == 2 ? i.baseline_cyc_bounds[1]: i.min_reliable_cyc

        if (max_bsf < i.max_bsf_lb || scaled_max_bsf < i.scaled_max_bsf_lb ||
            max_dr2 < max_dr2_lb || max_dr1 < max_dr1_lb  ||
            scaled_max_dr2 < i.scaled_max_dr2_lb ||
            scaled_max_dr1 < i.scaled_max_dr1_lb || o.cq[well_i, channel_i] > i.num_cycles ||
            o.cq[well_i, channel_i] < min_reliable_cyc_ || postbl_status == :Error ||
            b_ > 0)

            why_cq_failed = "Cq for W:$well_i, C:$channel_i failed due to: "
            why_cq_failed *= max_bsf < i.max_bsf_lb ? "max_bsf($max_bsf) < i.max_bsf_lb($(i.max_bsf_lb)) |" : ""
            why_cq_failed *= scaled_max_bsf < i.scaled_max_bsf_lb ? "scaled_max_bsf($scaled_max_bsf) < i.scaled_max_bsf_lb($(i.scaled_max_bsf_lb)) |" : ""
            why_cq_failed *= max_dr2 < max_dr2_lb ? "max_dr2($max_dr2) < max_dr2_lb($max_dr2_lb) |" : ""
            why_cq_failed *= max_dr1 < max_dr1_lb ? "max_dr1 < max_dr1_lb : " * "$max_dr1 < $max_dr1_lb |" : ""
            why_cq_failed *= scaled_max_dr2 < i.scaled_max_dr2_lb ? "scaled_max_dr2($scaled_max_dr2) < i.scaled_max_dr2_lb($(i.scaled_max_dr2_lb)) |" : ""
            why_cq_failed *= scaled_max_dr1 < i.scaled_max_dr1_lb ? "scaled_max_dr1($scaled_max_dr1) < i.scaled_max_dr1_lb($(i.scaled_max_dr1_lb)) |" : ""
            why_cq_failed *= o.cq[well_i, channel_i] > i.num_cycles ? "o.cq[well_i, channel_i]($(o.cq[well_i, channel_i])) > i.num_cycles($(i.num_cycles)) |" : ""
            why_cq_failed *= o.cq[well_i, channel_i] < min_reliable_cyc_ ? "o.cq[well_i, channel_i]($(o.cq[well_i, channel_i])) < min_reliable_cyc_($(min_reliable_cyc_)) |" : ""
            why_cq_failed *= postbl_status == :Error ? "postbl_status == :Error |" : ""
            why_cq_failed *= b_ > 0 ? "b_($b_) > 0 |" : ""
            debug(logger, why_cq_failed)

            o.blsub_fluos[:,well_i, channel_i]=o.blsub_fluos_flb[:,well_i, channel_i]
            o.dr1_pred[:,well_i, channel_i]=o.dr1_pred1[:,well_i, channel_i]
            o.dr2_pred[:,well_i, channel_i]=o.dr2_pred1[:,well_i, channel_i]

        else
            temp_cq[well_i, channel_i] = o.cq[well_i, channel_i]
            debug(logger, "Cq for W:$well_i, C:$channel_i >> $(o.cq[well_i, channel_i])")
        end
    end
    o.cq=convert(typeof(o.cq), temp_cq)
    return nothing
end

"Report amplification output fields relating to the calculation of `cq`."
function report_cq!(
    o                       ::AmpLongOutput,
    i                       ::AmpInput,
    well_i                  ::Int_T,
    channel_i               ::Int_T,
)
    if i.before_128x
        max_dr1_lb, max_dr2_lb, max_bsf_lb = [i.max_dr1_lb, i.max_dr2_lb, i.max_bsf_lb] ./ 128
    else
        max_dr1_lb, max_dr2_lb, max_bsf_lb = i.max_dr1_lb, i.max_dr2_lb, i.max_bsf_lb
    end
    #
    num_cycles = size(o.raw_data, 1)
    (postbl_status, cq_raw, max_dr1, max_dr2) =
        map([ :postbl_status, :cq_raw, :max_dr1, :max_dr2 ]) do fieldname
            fieldname -> getfield(o, fieldname)[well_i, channel_i]
        end
    max_bsf = maximum(o.blsub_fluos[:, well_i, channel_i])
    b_ = o.coefs[1, well_i, channel_i]
    (scaled_max_dr1, scaled_max_dr2, scaled_max_bsf) =
        [max_dr1, max_dr2, max_bsf] ./ o.max_qt_fluo
    why_NaN =
        if postbl_status == :Error
            "postbl_status == :Error"
        elseif b_ > 0
            "b > 0"
        elseif o.cq_method == ct && o.cq_raw == AMP_CT_VAL_DOMAINERROR
            "DomainError when calculating Ct"
        elseif o.cq_raw <= 0.1 || o.cq_raw >= num_cycles
            "cq_raw <= 0.1 || cq_raw >= num_cycles"
        elseif max_dr1 < max_dr1_lb
            "max_dr1 $max_dr1 < max_dr1_lb $max_dr1_lb"
        elseif max_dr2 < max_dr2_lb
            "max_dr2 $max_dr2 < max_dr2_lb $max_dr2_lb"
        elseif max_bsf < max_bsf_lb
            "max_bsf $max_bsf < max_bsf_lb $max_bsf_lb"
        elseif scaled_max_dr1 < i.scaled_max_dr1_lb
            "scaled_max_dr1 $scaled_max_dr1 < scaled_max_dr1_lb $(i.scaled_max_dr1_lb)"
        elseif scaled_max_dr2 < i.scaled_max_dr2_lb
            "scaled_max_dr2 $scaled_max_dr2 < scaled_max_dr2_lb $(i.scaled_max_dr2_lb)"
        elseif scaled_max_bsf < i.scaled_max_bsf_lb
            "scaled_max_bsf $scaled_max_bsf < scaled_max_bsf_lb $(i.scaled_max_bsf_lb)"
        else
            ""
        end ## why_NaN
    (why_NaN != "") && (o.cq[well_i, channel_i] = NaN_T)
    #
    for tup in (
        (:max_bsf,        max_bsf),
        (:scaled_max_dr1, scaled_max_dr1),
        (:scaled_max_dr2, scaled_max_dr2),
        (:scaled_max_bsf, scaled_max_bsf),
        (:why_NaN,        why_NaN))
        getfield(o, tup[1])[well_i, channel_i] = tup[2]
    end
    return nothing ## side effects only
end ## report_cq!

#===============================================================================
    helper functions >>
===============================================================================#

amp_init(i ::AmpInput, x) =
    SMatrix{i.num_wells, i.num_channels, typeof(x)}(
        fill(x, i.num_wells, i.num_channels))

function check_bl_cyc_bounds(
    i               ::AmpInput,
    bl_cyc_bounds   ::Union{Vector{I},Array{Vector{I},2}} where {I <: Int_T},
)
    debug(logger, "at check_bl_cyc_bounds()")
    (i.num_cycles <= 2) && return bl_cyc_bounds
    size_bcb = size(bl_cyc_bounds)
    if size_bcb == (0,) || size_bcb == (2,)
        return amp_init(i, bl_cyc_bounds)
    elseif size_bcb == (i.num_wells, i.num_channels) &&
        eltype(bl_cyc_bounds) <: AbstractVector ## final format of `baseline_cyc_bounds`
            return bl_cyc_bounds
    end
    throw(ArgumentError("`baseline_cyc_bounds` is not in the right format"))
end ## check_bl_cyc_bounds()

## baseline estimation parameters
# kwargs_bl(i ::AmpInput) =
#     Dict{Symbol,Any}(
#         :bl_method          => i.bl_method,
#         :bl_fallback_func   => i.bl_fallback_func,
#         :min_reliable_cyc   => i.min_reliable_cyc,
#     )

## quantitation parameters
# kwargs_quant(i ::AmpInput) =
#     Dict{Symbol,Any}(
#         :cq_method          => i.cq_method,
#         :denser_factor      => i.denser_factor,
#     )

## arguments for process_ad()
# kwargs_ad(i ::AmpInput) =
#     Dict{Symbol,Any}(
#         :ctrl_well_dict     => i.ctrl_well_dict,
#         # :cluster_method     => i.cluster_method,
#         # :norm_l             => i.norm_l,
#         # :encgr              => i.encgr,
#         # :categ_well_vec     => i.categ_well_vec
#     )


# Implementation of main console executable and related I/O.

using YAML
using JLD

import .ArgumentParse
import .Model: Image
import .ParallelRun: infer_box
using .SDSSIO

# read universal and dataset-specific configurations
function read_config(filename)
    data = open(YAML.load, filename)

    # universal settings
    cfg_keys = ["min_radius_pix",
                "num_ais_temperatures",
                "num_ais_samples",
                "num_joint_vi_iters"]
    # dictionary of cfg_keys actually present in configuration file.
    cfg_dict = Dict(k=>data[k] for k in cfg_keys if haskey(data, k))
    cfg = Config(;cfg_dict...)

    # dataset-specific settings
    datasets = Dict()
    for (name, dataset_config) in data["datasets"]
        if name == "sdss"
            basedir = dataset_config["basedir"]
            iostrategy = Symbol(get(dataset_config, "iostrategy", "plain"))
            dirlayout = Symbol(get(dataset_config, "dirlayout", "celeste"))
            slurp = get(dataset_config, "slurp", false)
            compressed = get(dataset_config, "compressed", false)
            datasets[name] = SDSSDataSet(basedir;
                                         dirlayout = dirlayout,
                                         compressed = compressed,
                                         iostrategy = iostrategy,
                                         slurp = slurp)
        else
            error("unrecognized dataset name: $(name)")
        end
    end

    return cfg, datasets
end

"""
Save provided results to a JLD file.
"""
function save_results(outdir, ramin, ramax, decmin, decmax, results)
    fname = @sprintf("%s/celeste-%.4f-%.4f-%.4f-%.4f.jld",
                     outdir, ramin, ramax, decmin, decmax)
    JLD.save(fname, "results", results)
end
save_results(outdir, box::BoundingBox, results) =
    save_results(outdir, box.ramin, box.ramax, box.decmin, box.decmax, results)

# command-line script
function main(console_args::Vector{String}=ARGS)

    parser = ArgumentParse.ArgumentParser(program_name="infer-box")
    ArgumentParse.add_argument(parser, "datasets",
                               help="comma-separated dataset names. Each dataset must appear in config file.")
    ArgumentParse.add_argument(parser, "ramin", arg_type=Float64)
    ArgumentParse.add_argument(parser, "ramax", arg_type=Float64)
    ArgumentParse.add_argument(parser, "decmin", arg_type=Float64)
    ArgumentParse.add_argument(parser, "decmax", arg_type=Float64)
    ArgumentParse.add_argument(parser, "--config",
                               help="YAML configuration file.",
                               default="celeste.yml")
    ArgumentParse.add_argument(parser, "--method",
                               help="inference method: {joint_vi, single_vi, mcmc}",
                               default="joint_vi")

    args = ArgumentParse.parse_args(parser, console_args)

    # read config
    config, known_datasets = read_config(args["config"])

    # initialize data sets
    names = split(args["datasets"], ",")
    datasets = [known_datasets[name] for name in names]

    box = BoundingBox(args["ramin"], args["ramax"],
                      args["decmin"], args["decmax"])

    # load images from all datasets
    images = Image[]
    for dataset in datasets
        append!(images, load_images(dataset, box))
    end

    results = infer_box(images, box; config = config,
                        method = Symbol(args["method"]))
    save_results(".", box, results)
    return 0
end


function set_wally(nodes, bdcon, wally, cellcenter, cellxmax, cellymax, cellzmax, icell)
	# serch wall
	swith_wall = zeros(6)   # x-, x+, y-, y+, z-, z+を検索する
							# 1:壁
	for i in 1:6
		if Int(bdcon[i][1]) == 2 || Int(bdcon[i][1]) == 7
			swith_wall[i] = 1
		end
	end

	# cal number of ponits
	nop = 0
	if swith_wall[1] == 1
		nop = nop + (cellymax-1)*(cellzmax-1)
	end
	if swith_wall[2] == 1
		nop = nop + (cellymax-1)*(cellzmax-1)
	end
	if swith_wall[3] == 1
		nop = nop + (cellzmax-1)*(cellxmax-1)
	end
	if swith_wall[4] == 1
		nop = nop + (cellzmax-1)*(cellxmax-1)
	end
	if swith_wall[5] == 1
		nop = nop + (cellxmax-1)*(cellzmax-1)
	end
	if swith_wall[6] == 1
		nop = nop + (cellxmax-1)*(cellzmax-1)
	end

	wallpoint = zeros(nop, 3) # (格子数, 座標)
	ite = 1

	# x-
	if swith_wall[1] == 1
		for k in 1+icell:cellzmax-icell
			for j in 1+icell:cellymax-icell
				wallpoint[ite,1] = nodes[icell+1,j,k,1]
				wallpoint[ite,2] = nodes[icell+1,j,k,2]
				wallpoint[ite,3] = nodes[icell+1,j,k,3]
				ite = ite + 1
			end
		end
	end
	# x+
	if swith_wall[2] == 1
		for k in 1+icell:cellzmax-icell
			for j in 1+icell:cellymax-icell
				wallpoint[ite,1] = nodes[cellxmax-icell+1,j,k,1]
				wallpoint[ite,2] = nodes[cellxmax-icell+1,j,k,2]
				wallpoint[ite,3] = nodes[cellxmax-icell+1,j,k,3]
				ite = ite + 1
			end
		end
	end
	# y-
	if swith_wall[3] == 1
		for k in 1+icell:cellzmax-icell
			for i in 1+icell:cellxmax-icell
				wallpoint[ite,1] = nodes[i,icell+1,k,1]
				wallpoint[ite,2] = nodes[i,icell+1,k,2]
				wallpoint[ite,3] = nodes[i,icell+1,k,3]
				ite = ite + 1
			end
		end
	end
	# y+
	if swith_wall[4] == 1
		for k in 1+icell:cellzmax-icell
			for i in 1+icell:cellxmax-icell
				wallpoint[ite,1] = nodes[i,cellymax-icell+1,k,1]
				wallpoint[ite,2] = nodes[i,cellymax-icell+1,k,2]
				wallpoint[ite,3] = nodes[i,cellymax-icell+1,k,3]
				ite = ite + 1
			end
		end
	end
	# z-
	if swith_wall[5] == 1
		for j in 1+icell:cellymax-icell
			for i in 1+icell:cellxmax-icell
				wallpoint[ite,1] = nodes[i,j,icell+1,1]
				wallpoint[ite,2] = nodes[i,j,icell+1,2]
				wallpoint[ite,3] = nodes[i,j,icell+1,3]
				ite = ite + 1
			end
		end
	end
	# z+
	if swith_wall[6] == 1
		for j in 1+icell:cellymax-icell
			for i in 1+icell:cellxmax-icell
				wallpoint[ite,1] = nodes[i,j,cellzmax-icell+1,1]
				wallpoint[ite,2] = nodes[i,j,cellzmax-icell+1,2]
				wallpoint[ite,3] = nodes[i,j,cellzmax-icell+1,3]
				ite = ite + 1
			end
		end
	end

	# 距離格納用
	distance_list = zeros(nop, 3)
	distance = 0

	# 最大値を2つ抽出
	for k in 1+icell:cellzmax-icell
		for j in 1+icell:cellymax-icell
			for i in 1+icell:cellxmax-icell
				for n in 1:nop
					distance = ((cellcenter[i,j,k,1]-wallpoint[n,1])^2 +
								(cellcenter[i,j,k,2]-wallpoint[n,2])^2 +
								(cellcenter[i,j,k,3]-wallpoint[n,3])^2
								)^0.5
					distance_list[n,1] = n
					distance_list[n,2] = distance
				end

				qdistance_list = copy(distance_list)
				qdistance_list = quicksort(qdistance_list,1,nop,2,2)

				# wall_point
				point1 = Int(qdistance_list[1,1])
				point2 = Int(qdistance_list[2,1])

				dis_bottom = ((wallpoint[point1, 1] - wallpoint[point2, 1])^2 +
							  (wallpoint[point1, 2] - wallpoint[point2, 2])^2 +
							  (wallpoint[point1, 3] - wallpoint[point2, 3])^2)^0.5

				# O型格子の関係で同じ点が二つあるため
				if dis_bottom == 0.0
					point1 = Int(qdistance_list[1,1])
					point2 = Int(qdistance_list[3,1])

					dis_bottom = ((wallpoint[point1, 1] - wallpoint[point2, 1])^2 +
				   			      (wallpoint[point1, 2] - wallpoint[point2, 2])^2 +
							      (wallpoint[point1, 3] - wallpoint[point2, 3])^2)^0.5
				end

				dis1 = distance_list[point1, 2]
				dis2 = distance_list[point2, 2]

				# ヘロンの公式から距離を算出
				# 正確な平面への投影はしてないので、注意
				s = (dis_bottom + dis1 + dis2)*0.5
				surface = (s*(s-dis_bottom)*(s-dis1)*(s-dis2))^0.5
				wally[i,j,k] = surface/dis_bottom * 2
			end
		end
	end

	# 仮想セルへの代入
	for k in 1+icell:cellzmax-icell
		for j in 1:icell
			for i in 1+icell:cellxmax-icell
				wally[i,j,k] = wally[i,icell+1,k]
				wally[i,cellymax-j+1,k] = wally[i,cellymax-icell,k]
			end
		end
	end
	for k in 1+icell:cellzmax-icell
		for j in 1+icell:cellymax-icell
			for i in 1:icell
				wally[i,j,k] = wally[icell+1,j,k]
				wally[cellxmax-i+1,j,k] = wally[cellxmax-icell,j,k]
			end
		end
	end
	for k in 1:icell
		for j in 1+icell:cellymax-icell
			for i in 1+icell:cellxmax-icell
				wally[i,j,k] = wally[i,j,icell+1]
				wally[i,j,cellzmax-k+1] = wally[i,j,cellzmax-icell]
			end
		end
	end

	return wally, swith_wall
end

function cal_yplus(yplus, Qbase, wally, swith_wall, mu, cellxmax, cellymax, cellzmax, vecAx, vecAy, vecAz, volume, icell)
	tvA = zeros(3)

	# 壁に並行な方向の速度を求める
	# 速度ベクトル内壁に垂直な成分は、ux = vec{u}*vecAx
	# u_purpose = (u^2+v^2+w^2 - ux^2)^2
	for k in 1+icell:cellzmax-icell
		for j in 1+icell:cellymax-icell
			for i in 1+icell:cellzmax-icell
				if swith_wall[1] == 1 || swith_wall[2] == 1
					tvA[1] = 0.5*(vecAx[i-1,j,k,1] + vecAx[i,j,k,1])
					tvA[2] = 0.5*(vecAx[i-1,j,k,2] + vecAx[i,j,k,2])
					tvA[3] = 0.5*(vecAx[i-1,j,k,3] + vecAx[i,j,k,3])
				elseif swith_wall[3] == 1 || swith_wall[4] == 1
					tvA[1] = 0.5*(vecAy[i,j-1,k,1] + vecAy[i,j,k,1])
					tvA[2] = 0.5*(vecAy[i,j-1,k,2] + vecAy[i,j,k,2])
					tvA[3] = 0.5*(vecAy[i,j-1,k,3] + vecAy[i,j,k,3])
				elseif swith_wall[5] == 1 || swith_wall[6] == 1
					tvA[1] = 0.5*(vecAz[i,j,k-1,1] + vecAz[i,j,k,1])
					tvA[2] = 0.5*(vecAz[i,j,k-1,2] + vecAz[i,j,k,2])
					tvA[3] = 0.5*(vecAz[i,j,k-1,3] + vecAz[i,j,k,3])
				end

				sqA = (tvA[1]^2 + tvA[2]^2 + tvA[3]^2)^0.5
				tvA[1] = tvA[1] / sqA
				tvA[2] = tvA[2] / sqA
				tvA[3] = tvA[3] / sqA

				nu = mu[i,j,k] / Qbase[i,j,k,1]
				u = tvA[1]*Qbase[i,j,k,2] + tvA[2]*Qbase[i,j,k,3] + tvA[3]*Qbase[i,j,k,4]
				ubar = (Qbase[i,j,k,2]^2 + Qbase[i,j,k,3]^2 + Qbase[i,j,k,4]^2 - u^2)^0.5
				yplus[i,j,k], up = cal_wall_val_spalding(ubar, wally[i,j,k], nu)
			end
		end
	end

	return yplus
end

function cal_wall_val_spalding(u,y,nu)
	# spalding 定数
	k = 0.4
	B = 5.5

	# ニュートン法
	up = 200.0
	for i in 1:100
		old = up

		up = up - spalding(up,u,y,nu,k,B)/spalding_dash(up,u,y,nu,k,B)

		if abs(up-old)/up < 10^(-4)
			break
		end
	end

	yplus = u*y/nu/up
	return yplus, up
end

function spalding(up,u,y,nu,k,B)
	t = k*up
	F = - 1/up* u*y/nu + up + exp(-k*B)*(exp(t)-1-t-t^2/2-t^3/6)
	return F
end

function spalding_dash(up,u,y,nu,k,B)
	t = k*up
	Fdash = 1/up^2* u*y/nu + 1 + exp(-k*B)*(k*exp(t)-k-k^2*up-k^3*up^2/2)
	return Fdash
end

# test
#=
uu = 6.64037328
yy = 0.06299976860
nunu = 1.406107133e-5
println(cal_wall_val_spalding(uu,yy,nunu))
throw(UndefVarError(:x))
=#

function wallf_Van_Driest(yplus)
	f = 1 - exp(-yplus/26)
	return f
end


using DocStringExtensions
using KernelAbstractions

using ..Mesh.Geometry
using ..VariableTemplates

import ..Mesh.Grids:
    _ξ1x1, _ξ2x1, _ξ3x1, _ξ1x2, _ξ2x2, _ξ3x2, _ξ1x3, _ξ2x3, _ξ3x3
import ..MPIStateArrays: array_device

"""
    VectorGradients{
        FT <: AbstractFloat,
        FTA2D <: AbstractArray{FT, 2},
        FTA3D <: AbstractArray{FT, 3},
    }

This data structure stores the spatial gradients of a velocity field.

# Fields

$(DocStringExtensions.FIELDS)

# Usage

    VectorGradients(data)

# Arguments for the inner constructor
 - `data`: 3-dimensional device array containing the spatial gradients
   (the second dimension must be 9)
"""
struct VectorGradients{
    FT <: AbstractFloat,
    FTA2D <: AbstractArray{FT, 2},
    FTA3D <: AbstractArray{FT, 3},
}
    "Device array storing the spatial gradient data"
    data::FTA3D
    "View of ∂u₁/∂x₁"
    ∂₁u₁::FTA2D
    "View of ∂u₁/∂x₂"
    ∂₂u₁::FTA2D
    "View of ∂u₁/∂x₃"
    ∂₃u₁::FTA2D
    "View of ∂u₂/∂x₁"
    ∂₁u₂::FTA2D
    "View of ∂u₂/∂x₂"
    ∂₂u₂::FTA2D
    "View of ∂u₂/∂x₃"
    ∂₃u₂::FTA2D
    "View of ∂u₃/∂x₁"
    ∂₁u₃::FTA2D
    "View of ∂u₃/∂x₂"
    ∂₂u₃::FTA2D
    "View of ∂u₃/∂x₃"
    ∂₃u₃::FTA2D

    function VectorGradients(
        data::AbstractArray{FT, 3},
    ) where {FT <: AbstractFloat}
        ∂₁u₁, ∂₂u₁, ∂₃u₁ =
            view(data, :, 1, :), view(data, :, 2, :), view(data, :, 3, :)
        ∂₁u₂, ∂₂u₂, ∂₃u₂ =
            view(data, :, 4, :), view(data, :, 5, :), view(data, :, 6, :)
        ∂₁u₃, ∂₂u₃, ∂₃u₃ =
            view(data, :, 7, :), view(data, :, 8, :), view(data, :, 9, :)

        return new{FT, typeof(∂₁u₁), typeof(data)}(
            data,
            ∂₁u₁,
            ∂₂u₁,
            ∂₃u₁,
            ∂₁u₂,
            ∂₂u₂,
            ∂₃u₂,
            ∂₁u₃,
            ∂₂u₃,
            ∂₃u₃,
        )
    end
end

"""
    VectorGradients(dg::DGModel, Q::MPIStateArray)

This constructor computes the spatial gradients of the velocity field.

# Arguments
 - `dg`: DGModel
 - `Q`: MPIStateArray containing the conservative state variables
"""
function VectorGradients(dg::DGModel, Q::MPIStateArray)
    bl = dg.balance_law
    FT = eltype(dg.grid)
    N = polynomialorder(dg.grid)
    Nq = N + 1
    npoints = Nq^3
    nrealelem = length(dg.grid.topology.realelems)

    g = similar(Q.realdata, npoints, nrealelem, 3, 3)
    data = similar(Q.realdata, npoints, 9, nrealelem)

    ind = [
        varsindex(vars_state_conservative(bl, FT), :ρ)
        varsindex(vars_state_conservative(bl, FT), :ρu)
    ]
    _ρ, _ρu, _ρv, _ρw = ind[1], ind[2], ind[3], ind[4]

    device = array_device(Q)
    workgroup = (Nq, Nq)
    ndrange = (nrealelem * Nq, Nq)

    kernel = vector_gradients_kernel!(device, workgroup)
    event = kernel(
        Q.realdata,
        dg.grid.D,
        dg.grid.vgeo,
        g,
        data,
        _ρ,
        _ρu,
        _ρv,
        _ρw,
        Val(Nq),
        ndrange = ndrange,
    )
    wait(event)

    return VectorGradients(data)
end

@kernel function vector_gradients_kernel!(
    sv::AbstractArray{FT},
    D::AbstractArray{FT, 2},
    vgeo::AbstractArray{FT},
    g::AbstractArray{FT, 4},
    vgrad_data::AbstractArray{FT, 3},
    _ρ::Int,
    _ρu::Int,
    _ρv::Int,
    _ρw::Int,
    ::Val{qm1},
) where {qm1, FT <: AbstractFloat}

    e = @index(Group, Linear)
    i, j = @index(Local, NTuple)

    s_D = @localmem FT (qm1, qm1)
    s_U = @localmem FT (qm1, qm1)
    s_V = @localmem FT (qm1, qm1)
    s_W = @localmem FT (qm1, qm1)

    s_D[i, j] = D[i, j]
    @synchronize
    # computing derivatives with respect to ξ1
    for t in 1:qm1, s in 1:qm1
        ijk = j + ((s - 1) + (t - 1) * qm1) * qm1
        s_U[i, j] = s_D[i, j] * (sv[ijk, _ρu, e] / sv[ijk, _ρ, e])
        s_V[i, j] = s_D[i, j] * (sv[ijk, _ρv, e] / sv[ijk, _ρ, e])
        s_W[i, j] = s_D[i, j] * (sv[ijk, _ρw, e] / sv[ijk, _ρ, e])
        @synchronize
        if j == 1
            for r in 2:qm1
                s_U[i, 1] += s_U[i, r]
                s_V[i, 1] += s_V[i, r]
                s_W[i, 1] += s_W[i, r]
            end
        end
        @synchronize
        if j == 1
            g[i + ((s - 1) + (t - 1) * qm1) * qm1, e, 1, 1] = s_U[i, 1] # ∂u₁∂ξ₁
            g[i + ((s - 1) + (t - 1) * qm1) * qm1, e, 2, 1] = s_V[i, 1] # ∂u₂∂ξ₁
            g[i + ((s - 1) + (t - 1) * qm1) * qm1, e, 3, 1] = s_W[i, 1] # ∂u₃∂ξ₁
        end
    end
    @synchronize
    # computing derivatives with respect to ξ2
    for t in 1:qm1, r in 1:qm1
        ijk = r + ((j - 1) + (t - 1) * qm1) * qm1
        s_U[i, j] = s_D[i, j] * (sv[ijk, _ρu, e] / sv[ijk, _ρ, e])
        s_V[i, j] = s_D[i, j] * (sv[ijk, _ρv, e] / sv[ijk, _ρ, e])
        s_W[i, j] = s_D[i, j] * (sv[ijk, _ρw, e] / sv[ijk, _ρ, e])
        @synchronize
        if j == 1
            for s in 2:qm1
                s_U[i, 1] += s_U[i, s]
                s_V[i, 1] += s_V[i, s]
                s_W[i, 1] += s_W[i, s]
            end
        end
        @synchronize
        if j == 1
            g[r + ((i - 1) + (t - 1) * qm1) * qm1, e, 1, 2] = s_U[i, 1] # ∂u₁∂ξ₂
            g[r + ((i - 1) + (t - 1) * qm1) * qm1, e, 2, 2] = s_V[i, 1] # ∂u₂∂ξ₂
            g[r + ((i - 1) + (t - 1) * qm1) * qm1, e, 3, 2] = s_W[i, 1] # ∂u₃∂ξ₂
        end
    end
    @synchronize
    # computing derivatives with respect to ξ3
    for s in 1:qm1, r in 1:qm1
        ijk = r + ((s - 1) + (j - 1) * qm1) * qm1
        s_U[i, j] = s_D[i, j] * (sv[ijk, _ρu, e] / sv[ijk, _ρ, e])
        s_V[i, j] = s_D[i, j] * (sv[ijk, _ρv, e] / sv[ijk, _ρ, e])
        s_W[i, j] = s_D[i, j] * (sv[ijk, _ρw, e] / sv[ijk, _ρ, e])
        @synchronize
        if j == 1
            for t in 2:qm1
                s_U[i, 1] += s_U[i, t]
                s_V[i, 1] += s_V[i, t]
                s_W[i, 1] += s_W[i, t]
            end
        end
        @synchronize
        if j == 1
            g[r + ((s - 1) + (i - 1) * qm1) * qm1, e, 1, 3] = s_U[i, 1] # ∂u₁∂ξ₃
            g[r + ((s - 1) + (i - 1) * qm1) * qm1, e, 2, 3] = s_V[i, 1] # ∂u₂∂ξ₃
            g[r + ((s - 1) + (i - 1) * qm1) * qm1, e, 3, 3] = s_W[i, 1] # ∂u₃∂ξ₃
        end
    end
    @synchronize

    ∂₁u₁, ∂₂u₁, ∂₃u₁ = 1, 2, 3
    ∂₁u₂, ∂₂u₂, ∂₃u₂ = 4, 5, 6
    ∂₁u₃, ∂₂u₃, ∂₃u₃ = 7, 8, 9

    for k in 1:qm1
        ijk = i + ((j - 1) + (k - 1) * qm1) * qm1

        ξ1x1 = vgeo[ijk, _ξ1x1, e]
        ξ1x2 = vgeo[ijk, _ξ1x2, e]
        ξ1x3 = vgeo[ijk, _ξ1x3, e]
        ξ2x1 = vgeo[ijk, _ξ2x1, e]
        ξ2x2 = vgeo[ijk, _ξ2x2, e]
        ξ2x3 = vgeo[ijk, _ξ2x3, e]
        ξ3x1 = vgeo[ijk, _ξ3x1, e]
        ξ3x2 = vgeo[ijk, _ξ3x2, e]
        ξ3x3 = vgeo[ijk, _ξ3x3, e]

        vgrad_data[ijk, ∂₁u₁, e] =
            g[ijk, e, 1, 1] * ξ1x1 +
            g[ijk, e, 1, 2] * ξ2x1 +
            g[ijk, e, 1, 3] * ξ3x1
        vgrad_data[ijk, ∂₂u₁, e] =
            g[ijk, e, 1, 1] * ξ1x2 +
            g[ijk, e, 1, 2] * ξ2x2 +
            g[ijk, e, 1, 3] * ξ3x2
        vgrad_data[ijk, ∂₃u₁, e] =
            g[ijk, e, 1, 1] * ξ1x3 +
            g[ijk, e, 1, 2] * ξ2x3 +
            g[ijk, e, 1, 3] * ξ3x3

        vgrad_data[ijk, ∂₁u₂, e] =
            g[ijk, e, 2, 1] * ξ1x1 +
            g[ijk, e, 2, 2] * ξ2x1 +
            g[ijk, e, 2, 3] * ξ3x1
        vgrad_data[ijk, ∂₂u₂, e] =
            g[ijk, e, 2, 1] * ξ1x2 +
            g[ijk, e, 2, 2] * ξ2x2 +
            g[ijk, e, 2, 3] * ξ3x2
        vgrad_data[ijk, ∂₃u₂, e] =
            g[ijk, e, 2, 1] * ξ1x3 +
            g[ijk, e, 2, 2] * ξ2x3 +
            g[ijk, e, 2, 3] * ξ3x3

        vgrad_data[ijk, ∂₁u₃, e] =
            g[ijk, e, 3, 1] * ξ1x1 +
            g[ijk, e, 3, 2] * ξ2x1 +
            g[ijk, e, 3, 3] * ξ3x1
        vgrad_data[ijk, ∂₂u₃, e] =
            g[ijk, e, 3, 1] * ξ1x2 +
            g[ijk, e, 3, 2] * ξ2x2 +
            g[ijk, e, 3, 3] * ξ3x2
        vgrad_data[ijk, ∂₃u₃, e] =
            g[ijk, e, 3, 1] * ξ1x3 +
            g[ijk, e, 3, 2] * ξ2x3 +
            g[ijk, e, 3, 3] * ξ3x3
    end
end

#--------------------------------------------------------------------------------------------------

"""
    Vorticity{
        FT <: AbstractFloat,
        FTA2D <: AbstractArray{FT, 2},
        FTA3D <: AbstractArray{FT, 3},
    }

This data structure stores the vorticity of a velocity field.

# Fields

$(DocStringExtensions.FIELDS)

# Usage

    Vorticity(data)

# Arguments for the inner constructor
 - `data`: 3-dimensional device array containing the vorticity data
   (the second dimension must be 3)
"""
struct Vorticity{
    FT <: AbstractFloat,
    FTA2D <: AbstractArray{FT, 2},
    FTA3D <: AbstractArray{FT, 3},
}
    "Device array storing the vorticity data"
    data::FTA3D
    "View of x1 component of vorticity"
    Ω₁::FTA2D
    "View of x2 component of vorticity"
    Ω₂::FTA2D
    "View of x3 component of vorticity"
    Ω₃::FTA2D
    function Vorticity(data::AbstractArray{FT, 3}) where {FT <: AbstractFloat}
        Ω₁ = view(data, :, 1, :)
        Ω₂ = view(data, :, 2, :)
        Ω₃ = view(data, :, 3, :)
        return new{FT, typeof(Ω₁), typeof(data)}(data, Ω₁, Ω₂, Ω₃)
    end
end

"""
    Vorticity(
        dg::DGModel,
        vgrad::VectorGradients,
    )

This function computes the vorticity of the velocity field.

# Arguments
 - `dg`: DGModel
 - `vgrad`: vector gradients
"""
function Vorticity(dg::DGModel, vgrad::VectorGradients)
    bl = dg.balance_law
    FT = eltype(dg.grid)
    N = polynomialorder(dg.grid)
    Nq = N + 1
    npoints = Nq^3
    nrealelem = length(dg.grid.topology.realelems)

    data = similar(vgrad.data, npoints, 3, nrealelem)

    device = array_device(data)
    workgroup = (Nq, Nq)
    ndrange = (nrealelem * Nq, Nq)

    kernel = vorticity_kernel!(device, workgroup)
    event = kernel(vgrad.data, data, Val(Nq), ndrange = ndrange)
    wait(event)

    return Vorticity(data)
end

@kernel function vorticity_kernel!(
    vgrad_data::AbstractArray{FT, 3},
    vort_data::AbstractArray{FT, 3},
    ::Val{qm1},
) where {qm1, FT <: AbstractFloat}

    e = @index(Group, Linear)
    i, j = @index(Local, NTuple)

    ∂₁u₁, ∂₂u₁, ∂₃u₁ = 1, 2, 3
    ∂₁u₂, ∂₂u₂, ∂₃u₂ = 4, 5, 6
    ∂₁u₃, ∂₂u₃, ∂₃u₃ = 7, 8, 9
    Ω₁, Ω₂, Ω₃ = 1, 2, 3

    for k in 1:qm1
        ijk = i + ((j - 1) + (k - 1) * qm1) * qm1

        vort_data[ijk, Ω₁, e] =
            vgrad_data[ijk, ∂₂u₃, e] - vgrad_data[ijk, ∂₃u₂, e]
        vort_data[ijk, Ω₂, e] =
            vgrad_data[ijk, ∂₃u₁, e] - vgrad_data[ijk, ∂₁u₃, e]
        vort_data[ijk, Ω₃, e] =
            vgrad_data[ijk, ∂₁u₂, e] - vgrad_data[ijk, ∂₂u₁, e]
    end
end