using KernelAbstractions
using Oceananigans.Utils
using Oceananigans.Architectures: device, @hascuda, CPU, GPU, array_type
using Oceananigans.Fields: datatuple
using Oceananigans.BoundaryConditions: apply_x_bcs!, apply_y_bcs!, apply_z_bcs!

@kernel function update_total_density!(total_density, grid, gases, tracers)
    i, j, k = @index(Global, NTuple)

    @inbounds total_density[i, j, k] = diagnose_density(i, j, k, grid, gases, tracers)
end

# This is for users. Do not use in time_stepping.jl as it doesn't make use of intermediate fields.
function update_total_density!(model)

    density_update_event =
        launch!(model.architecture, model.grid, :xyz, update_total_density!,
                datatuple(model.total_density), model.grid, model.gases, datatuple(model.tracers),
                dependencies=Event(device(model.architecture)))

    wait(device(model.architecture), density_update_event)

    return nothing
end

#####
##### Computing slow source terms (viscous dissipation, diffusion, and Coriolis terms).
#####

function compute_slow_source_terms!(slow_source_terms, arch, grid, thermodynamic_variable, gases, gravity, coriolis, closure, total_density, momenta, tracers, diffusivities, forcing, clock)

    slow_source_terms, total_density, momenta, tracers, diffusivities =
        datatuples(slow_source_terms, total_density, momenta, tracers, diffusivities)

    workgroup, worksize = work_layout(grid, :xyz)
    barrier = Event(device(arch))

    momentum_kernel! = compute_slow_momentum_source_terms!(device(arch), workgroup, worksize)
    tracer_kernel! = compute_slow_tracer_source_terms!(device(arch), workgroup, worksize)
    thermodynamic_variable_kernel! = compute_slow_thermodynamic_variable_source_terms!(device(arch), workgroup, worksize)

    momentum_event = momentum_kernel!(slow_source_terms, grid, coriolis, closure, total_density, momenta, tracers, diffusivities, forcing, clock, dependencies=barrier)

    events = [momentum_event]

    for (tracer_index, ρc_name) in enumerate(propertynames(tracers))
        ρc   = getproperty(tracers, ρc_name)
        S_ρc = getproperty(slow_source_terms.tracers, ρc_name)
        forcing_ρc = getproperty(forcing, ρc_name)

        tracer_event = tracer_kernel!(S_ρc, grid, closure, tracer_index, total_density, ρc, momenta, tracers, diffusivities, forcing_ρc, clock, dependencies=barrier)
        push!(events, tracer_event)
    end

    thermodynamic_variable_event = thermodynamic_variable_kernel!(slow_source_terms.tracers[1], grid, thermodynamic_variable, gases, gravity, closure, total_density, momenta, tracers, diffusivities, dependencies=barrier)
    push!(events, thermodynamic_variable_event)

    wait(device(arch), MultiEvent(Tuple(events)))

    return nothing
end

@kernel function compute_slow_momentum_source_terms!(slow_source_terms, grid, coriolis, closure, total_density, momenta, tracers, diffusivities, forcing, clock)
    i, j, k = @index(Global, NTuple)

    @inbounds slow_source_terms.ρu[i, j, k] = ρu_slow_source_term(i, j, k, grid, coriolis, closure, total_density, momenta, diffusivities) + forcing.ρu(i, j, k, grid, clock, merge(momenta, tracers))
    @inbounds slow_source_terms.ρv[i, j, k] = ρv_slow_source_term(i, j, k, grid, coriolis, closure, total_density, momenta, diffusivities) + forcing.ρv(i, j, k, grid, clock, merge(momenta, tracers))
    @inbounds slow_source_terms.ρw[i, j, k] = ρw_slow_source_term(i, j, k, grid, coriolis, closure, total_density, momenta, diffusivities) + forcing.ρw(i, j, k, grid, clock, merge(momenta, tracers))
end

@kernel function compute_slow_tracer_source_terms!(S_ρc, grid, closure, tracer_index, total_density, ρc, momenta, tracers, diffusivities, forcing, clock)
    i, j, k = @index(Global, NTuple)

    @inbounds S_ρc[i, j, k] = ρc_slow_source_term(i, j, k, grid, closure, tracer_index, total_density, ρc, diffusivities) + forcing(i, j, k, grid, clock, merge(momenta, tracers))
end

@kernel function compute_slow_thermodynamic_variable_source_terms!(S_ρt, grid, thermodynamic_variable, gases, gravity, closure, total_density, momenta, tracers, diffusivities)
    i, j, k = @index(Global, NTuple)

    @inbounds S_ρt[i, j, k] += ρt_slow_source_term(i, j, k, grid, closure, thermodynamic_variable, gases, gravity, total_density, momenta, tracers, diffusivities)
end

#####
##### Computing fast source terms (advection, pressure gradient, and buoyancy terms).
#####

function compute_fast_source_terms!(fast_source_terms, arch, grid, thermodynamic_variable, gases, gravity, advection_scheme, total_density, momenta, tracers, slow_source_terms)

    fast_source_terms, total_density, momenta, tracers, slow_source_terms =
        datatuples(fast_source_terms, total_density, momenta, tracers, slow_source_terms)

    workgroup, worksize = work_layout(grid, :xyz)
    barrier = Event(device(arch))

    momentum_kernel! = compute_fast_momentum_source_terms!(device(arch), workgroup, worksize)
    tracer_kernel! = compute_fast_tracer_source_terms!(device(arch), workgroup, worksize)
    thermodynamic_variable_kernel! = compute_fast_thermodynamic_variable_source_terms!(device(arch), workgroup, worksize)

    momentum_event = momentum_kernel!(fast_source_terms, grid, thermodynamic_variable, gases, gravity, advection_scheme, total_density, momenta, tracers, slow_source_terms, dependencies=barrier)

    events = [momentum_event]

    for ρc_name in propertynames(tracers)
        ρc   = getproperty(tracers, ρc_name)
        F_ρc = getproperty(fast_source_terms.tracers, ρc_name)
        S_ρc = getproperty(slow_source_terms.tracers, ρc_name)

        tracer_event = tracer_kernel!(F_ρc, grid, advection_scheme, total_density, momenta, ρc, S_ρc, dependencies=barrier)
        push!(events, tracer_event)
    end

    thermodynamic_variable_event = thermodynamic_variable_kernel!(fast_source_terms.tracers[1], grid, thermodynamic_variable, gases, gravity, total_density, momenta, tracers, dependencies=barrier)
    push!(events, thermodynamic_variable_event)

    wait(device(arch), MultiEvent(Tuple(events)))

    return nothing
end

@kernel function compute_fast_momentum_source_terms!(fast_source_terms, grid, thermodynamic_variable, gases, gravity, advection_scheme, total_density, momenta, tracers, slow_source_terms)
    i, j, k = @index(Global, NTuple)

    @inbounds fast_source_terms.ρu[i, j, k] = ρu_fast_source_term(i, j, k, grid, thermodynamic_variable, gases, gravity, advection_scheme, total_density, momenta, tracers, slow_source_terms.ρu)
    @inbounds fast_source_terms.ρv[i, j, k] = ρv_fast_source_term(i, j, k, grid, thermodynamic_variable, gases, gravity, advection_scheme, total_density, momenta, tracers, slow_source_terms.ρv)
    @inbounds fast_source_terms.ρw[i, j, k] = ρw_fast_source_term(i, j, k, grid, thermodynamic_variable, gases, gravity, advection_scheme, total_density, momenta, tracers, slow_source_terms.ρw)
end

@kernel function compute_fast_tracer_source_terms!(F_ρc, grid, advection_scheme, total_density, momenta, ρc, S_ρc)
    i, j, k = @index(Global, NTuple)

    @inbounds F_ρc[i, j, k] = ρc_fast_source_term(i, j, k, grid, advection_scheme, total_density, momenta, ρc, S_ρc)
end

@kernel function compute_fast_thermodynamic_variable_source_terms!(F_ρt, grid, thermodynamic_variable, gases, gravity, total_density, momenta, tracers)
    i, j, k = @index(Global, NTuple)

    @inbounds F_ρt[i, j, k] += ρt_fast_source_term(i, j, k, grid, thermodynamic_variable, gases, gravity, total_density, momenta, tracers)
end

#####
##### Calculating boundary tendency contributions
#####

function calculate_boundary_tendency_contributions!(source_terms, arch, momenta, tracers, clock, model_fields)

    barrier = Event(device(arch))

    events = []

    # Momentum fields
    momentum_source_terms = (source_terms.ρu, source_terms.ρv, source_terms.ρw)

    for (ρϕ_source_term, ρϕ) in zip(momentum_source_terms, momenta)
        x_bcs_event = apply_x_bcs!(ρϕ_source_term, ρϕ, arch, barrier, clock, model_fields)
        y_bcs_event = apply_y_bcs!(ρϕ_source_term, ρϕ, arch, barrier, clock, model_fields)
        z_bcs_event = apply_z_bcs!(ρϕ_source_term, ρϕ, arch, barrier, clock, model_fields)

        push!(events, x_bcs_event, y_bcs_event, z_bcs_event)
    end

    # Tracer fields
    for (ρϕ_source_term, ρϕ) in zip(source_terms.tracers, tracers)
        x_bcs_event = apply_x_bcs!(ρϕ_source_term, ρϕ, arch, barrier, clock, model_fields)
        y_bcs_event = apply_y_bcs!(ρϕ_source_term, ρϕ, arch, barrier, clock, model_fields)
        z_bcs_event = apply_z_bcs!(ρϕ_source_term, ρϕ, arch, barrier, clock, model_fields)

        push!(events, x_bcs_event, y_bcs_event, z_bcs_event)
    end

    events = filter(e -> typeof(e) <: Event, events)

    wait(device(arch), MultiEvent(Tuple(events)))

    return nothing
end

#####
##### Advancing state variables
#####

function advance_state_variables!(state_variables, arch, grid, momenta, tracers, fast_source_terms; Δt)

    state_variables, momenta, tracers, fast_source_terms =
        datatuples(state_variables, momenta, tracers, fast_source_terms )

    workgroup, worksize = work_layout(grid, :xyz)
    barrier = Event(device(arch))

    momentum_kernel! = advance_momentum!(device(arch), workgroup, worksize)
    tracer_kernel! = advance_tracer!(device(arch), workgroup, worksize)

    momentum_event = momentum_kernel!(state_variables, grid, momenta, fast_source_terms, Δt, dependencies=barrier)

    events = [momentum_event]

    for ρc_name in propertynames(tracers)
        ρc   = getproperty(tracers, ρc_name)
        ρc⁺  = getproperty(state_variables.tracers, ρc_name)
        F_ρc = getproperty(fast_source_terms.tracers, ρc_name)

        tracer_event = tracer_kernel!(ρc⁺, grid, ρc, F_ρc, Δt, dependencies=barrier)
        push!(events, tracer_event)
    end

    wait(device(arch), MultiEvent(Tuple(events)))

    return nothing
end

@kernel function advance_momentum!(momenta⁺, grid, momenta, fast_source_terms, Δt)
    i, j, k = @index(Global, NTuple)

    @inbounds momenta⁺.ρu[i, j, k] = momenta.ρu[i, j, k] + Δt * fast_source_terms.ρu[i, j, k]
    @inbounds momenta⁺.ρv[i, j, k] = momenta.ρv[i, j, k] + Δt * fast_source_terms.ρv[i, j, k]
    @inbounds momenta⁺.ρw[i, j, k] = momenta.ρw[i, j, k] + Δt * fast_source_terms.ρw[i, j, k]
end

@kernel function advance_tracer!(ρc⁺, grid, ρc, F_ρc, Δt)
    i, j, k = @index(Global, NTuple)

    @inbounds ρc⁺[i, j, k] = ρc[i, j, k] + Δt * F_ρc[i, j, k]
end


export DenseConvDims

"""
    DenseConvDims

Concrete subclass of `ConvDims` for a normal, dense, conv2d/conv3d.
"""
struct DenseConvDims{N, K, S, P, D} <: ConvDims{N}
    input_size::NTuple{N, Int}

    kernel_size::NTuple{K, Int}
    channels_in::Int
    channels_out::Int
    groupcount::Int

    stride::NTuple{S, Int}
    padding::NTuple{P, Int}
    dilation::NTuple{D, Int}
    flipkernel::Bool
end

function DenseConvDims(
    x_size::NTuple{M}, w_size::NTuple{M};
    stride = 1, padding = 0, dilation = 1, groups = 1,
    flipkernel::Bool = false,
) where {M}
    sstride, ppadding, ddilation = check_spdf(
        x_size, w_size, stride, padding, dilation)

    # Ensure channels are equal
    if x_size[end - 1] != w_size[end - 1] * groups
        xs = x_size[end - 1]
        ws = w_size[end - 1]
        throw(DimensionMismatch("Input channels must match! ($xs vs. $ws)"))
    end

    # Ensure groups are valid
    if x_size[end - 1] % w_size[end - 1] != 0 || w_size[end] % groups != 0
        throw(DimensionMismatch(
            "Group count should be divisble by input and output channels ($groups vs. $(w_size[end-1:end]))"))
    end

    DenseConvDims(
        x_size[1:(end - 2)],
        w_size[1:(end - 2)], x_size[end - 1], w_size[end], groups,
        sstride, ppadding, ddilation, flipkernel)
end

function DenseConvDims(x::AbstractArray, w::AbstractArray; kwargs...)
    if ndims(x) != ndims(w)
        throw(DimensionMismatch(
            "Rank of x and w must match! ($(ndims(x)) vs. $(ndims(w)))"))
    end
    return DenseConvDims(size(x), size(w); kwargs...)
end

# Useful for constructing a new DenseConvDims that has only a few elements different
# from the original progenitor object that it inherits shapes from.
@inline DenseConvDims(
    c::C; I=input_size(c), K=kernel_size(c),
    C_in=channels_in(c), C_out=channels_out(c), S=stride(c),
    P=padding(c), D=dilation(c), F=flipkernel(c), G=groupcount(c),
) where C <: ConvDims = DenseConvDims(
    I,
    K, C_in, C_out, G,
    S, P, D, F)

@inline groupcount(c::DenseConvDims) = c.groupcount
@inline channels_in(c::DenseConvDims) = c.channels_in
@inline channels_out(c::DenseConvDims) = c.channels_out

@inline input_size(c::DenseConvDims) = c.input_size
@inline kernel_size(c::DenseConvDims) = c.kernel_size

@inline stride(c::DenseConvDims) = c.stride
@inline padding(c::DenseConvDims) = c.padding
@inline dilation(c::DenseConvDims) = c.dilation
@inline flipkernel(c::DenseConvDims) = c.flipkernel

function check_dims(x::NTuple{M}, w::NTuple{M}, y::NTuple{M}, cdims::DenseConvDims) where {M}
    # First, check that channel counts are all correct:
    @assert x[M-1] * groupcount(cdims) == channels_in(cdims) DimensionMismatch("Data input channel count ($(x[M-1]) vs. $(channels_in(cdims)))")
    @assert y[M-1] == channels_out(cdims) ÷ groupcount(cdims)  DimensionMismatch("Data output channel count ($(y[M-1]) vs. $(channels_out(cdims)))")
    @assert w[M-1] * groupcount(cdims) == channels_in(cdims) DimensionMismatch("Kernel input channel count ($(w[M-1]) vs. $(channels_in(cdims)))")
    @assert w[M] * groupcount(cdims) == channels_out(cdims) DimensionMismatch("Kernel output channel count ($(w[M]) vs. $(channels_out(cdims)))")

    # Next, check that the spatial dimensions match up
    @assert x[1:M-2] == input_size(cdims) DimensionMismatch("Data input spatial size ($(x[1:M-2]) vs. $(input_size(cdims)))")
    @assert y[1:M-2] == output_size(cdims) DimensionMismatch("Data output spatial size ($(y[1:M-2]) vs. $(output_size(cdims)))")
    @assert w[1:M-2] == kernel_size(cdims) DimensionMismatch("Kernel spatial size ($(w[1:M-2]) vs. $(kernel_size(cdims)))")

    # Check the groups match
    @assert channels_in(cdims) % groupcount(cdims) == 0 DimensionMismatch("Groups ($(groupcount(cdims))) should be divisble by input channels $(channels_in(cdims))")

    # Finally, check that the batch size matches
    @assert x[M] == y[M] DimensionMismatch("Batch size ($(x[M]) vs. $(y[M]))")
end


function mcmc(setup_mcmc::MCMCSetup, schedule::MCMCSchedule, setups::Vector{T}) where T <: ModelSetup
    workspaces = [create_workspace(setup) for setup in setups]
    wss, lls = [ws[1] for ws in workspaces], [ws[2] for ws in workspaces]
    θ = workspaces[1][3]
    ws_mcmc = create_workspace(setup_mcmc, schedule, θ)
    # adpt = adaptation_object(setup, ws) # adaptive auxiliary law currently not supported

    aux = nothing
    for step in schedule
        step.save && [save_imputed!(ws) for ws in wss]
        for i in step.idx
            if step.param_updt || typeof(ws_mcmc.updates[i]) <: Imputation
                wss = [next(ws, ws_mcmc.updates[i]) for ws in wss]
                lls, acc, θ = update!(ws_mcmc.updates[i], wss, θ, lls, step, aux)
                aux = aux_params(ws_mcmc.updates[i], aux)
                #TODO: make it nicer, currently we need a little dance here
                dance = typeof(ws_mcmc.updates[i]) <: Imputation
                dance && begin
                    update!(ws_mcmc, acc[1], θ, i)
                    for k in 2:length(setups)
                        register_accpt!(ws_mcmc.updates[i], acc[k])
                    end
                end
                !dance && update!(ws_mcmc, acc, θ, i)
                step.verbose && print("\n")
            end
        end
        step.verbose && print("-----------------------------------------------",
                              "------\n")
        step.readjust && readjust!(ws_mcmc, step.iter)
        step.fuse && fuse!(ws_mcmc, schedule)
        # ll = adaptation!(ws, adpt, step.iter, ll) adaptation currently not supported
    end
    [display_summary(ws, ws_mcmc) for ws in wss]
    wss, ws_mcmc
end

function update!(updt::Imputation{NoBlocking}, wss::Vector{<:Workspace{OS}}, θ,
                 lls, step, ::Any, headstart=false) where OS
    ρ, K = updt.ρs[1], length(wss)
    accpts = repeat([false], K)

    for k in 1:K
        ws = wss[k]
        # sample proposal starting point
        zᵒ, yᵒ = proposal_start_pt(ws, ws.P[1], ρ)

        # sample proposal path
        m = length(ws.WWᵒ)
        success = sample_segments!(1:m, ws, yᵒ, ρ, Val{headstart}())
        llᵒ = success ? (logpdf(ws.x0_prior, yᵒ) +
                         path_log_likhd(OS(), ws.XXᵒ, ws.P, 1:m, ws.fpt) +
                         lobslikelihood(ws.P[1], yᵒ)) : -Inf

        print_info(step, value(lls[k]), value(llᵒ), "impute")

        if accept_sample(llᵒ-lls[k], step.verbose)
            swap!(ws.XX, ws.XXᵒ, ws.WW, ws.WWᵒ, 1:m)
            set!(ws.z, zᵒ)
            lls[k] = llᵒ
            accpts[k] = true
        end
    end
    lls, accpts, θ
end

function update!(pu::ParamUpdate{MetropolisHastingsUpdt},
                 wss::Vector{<:Workspace{OS}}, θ, lls, step,
                 blocking::NoBlocking) where OS
    K = length(wss)

    #NOTE for now, let's sample parameter corresponding to a single path, this
    # will need to be changed for Mixed effect models
    θᵒ = rand(pu.t_kernel, θ, pu.updt_coord)               # sample new parameter
    (logpdf(pu.priors, θᵒ) === -Inf) && (return lls, false, θ)

    llᵒs = copy(lls)
    zᵒs = [copy(ws.z.val) for ws in wss]
    for k in 1:K
        ws = wss[k]
        WW, Pᵒ, P, XXᵒ, XX, fpt = ws.WW, ws.Pᵒ, ws.P, ws.XXᵒ, ws.XX, ws.fpt
        m = length(WW)
        update_laws!(Pᵒ, θᵒ)
        pu.aux.recompute_ODEs && solve_back_rec!(blocking, pu.aux.solver, Pᵒ) # compute (H, Hν, c)

        # find white noise which for a given θᵒ gives a correct starting point
        y = XX[1].yy[1]
        zᵒs[k] = inv_start_pt(y, ws.x0_prior, Pᵒ[1])

        success = find_path_from_wiener!(XXᵒ, y, WW, Pᵒ, 1:m)

        llᵒs[k] = success ? (logpdf(ws.x0_prior, y) +
                             path_log_likhd(OS(), XXᵒ, Pᵒ, 1:m, fpt) +
                             lobslikelihood(Pᵒ[1], y)) : -Inf
    end
    ll, llᵒ = sum(lls), sum(llᵒs)
    print_info(step, value(ll), value(llᵒ))

    llr = ( llᵒ - ll + prior_kernel_contrib(pu, θ, θᵒ))

    # Accept / reject
    if accept_sample(llr, step.verbose)
        for k in 1:K
            ws = wss[k]
            WW, Pᵒ, P, XXᵒ, XX, fpt = ws.WW, ws.Pᵒ, ws.P, ws.XXᵒ, ws.XX, ws.fpt
            m = length(WW)
            swap!(XX, XXᵒ, P, Pᵒ, 1:m)
            set!(ws.z, zᵒs[k])
        end
        return llᵒs, true, θᵒ
    else
        return lls, false, θ
    end
end

function update!(pu::ParamUpdate{ConjugateUpdt}, wss::Vector{<:Workspace{OS}},
                 θ, lls, step, blocking::NoBlocking) where OS
    K = length(wss)

    θᵒ = conjugate_draw(θ, [ws.XX for ws in wss], wss[1].P[1].Target, pu.priors[1], pu.updt_coord)

    total_ll_old = sum(lls)
    for k in 1:K
        ws = wss[k]
        WW, P, XX, fpt = ws.WW, ws.P, ws.XX, ws.fpt
        m = length(WW)

        update_laws!(P, θᵒ)
        pu.aux.recompute_ODEs && solve_back_rec!(blocking, pu.aux.solver, P) # compute (H, Hν, c)

        for i in 1:m    # compute wiener path WW that generates XX
            inv_solve!(EulerMaruyamaBounded(), XX[i], WW[i], P[i])
        end
        # compute white noise that generates starting point
        y = XX[1].yy[1]
        z = inv_start_pt(y, ws.x0_prior, P[1])

        llᵒ = logpdf(ws.x0_prior, y)
        llᵒ += path_log_likhd(OS(), XX, P, 1:m, fpt; skipFPT=true)
        llᵒ += lobslikelihood(P[1], y)
        lls[k] = llᵒ
        set!(ws.z, z)
    end
    print_info(step, sum(total_ll_old), sum(lls))
    return lls, true, θᵒ
end

function conjugate_draw(θ, XX::Vector{<:Vector}, PT, prior, updtIdx)
    μ = mustart(updtIdx)
    𝓦 = μ*μ'
    ϑ = SVector(thetaex(updtIdx, θ))
    for X in XX
        μ, 𝓦 = _conjugate_draw(ϑ, μ, 𝓦, X, PT, updtIdx)
    end
    Σ = inv(𝓦 + inv(Matrix(prior.Σ)))
    Σ = (Σ + Σ')/2 # eliminates numerical inconsistencies
    μ_post = Σ * (μ + Vector(prior.Σ\prior.μ))
    ϑ = rand(Gaussian(μ_post, Σ))
    move_to_proper_place(ϑ, θ, updtIdx)     # align so that dimensions agree
end


module Processing
import FFTW
using EllipsisNotation
import Glob: glob
import Luna: Maths, Fields, PhysData
import Luna.PhysData: wlfreq, c
import Luna.Grid: AbstractGrid, RealGrid, EnvGrid, from_dict
import Luna.Output: AbstractOutput, HDF5Output

"""
    Common(val)

Wrapper type to tell `scanproc` that `val` is the same for each simulation being processed,
and so only needs to be returned once rather than for each simulation in the scan.
"""
struct Common{dT}
    data::dT
end

"""
    VarLength(val)

Wrapper type to tell `scanproc` that the shape of `val` is different for each simulation being
processed. Return values wrapped in `VarLength` will be placed in an array of arrays.

!!! note
    While the **shape** of `val` can be different between simulations, the **type** must be
    the same, including the dimensionality and element type of arrays.
"""
struct VarLength{dT}
    data::dT
end

"""
    scanproc(f, scanfiles)
    scanproc(f, directory)
    scanproc(f, directory, pattern)
    scanproc(f)

Iterate over the scan output files, apply the processing function `f(o::AbstractOutput)`,
and collect the results in arrays.

The files can be given as:

- a `Vector` of `AbstractString`s containing file paths
- a directory to search for files according to the naming pattern of
    `Output.ScanHDF5Output`
- a directory and a `glob` pattern

If nothing is specified, `scanproc` uses the current working directory.

`f` can return a single value, an array, or a tuple/array of arrays/numbers. Arrays returned
by `f` must either be of the same size for each processed file, or wrapped in a `VarLength`.
Values returned by `f` which are guaranteed to be identical for each processed file can be
wrapped in a `Common`, and `scanproc` only returns these once.

# Example
```julia
Et, Eω = scanproc("path/to/scandir") do output
    t, Et = getEt(output)
    ω, Eω = getEω(output)
    energyout = energyout = Processing.VarLength(output["stats"]["energy"])
    Common(t), Et, Common(ω), Eω, energyout
end
```
"""
function scanproc(f, scanfiles::AbstractVector{<:AbstractString}; shape=nothing)
    local scanidcs, arrays
    scanfiles = sort(scanfiles)
    for (idx, fi) in enumerate(scanfiles)
        o = HDF5Output(fi)
        # wraptuple makes sure we definitely have a Tuple, even if f only returns one thing
        ret = wraptuple(f(o))
        if idx == 1 # initialise arrays
            isnothing(shape) && (shape = Tuple(o["meta"]["scanshape"]))
            scanidcs = CartesianIndices(shape)
            arrays = _arrays(ret, shape)
        end
        for (ridx, ri) in enumerate(ret)
            _addret!(arrays[ridx], scanidcs[idx], ri)
        end
    end
    unwraptuple(arrays) # if f only returns one thing, we also only return one array
end

wraptuple(x::Tuple) = x
wraptuple(x) = (x,)

unwraptuple(x::Tuple{<:Any}) = x[1] # single-element Tuple
unwraptuple(x) = x

function _addret!(array, aidcs, ri)
    array[aidcs] = ri
end

function _addret!(array, aidcs, ri::AbstractArray)
    idcs = CartesianIndices(ri)
    array[idcs, aidcs] .= ri
end

function _addret!(array, aidcs, ri::VarLength)
    array[aidcs] = ri.data
end

_addret!(array, aidcs, ri::Common) = nothing

# Default pattern for files named by ScanHDF5Output is [name]_[scanidx].h5 with 5 digits
defpattern = "*_[0-9][0-9][0-9][0-9][0-9].h5"

function scanproc(f, directory::AbstractString=pwd(), pattern::AbstractString=defpattern;
                  shape=nothing)
    scanfiles = glob(pattern, directory) # this returns absolute paths if directory given
    scanproc(f, scanfiles; shape=shape)
end

# Make array(s) with correct size to hold processing results
_arrays(ret, shape) = Array{typeof(ret)}(undef, shape)
_arrays(ret::AbstractArray, shape) = zeros(eltype(ret), (size(ret)..., shape...))
_arrays(ret::Tuple, shape) = Tuple([_arrays(ri, shape) for ri in ret])
_arrays(com::Common, shape) = com.data
_arrays(vl::VarLength, shape) = Array{typeof(vl.data), length(shape)}(undef, shape)

"""
    coherence(Eω; ndim=1)

Calculate the first-order coherence function g₁₂ of the set of fields `Eω`. The ensemble
average is taken over the last `ndim` dimensions of `Eω`, other dimensions are preserved.

See J. M. Dudley and S. Coen, Optics Letters 27, 1180 (2002).
"""
function coherence(Eω; ndim=1)
    dimsize = size(Eω)[end-ndim+1:end]
    outsize = size(Eω)[1:end-ndim]
    prodidcs = CartesianIndices(dimsize)
    restidcs = CartesianIndices(outsize)
    coherence(Eω, prodidcs, restidcs)
end

# function barrier for speedup
function coherence(Eω, prodidcs, restidcs)
    num = zeros(ComplexF64, size(restidcs))
    den1 = zeros(ComplexF64, size(restidcs))
    den2 = zeros(ComplexF64, size(restidcs))
    it = Iterators.product(prodidcs, prodidcs)
    for (idx1, idx2) in it
        Eω1 = Eω[restidcs, idx1]
        Eω2 = Eω[restidcs, idx2]
        @. num += conj(Eω1)*Eω2
        @. den1 += abs2(Eω1)
        @. den2 += abs2(Eω2)
    end
    @. abs(num/sqrt(den1*den2))
end

"""
    arrivaltime(grid, Eω; bandpass=nothing, method=:moment, oversampling=1)

Extract the arrival time of the pulse in the wavelength limits `λlims`.

# Arguments
- `bandpass` : method to bandpass the field if required. See [`window_maybe`](@ref)
- `method::Symbol` : `:moment` to use 1st moment to extract arrival time, `:peak` to use
                    the time of peak power
- `oversampling::Int` : If >1, oversample the time-domain field before extracting delay
- `sumdims` : Single `Int` or `Tuple` of `Int`s. The time-domain power will be summed over
            these dimensions (e.g. modes) before extracting the arrival time.
"""
function arrivaltime(grid::AbstractGrid, Eω;
                     bandpass=nothing, method=:moment, oversampling=1, sumdims=nothing)
    to, Eto = getEt(grid, Eω; oversampling=oversampling, bandpass=bandpass)
    Pt = abs2.(Eto)
    if !isnothing(sumdims)
        Pt = dropdims(sum(Pt; dims=sumdims); dims=sumdims)
    end
    arrivaltime(to, Pt; method=method)
end

function arrivaltime(t::AbstractVector, It::AbstractVector; method)
    if method == :moment
        Maths.moment(t, It)
    elseif method == :peak
        t[argmax(It)]
    else
        error("Unknown arrival time method $method")
    end
end

function arrivaltime(t::AbstractVector, It::AbstractArray; method)
    out = Array{Float64, ndims(It)-1}(undef, size(It)[2:end])
    cidcs = CartesianIndices(size(It)[2:end])
    for ii in cidcs
        out[ii] = arrivaltime(t, It[:, ii]; method=method)
    end
    out
end

"""
    time_bandwidth(grid, Eω; bandpass=nothing, oversampling=1)

Extract the time-bandwidth product, after bandpassing if required. The TBP
is defined here as ΔfΔt where Δx is the FWHM of x. (In this definition, the TBP of
a perfect Gaussian pulse is ≈0.44). If `oversampling` > 1, the time-domain field is
oversampled before extracting the FWHM.
"""
function time_bandwidth(grid, Eω; bandpass=nothing, oversampling=1, sumdims=nothing)
    fwt = fwhm_t(grid, Eω; bandpass=bandpass, oversampling=oversampling, sumdims=nothing)
    fwf = fwhm_f(grid, Eω; bandpass=bandpass)
    fwt.*fwf
end

"""
    fwhm_t(grid::AbstractGrid, Eω; bandpass=nothing, oversampling=1)

Extract the temporal FWHM. If `bandpass` is given, bandpass the fieldaccording to
[`window_maybe`](@ref). If `oversampling` > 1, the  time-domain field is oversampled before
extracting the FWHM. If `sumdims` is given, the time-domain power is summed over these
dimensions (e.g. modes) before extracting the FWHM.
"""
function fwhm_t(grid::AbstractGrid, Eω; bandpass=nothing, oversampling=1, sumdims=nothing)
    to, Eto = getEt(grid, Eω; oversampling=oversampling, bandpass=bandpass)
    Pt = abs2.(Eto)
    if !isnothing(sumdims)
        Pt = dropdims(sum(Pt; dims=sumdims); dims=sumdims)
    end
    fwhm(to, Pt)
end

function fwhm_t(output::AbstractOutput; kwargs...)
    grid = makegrid(output)
    fwhm_t(grid, output["Eω"]; kwargs...)
end

"""
    fwhm_f(grid, Eω::Vector; bandpass=nothing, oversampling=1)

Extract the frequency FWHM. If `bandpass` is given, bandpass the field according to
[`window_maybe`](@ref). If `sumdims` is given, the energy density is summed over these
dimensions (e.g. modes) before extracting the FWHM.
"""
function fwhm_f(grid::AbstractGrid, Eω; bandpass=nothing, oversampling=1, sumdims=nothing)
    Eω = window_maybe(grid.ω, Eω, bandpass)
    f, If = getIω(getEω(grid, Eω)..., :f)
    if !isnothing(sumdims)
        If = dropdims(sum(If; dims=sumdims); dims=sumdims)
    end
    fwhm(f, If)
end

function fwhm(x, I)
    out = Array{Float64, ndims(I)-1}(undef, size(I)[2:end])
    cidcs = CartesianIndices(size(I)[2:end])
    for ii in cidcs
        out[ii] = fwhm(x, I[:, ii])
    end
    out
end

fwhm(x::Vector, I::Vector) = Maths.fwhm(x, I)

"""
    peakpower(grid, Eω; bandpass=nothing, oversampling=1)

Extract the peak power. If `bandpass` is given, bandpass the field according to
[`window_maybe`](@ref).
"""
function peakpower(grid, Eω; bandpass=nothing, oversampling=1)
    to, Eto = getEt(grid, Eω; oversampling=oversampling, bandpass=bandpass)
    dropdims(maximum(abs2.(Eto); dims=1); dims=1)
end

"""
    energy(grid, Eω; bandpass=nothing)

Extract energy. If `bandpass` is given, bandpass the field according to
[`window_maybe`](@ref).
"""
function energy(grid, Eω; bandpass=nothing)
    Eω = window_maybe(grid.ω, Eω, bandpass)
    _, energyω = Fields.energyfuncs(grid)
    _energy(Eω, energyω)
end

function energy(output::AbstractOutput; bandpass=nothing)
    grid = makegrid(output)
    energy(grid, output["Eω"]; bandpass=bandpass)
end

_energy(Eω::Vector, energyω) = energyω(Eω)

function _energy(Eω, energyω)
    out = Array{Float64, ndims(Eω)-1}(undef, size(Eω)[2:end])
    cidcs = CartesianIndices(size(Eω)[2:end])
    for ii in cidcs
        out[ii] = _energy(Eω[:, ii], energyω)
    end
    out
end

"""
    field_autocorrelation(Et; dims=1)

Calculate the field autocorrelation of `Et`.
"""
function field_autocorrelation(Et, grid::EnvGrid; dims=1)
    FFTW.fftshift(FFTW.ifft(abs2.(FFTW.fft(Et, dims)), dims), dims)
end

function field_autocorrelation(Et, grid::RealGrid; dims=1)
    fac = FFTW.fftshift(FFTW.irfft(abs2.(FFTW.rfft(Et, dims)), length(grid.t), dims), dims)
    Maths.hilbert(fac, dim=dims)
end

"""
    intensity_autocorrelation(Et, grid)

Calculate the intensity autocorrelation of `Et` over `grid`.
"""
function intensity_autocorrelation(Et, grid; dims=1)
    real.(FFTW.fftshift(FFTW.irfft(abs2.(FFTW.rfft(Fields.It(Et, grid), dims)), length(grid.t), dims), dims))
end

"""
    coherence_time(grid, Et; dims=1)

Get the coherence time of a field `Et` over `grid`.
"""
function coherence_time(grid, Et; dims=1)
    Maths.fwhm(grid.t, abs2.(field_autocorrelation(Et, grid, dims=dims)))
end

"""
    specres(ω, Iω, specaxis, resolution, specrange; window=nothing, nsamples=10)

Smooth the spectral energy density `Iω(ω)` to account for the given `resolution`
on the defined `specaxis` and `specrange`. The `window` function to use defaults
to a Gaussian function with FWHM of `resolution`, and by default we sample `nsamples=10`
times within each `resolution`.

Note that you should prefer the `resolution` keyword of [`getIω`](@ref) instead of calling
this function directly.

The input `ω` and `Iω` should be as returned by [`getIω`](@ref) with `specaxis = :ω`.

Returns the new specaxis grid and smoothed spectrum.
"""
function specres(ω, Iω, specaxis, resolution, specrange; window=nothing, nsamples=10)
    if isnothing(window)
        window = let ng=Maths.gaussnorm(fwhm=resolution), σ=resolution/(2*(2*log(2))^(1/2))
            (x,x0) -> exp(-0.5*((x - x0)/σ)^2)/ng
        end
    end
    if specaxis == :λ
        xg, Ix = _specres(ω, Iω, resolution, specrange, window, nsamples, wlfreq, wlfreq)
    elseif specaxis == :f
        xg, Ix = _specres(ω, Iω, resolution, specrange, window, nsamples, x -> x/(2π), x -> x*(2π))
    else
        error("`specaxis` must be one of `:λ` or `:f`")
    end
    xg, Ix
end

function _specres(ω, Iω, resolution, xrange, window, nsamples, ωtox, xtoω)
    # build output grid and array
    x = ωtox.(ω)
    fxrange = extrema(x[(x .> 0) .& isfinite.(x)])
    if isnothing(xrange)
        xrange = fxrange
    else
        xrange = extrema(xrange)
        xrange = (max(xrange[1], fxrange[1]), min(xrange[2], fxrange[2]))
    end
    nxg = ceil(Int, (xrange[2] - xrange[1])/resolution*nsamples)
    xg = collect(range(xrange[1], xrange[2], length=nxg))
    rdims = size(Iω)[2:end]
    Ix = Array{Float64, ndims(Iω)}(undef, ((nxg,)..., rdims...))
    fill!(Ix, 0.0)
    cidcs = CartesianIndices(rdims)
    # we find a suitable nspan
    nspan = 1
    while window(nspan*resolution, 0.0)/window(0.0, 0.0) > 1e-8
        nspan += 1
    end
    # now we build arrays of start and end indices for the relevant frequency
    # band for each output. For a frequency grid this is a little inefficient
    # but for a wavelength grid, which has varying index ranges, this is essential
    # and I think having a common code is simpler/cleaner.
    istart = Array{Int,1}(undef,nxg)
    iend = Array{Int,1}(undef,nxg)
    δω = ω[2] - ω[1]
    i0 = argmin(abs.(ω))
    ωs = ω[i0]
    for i in 1:nxg
        i1 = i0 + round(Int, (xtoω(xg[i] + resolution*nspan) - ωs)/δω)
        i2 = i0 + round(Int, (xtoω(xg[i] - resolution*nspan) - ωs)/δω)
        # we want increasing indices
        if i1 > i2
            i1,i2 = i2,i1
        end
        # handle boundaries
        if i2 > length(ω)
            i2 = length(ω)
        end
        if i1 < i0
            i1 = i0
        end
        istart[i] = i1
        iend[i] = i2
    end
    # run the convolution kernel - the function barrier massively improves performance
    _specres_kernel!(Ix, cidcs, istart, iend, Iω, window, x, xg, δω)
    xg, Ix
end

"""
Convolution kernel for each output point. We simply loop over all outer indices
and output points. The inner loop adds up the contributions from the specified window
around the target point. Note that this works without scaling also for wavelength ranges
because the integral is still over a frequency grid (with appropriate frequency dependent
integration bounds).
"""
function _specres_kernel!(Ix, cidcs, istart, iend, Iω, window, x, xg, δω)
    @inbounds @fastmath for ii in cidcs
        for j in 1:size(Ix, 1)
            for k in istart[j]:iend[j]
                Ix[j,ii] += Iω[k,ii] * window(x[k], xg[j]) * δω
            end
        end
    end
    Ix[Ix .<= 0.0] .= minimum(Ix[Ix .> 0.0])
end

function _specrangeselect(x, Ix; specrange=nothing, sortx=false)
    cidcs = CartesianIndices(size(Ix)[2:end])
    if !isnothing(specrange)
        specrange = extrema(specrange)
        idcs = ((x .>= specrange[1]) .& (x .<= specrange[2]))
        x = x[idcs]
        Ix = Ix[idcs, cidcs]
    end
    if sortx
        idcs = sortperm(x)
        x = x[idcs]
        Ix = Ix[idcs, cidcs]
    end
    x, Ix
end

"""
    ωwindow_λ(ω, λlims; winwidth=:auto)

Create a ω-axis filtering window to filter in `λlims`. `winwidth`, if a `Number`, sets
the smoothing width of the window in rad/s.
"""
function ωwindow_λ(ω, λlims; winwidth=:auto)
    ωmin, ωmax = extrema(wlfreq.(λlims))
    winwidth == :auto && (winwidth = 64*abs(ω[2] - ω[1]))
    window = Maths.planck_taper(ω, ωmin-winwidth, ωmin, ωmax, ωmax+winwidth)
end

"""
    getIω(ω, Eω, specaxis; specrange=nothing, resolution=nothing)

Get spectral energy density and x-axis given a frequency array `ω` and frequency-domain field
`Eω`, assumed to be correctly normalised (see [`getEω`](@ref)). `specaxis` determines the
x-axis:

- :f -> x-axis is frequency in Hz and Iω is in J/Hz
- :ω -> x-axis is angular frequency in rad/s and Iω is in J/(rad/s)
- :λ -> x-axis is wavelength in m and Iω is in J/m

# Keyword arguments
- `specrange::Tuple` can be set to a pair of limits on the spectral range (in `specaxis` units).
- `resolution::Real` is set, smooth the spectral energy density as defined by [`specres`](@ref).

Note that if `resolution` and `specaxis=:λ` is set it is highly recommended to also set `specrange`.
"""
function getIω(ω, Eω, specaxis; specrange=nothing, resolution=nothing)
    sortx = false
    if specaxis == :ω || !isnothing(resolution)
        specx = ω
        Ix = abs2.(Eω)
        if !isnothing(resolution)
            return specres(ω, Ix, specaxis, resolution, specrange)
        end
    elseif specaxis == :f
        specx = ω./2π
        Ix = abs2.(Eω)*2π
    elseif specaxis == :λ
        specx = wlfreq.(ω)
        Ix = @. ω^2/(2π*c) * abs2.(Eω)
        sortx = true
    else
        error("Unknown specaxis $specaxis")
    end
    if !isnothing(specrange) || sortx
        specx, Ix = _specrangeselect(specx, Ix, specrange=specrange, sortx=sortx)
    end
    return specx, Ix
end

"""
    getIω(output, specaxis[, zslice]; kwargs...)

Calculate the correctly normalised frequency-domain field and convert it to spectral
energy density on x-axis `specaxis` (`:f`, `:ω`, or `:λ`). If `zslice` is given,
returs only the slices of `Eω` closest to the given distances. `zslice` can be a single
number or an array. `specaxis` determines the
x-axis:

- :f -> x-axis is frequency in Hz and Iω is in J/Hz
- :ω -> x-axis is angular frequency in rad/s and Iω is in J/(rad/s)
- :λ -> x-axis is wavelength in m and Iω is in J/m

# Keyword arguments
- `specrange::Tuple` can be set to a pair of limits on the spectral range (in `specaxis` units).
- `resolution::Real` is set, smooth the spectral energy density as defined by [`specres`](@ref).

Note that `resolution` is set and `specaxis=:λ` it is highly recommended to also set `specrange`.
"""
getIω(output::AbstractOutput, specaxis; kwargs...) = getIω(getEω(output)..., specaxis; kwargs...)

function getIω(output::AbstractOutput, specaxis, zslice; kwargs...)
    ω, Eω, zactual = getEω(output, zslice)
    specx, Iω = getIω(ω, Eω, specaxis; kwargs...)
    return specx, Iω, zactual
end

"""
    getEω(output[, zslice])

Get frequency-domain modal field from `output` with correct normalisation (i.e.
`abs2.(Eω)`` gives angular-frequency spectral energy density in J/(rad/s)).
"""
getEω(output::AbstractOutput, args...) = getEω(makegrid(output), output, args...)
getEω(grid, output) = getEω(grid, output["Eω"])

function getEω(grid::RealGrid, Eω::AbstractArray)
    ω = grid.ω[grid.sidx]
    Eω = Eω[grid.sidx, CartesianIndices(size(Eω)[2:end])]
    return ω, Eω*fftnorm(grid)
end

function getEω(grid::EnvGrid, Eω::AbstractArray)
    idcs = FFTW.fftshift(grid.sidx)
    Eωs = FFTW.fftshift(Eω, 1)
    ω = FFTW.fftshift(grid.ω)[idcs]
    Eω = Eωs[idcs, CartesianIndices(size(Eω)[2:end])]
    return ω, Eω*fftnorm(grid)
end

function getEω(grid, output, zslice)
    zidx = nearest_z(output, zslice)
    ω, Eω = getEω(grid, output["Eω", .., zidx])
    return ω, Eω, output["z"][zidx]
end

fftnorm(grid::RealGrid) = Maths.rfftnorm(grid.t[2] - grid.t[1])
fftnorm(grid::EnvGrid) = Maths.fftnorm(grid.t[2] - grid.t[1])

"""
    getEt(output[, zslice]; kwargs...)

Get the envelope time-domain electric field (including the carrier wave) from the `output`.
If `zslice` is given, returs only the slices of `Eω` closest to the given distances. `zslice`
can be a single number or an array.
"""
getEt(output::AbstractOutput, args...; kwargs...) = getEt(
    makegrid(output), output, args...; kwargs...)

"""
    getEt(grid, Eω; trange=nothing, oversampling=4, bandpass=nothing, FTL=false)

Get the envelope time-domain electric field (including the carrier wave) from the frequency-
domain field `Eω`. The field can be cropped in time using `trange`, it is oversampled by
a factor of `oversampling` (default 4) and can be bandpassed with `bandpass`
(see [`window_maybe`](@ref)). If `FTL` is `true`, return the Fourier-transform limited pulse,
i.e. remove any spectral phase.

If `zslice` is given, returs only the slices of `Eω` closest to the given distances. `zslice`
can be a single number or an array.
"""
function getEt(grid::AbstractGrid, Eω::AbstractArray;
               trange=nothing, oversampling=4, bandpass=nothing,
               FTL=false, propagate=nothing)
    t = grid.t
    Eω = window_maybe(grid.ω, Eω, bandpass)
    if FTL
        τ = length(grid.t) * (grid.t[2] - grid.t[1])/2
        Eω .= abs.(Eω) .* exp.(-1im .* grid.ω .* τ)
    end
    Eω = prop_maybe(grid, Eω, propagate)
    Etout = envelope(grid, Eω)
    if isnothing(trange)
        idcs = 1:length(t)
    else
        idcs = @. (t < max(trange...)) & (t > min(trange...))
    end
    to, Eto = Maths.oversample(t[idcs], Etout[idcs, ..], factor=oversampling)
    return to, Eto
end

getEt(grid::AbstractGrid, output::AbstractOutput; kwargs...) = getEt(grid, output["Eω"]; kwargs...)

function getEt(grid::AbstractGrid, output::AbstractOutput, zslice;
               kwargs...)
    t = grid.t
    zidx = nearest_z(output, zslice)
    to, Eto = getEt(grid, output["Eω", .., zidx]; kwargs...)
    return to, Eto, output["z"][zidx]
end

"""
    AutoWindow(width, λmin, λmax, ω0fun; relative=false, ndims=1)

Window function generator which automatically tracks the central frequency in the spectral
region given by `λmin` and `λmax` and applies a window of a specific `width` around the peak.
The central frequency is found using the function `ω0fun(ω, Iω::AbstractVector)`, where
`ω` and `Iω` are already cropped to within the wavelength limits given.
If `relative` is `true`, `width` is relative bandwidth instead of the wavelength width.
`ndims` determines how many dimensions of the array to sum over. For a field array with size
`(Nω, N1, N2, ...)`, the first dimension is always assumed to be frequency. `ndim=1` means
each field to be analysed is 1-dimensional, so the window iterates over all of `(N1, N2, ...)`.
`ndim=2` means each field to be analysed is 2-dimensional, `(Nω, N1)` in size, and will be
summed over its second dimension before finding the central frequency. The window iterates
over all other dimensions, `(N2, ...)`.

A `AutoWindow` automatically stores the limits of the windows it applies in the field `lims`.
"""
mutable struct AutoWindow
    width::Float64
    λmin::Float64
    λmax::Float64
    ω0fun
    relative::Bool
    ndims
    lims
end

function AutoWindow(width, λmin, λmax, ω0fun; relative=false, ndims=1)
    AutoWindow(width, λmin, λmax, ω0fun, relative, ndims, nothing)
end

function (pw::AutoWindow)(ω, Eω)
    cidcs = CartesianIndices(size(Eω)[(pw.ndims+1):end])
    out = similar(Eω)
    cropidcs = (ω .> wlfreq(pw.λmax)) .& (ω .< wlfreq(pw.λmin))
    cropω = ω[cropidcs]
    Iω = abs2.(Eω)
    limsA = zeros((2, size(Eω)[(pw.ndims+1):end]...))
    for cidx in cidcs
        Iω_this = Iω[.., Tuple(cidx)...]
        Iωsum = sum(Iω_this; dims=2:ndims(Iω_this))
        λ0 = wlfreq(pw.ω0fun(cropω, Iωsum[cropidcs]))
        lims = pw.relative ? λ0.*(1 .+ (-0.5, 0.5).*pw.width) : λ0 .+ (-0.5, 0.5).*pw.width
        window = ωwindow_λ(ω, lims)
        limsA[:, Tuple(cidx)...] .= lims
        out[.., Tuple(cidx)...] .= Eω[.., Tuple(cidx)...] .* window
    end
    pw.lims = limsA
    out
end

"""
    PeakWindow(width, λmin, λmax; relative=false, ndims=1)

An [`AutoWindow`](@ref) which uses the peak of the spectral energy density as the central
frequency.
"""
function PeakWindow(width, λmin, λmax; relative=false, ndims=1)
    ω0fun = (ω, Iω) ->  ω[argmax(Iω)]
    AutoWindow(width, λmin, λmax, ω0fun; relative=relative, ndims=ndims)
end

"""
    CentroidWindow(width, λmin, λmax; relative=false, ndims=1, power=1)

An [`AutoWindow`](@ref) which uses the centroid (centre of mass or first moment) of the
spectral energy density as the central frequency. Before calculating the centroid, the
SED is raised to the `power` given.
"""
function CentroidWindow(width, λmin, λmax; relative=false, ndims=1, power=1)
    ω0fun = (ω, Iω) -> Maths.moment(ω, Iω.^power)
    AutoWindow(width, λmin, λmax, ω0fun; relative=relative, ndims=ndims)
end

"""
    window_maybe(ω, Eω, win)

Apply a frequency window to the field `Eω` if required. Possible values for `win`:

- `nothing` : no window is applied
- 4-`Tuple` of `Number`s : the 4 parameters for a `Maths.planck_taper` in **wavelength**
- 3-`Tuple` of `Number`s : minimum, maximum **wavelength**, and smoothing in **radial frequency**
- 2-`Tuple` of `Number`s : minimum and maximum **wavelength** with automatically chosen smoothing
- `Vector{<:Real}` : a pre-defined window function (shape must match `ω`)
- `PeakWindow` : automatically track the peak in a given range and apply the window around it
- `window(ω, Eω)` : an arbitrary user-supplied window function
"""
window_maybe(ω, Eω, ::Nothing) = Eω
window_maybe(ω, Eω, win::NTuple{4, Number}) = Eω.*Maths.planck_taper(
    ω, sort(wlfreq.(collect(win)))...)
window_maybe(ω, Eω, win::NTuple{2, Number}) = Eω .* ωwindow_λ(ω, win)
window_maybe(ω, Eω, win::NTuple{3, Number}) = Eω .* ωwindow_λ(ω, win[1:2]; winwidth=win[3])
window_maybe(ω, Eω, window) = window(ω, Eω)
window_maybe(ω, Eω, window::AbstractVector) = Eω.*window

prop_maybe(grid, Eω, ::Nothing) = Eω
prop_maybe(grid, Eω, propagator) = propagator(grid, Eω)

"""
    envelope(grid, Eω)

Get the envelope electric field including the carrier wave from the frequency-domain field
`Eω` sampled on `grid`.
"""
envelope(grid::RealGrid, Eω) = Maths.hilbert(FFTW.irfft(Eω, length(grid.t), 1))
envelope(grid::EnvGrid, Eω) = FFTW.ifft(Eω, 1) .* exp.(im.*grid.ω0.*grid.t)

"""
    makegrid(output)

Create an `AbstractGrid` from the `"grid"` dictionary saved in `output`.
"""
function makegrid(output)
    if output["simulation_type"]["field"] == "field-resolved"
        from_dict(RealGrid, output["grid"])
    else
        from_dict(EnvGrid, output["grid"])
    end
end

"""
    nearest_z(output, z)

Return the index of saved z-position(s) closest to the position(s) `z`. Output is always
an array, even if `z` is a number.
"""
nearest_z(output, z::Number) = [argmin(abs.(output["z"] .- z))]
nearest_z(output, z) = [argmin(abs.(output["z"] .- zi)) for zi in z]

end


function auglag_kernel(n::Int, major_iter::Int, max_auglag::Int,
                       line_start::Int, scale::Float64,
                       mu_max::Float64,
                       u_curr::CuDeviceArray{Float64,1}, v_curr::CuDeviceArray{Float64,1},
                       l_curr::CuDeviceArray{Float64,1}, rho::CuDeviceArray{Float64,1},
                       wRIij::CuDeviceArray{Float64,1},
                       param::CuDeviceArray{Float64,2},
                       _YffR::CuDeviceArray{Float64,1}, _YffI::CuDeviceArray{Float64,1},
                       _YftR::CuDeviceArray{Float64,1}, _YftI::CuDeviceArray{Float64,1},
                       _YttR::CuDeviceArray{Float64,1}, _YttI::CuDeviceArray{Float64,1},
                       _YtfR::CuDeviceArray{Float64,1}, _YtfI::CuDeviceArray{Float64,1},
                       frBound::CuDeviceArray{Float64,1}, toBound::CuDeviceArray{Float64,1})
    tx = threadIdx().x
    ty = threadIdx().y
    I = blockIdx().x

    x = @cuDynamicSharedMem(Float64, n)
    xl = @cuDynamicSharedMem(Float64, n, n*sizeof(Float64))
    xu = @cuDynamicSharedMem(Float64, n, (2*n)*sizeof(Float64))

    pij_idx = line_start + 8*(I-1)

    if tx == 1 && ty == 1
        @inbounds for i=1:n
            xl[i] = -Inf
            xu[i] = Inf
        end

        @inbounds begin
            xl[5] = frBound[2*(I-1)+1]
            xu[5] = frBound[2*I]
            xl[6] = toBound[2*(I-1)+1]
            xu[6] = toBound[2*I]
            xl[9] = -2*pi
            xu[9] = 2*pi
            xl[10] = -2*pi
            xu[10] = 2*pi

            x[1] = u_curr[pij_idx]
            x[2] = u_curr[pij_idx+1]
            x[3] = u_curr[pij_idx+2]
            x[4] = u_curr[pij_idx+3]
            x[5] = min(xu[5], max(xl[5], u_curr[pij_idx+4]))
            x[6] = min(xu[6], max(xl[6], u_curr[pij_idx+5]))
            x[7] = wRIij[2*(I-1)+1]
            x[8] = wRIij[2*I]
            x[9] = min(xu[9], max(xl[9], u_curr[pij_idx+6]))
            x[10] = min(xu[10], max(xl[10], u_curr[pij_idx+7]))
        end
    end

    @inbounds begin
        YffR = _YffR[I]; YffI = _YffI[I]
        YftR = _YftR[I]; YftI = _YftI[I]
        YttR = _YttR[I]; YttI = _YttI[I]
        YtfR = _YtfR[I]; YtfI = _YtfI[I]

        param[1,I] = l_curr[pij_idx]
        param[2,I] = l_curr[pij_idx+1]
        param[3,I] = l_curr[pij_idx+2]
        param[4,I] = l_curr[pij_idx+3]
        param[5,I] = l_curr[pij_idx+4]
        param[6,I] = l_curr[pij_idx+5]
        param[7,I] = rho[pij_idx]
        param[8,I] = rho[pij_idx+1]
        param[9,I] = rho[pij_idx+2]
        param[10,I] = rho[pij_idx+3]
        param[11,I] = rho[pij_idx+4]
        param[12,I] = rho[pij_idx+5]
        param[13,I] = v_curr[pij_idx]
        param[14,I] = v_curr[pij_idx+1]
        param[15,I] = v_curr[pij_idx+2]
        param[16,I] = v_curr[pij_idx+3]
        param[17,I] = v_curr[pij_idx+4]
        param[18,I] = v_curr[pij_idx+5]

        param[25,I] = l_curr[pij_idx+6]
        param[26,I] = l_curr[pij_idx+7]
        param[27,I] = rho[pij_idx+6]
        param[28,I] = rho[pij_idx+7]
        param[29,I] = v_curr[pij_idx+6]
        param[30,I] = v_curr[pij_idx+7]
    end

    if major_iter == 1
        @inbounds param[24,I] = 10.0
        mu = 10.0
    else
        @inbounds mu = param[24,I]
    end

    CUDA.sync_threads()

    eta = 1 / CUDA.pow(mu, 0.1)
    omega = 1 / mu
    max_feval = 500
    max_minor = 200
    gtol = 1e-6

    it = 0
    terminate = false

    while !terminate
        it += 1

        # Solve the branch problem.
        status, minor_iter = tron_kernel(n, 0, max_feval, max_minor, gtol, scale, false, x, xl, xu,
                                         param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)

        @inbounds begin
            # Check the termination condition.
            cviol1 = x[1] - (YffR*x[5] + YftR*x[7] + YftI*x[8])
            cviol2 = x[2] - (-YffI*x[5] - YftI*x[7] + YftR*x[8])
            cviol3 = x[3] - (YttR*x[6] + YtfR*x[7] - YtfI*x[8])
            cviol4 = x[4] - (-YttI*x[6] - YtfI*x[7] - YtfR*x[8])
            cviol5 = x[7]^2 + x[8]^2 - x[5]*x[6]
            cviol6 = x[9] - x[10] - atan(x[8], x[7])
        end

        cnorm = max(abs(cviol1), abs(cviol2), abs(cviol3), abs(cviol4), abs(cviol5), abs(cviol6))

        if cnorm <= eta
            if cnorm <= 1e-6
                terminate = true
            else
                if tx == 1 && ty == 1
                    @inbounds begin
                        param[19,I] += mu*cviol1
                        param[20,I] += mu*cviol2
                        param[21,I] += mu*cviol3
                        param[22,I] += mu*cviol4
                        param[23,I] += mu*cviol5
                        param[31,I] += mu*cviol6
                    end
                end

                eta = eta / CUDA.pow(mu, 0.9)
                omega  = omega / mu
            end
        else
            mu = min(mu_max, mu*10)
            eta = 1 / CUDA.pow(mu, 0.1)
            omega = 1 / mu
            @inbounds param[24,I] = mu
        end

        if it >= max_auglag
            terminate = true
        end

        CUDA.sync_threads()
    end

    @inbounds begin
        u_curr[pij_idx] = x[1]
        u_curr[pij_idx+1] = x[2]
        u_curr[pij_idx+2] = x[3]
        u_curr[pij_idx+3] = x[4]
        u_curr[pij_idx+4] = x[5]
        u_curr[pij_idx+5] = x[6]
        wRIij[2*(I-1)+1] = x[7]
        wRIij[2*I] = x[8]
        u_curr[pij_idx+6] = x[9]
        u_curr[pij_idx+7] = x[10]
        param[24,I] = mu
    end

    CUDA.sync_threads()

    return
end

function auglag_kernel_cpu(n::Int, nline::Int, major_iter::Int, max_auglag::Int,
                           line_start::Int, mu_max::Float64,
                           u_curr::Array{Float64}, v_curr::Array{Float64},
                           l_curr::Array{Float64}, rho::Array{Float64},
                           wRIij::Array{Float64},
                           param::Array{Float64},
                           YffR::Array{Float64}, YffI::Array{Float64},
                           YftR::Array{Float64}, YftI::Array{Float64},
                           YttR::Array{Float64}, YttI::Array{Float64},
                           YtfR::Array{Float64}, YtfI::Array{Float64},
                           frBound::Array{Float64}, toBound::Array{Float64})
    avg_auglag_it = 0
    avg_minor_it = 0

    x = zeros(n)
    xl = zeros(n)
    xu = zeros(n)

    @inbounds for I=1:nline
        @inbounds for i=1:n
            xl[i] = -Inf
            xu[i] = Inf
        end

        xl[5] = frBound[2*(I-1)+1]
        xu[5] = frBound[2*I]
        xl[6] = toBound[2*(I-1)+1]
        xu[6] = toBound[2*I]
        xl[9] = -2*pi
        xu[9] = 2*pi
        xl[10] = -2*pi
        xu[10] = 2*pi

        pij_idx = line_start + 8*(I-1)

        x[1] = u_curr[pij_idx]
        x[2] = u_curr[pij_idx+1]
        x[3] = u_curr[pij_idx+2]
        x[4] = u_curr[pij_idx+3]
        x[5] = min(xu[5], max(xl[5], u_curr[pij_idx+4]))
        x[6] = min(xu[6], max(xl[6], u_curr[pij_idx+5]))
        x[7] = wRIij[2*(I-1)+1]
        x[8] = wRIij[2*I]
        x[9] = min(xu[9], max(xl[9], u_curr[pij_idx+6]))
        x[10] = min(xu[10], max(xl[10], u_curr[pij_idx+7]))

        param[1,I] = l_curr[pij_idx]
        param[2,I] = l_curr[pij_idx+1]
        param[3,I] = l_curr[pij_idx+2]
        param[4,I] = l_curr[pij_idx+3]
        param[5,I] = l_curr[pij_idx+4]
        param[6,I] = l_curr[pij_idx+5]
        param[7,I] = rho[pij_idx]
        param[8,I] = rho[pij_idx+1]
        param[9,I] = rho[pij_idx+2]
        param[10,I] = rho[pij_idx+3]
        param[11,I] = rho[pij_idx+4]
        param[12,I] = rho[pij_idx+5]
        param[13,I] = v_curr[pij_idx]
        param[14,I] = v_curr[pij_idx+1]
        param[15,I] = v_curr[pij_idx+2]
        param[16,I] = v_curr[pij_idx+3]
        param[17,I] = v_curr[pij_idx+4]
        param[18,I] = v_curr[pij_idx+5]

        if major_iter == 1
            param[24,I] = 10.0
            mu = 10.0
        else
            mu = param[24,I]
        end

        param[25,I] = l_curr[pij_idx+6]
        param[26,I] = l_curr[pij_idx+7]
        param[27,I] = rho[pij_idx+6]
        param[28,I] = rho[pij_idx+7]
        param[29,I] = v_curr[pij_idx+6]
        param[30,I] = v_curr[pij_idx+7]

        function eval_f_cb(x)
            f= eval_f_kernel_cpu(I, x, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            return f
        end

        function eval_g_cb(x, g)
            eval_grad_f_kernel_cpu(I, x, g, param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            return
        end

        function eval_h_cb(x, mode, rows, cols, scale, lambda, values)
            eval_h_kernel_cpu(I, x, mode, scale, rows, cols, lambda, values,
                              param, YffR, YffI, YftR, YftI, YttR, YttI, YtfR, YtfI)
            return
        end

        eta = 1 / mu^0.1
        omega = 1 / mu
        max_feval = 500
        max_minor = 100
        gtol = 1e-6

        nele_hess = 33
        tron = ExaTron.createProblem(n, xl, xu, nele_hess, eval_f_cb, eval_g_cb, eval_h_cb;
                                     :tol => gtol, :matrix_type => :Dense, :max_minor => 200,
                                     :frtol => 1e-12)
        tron.x .= x
        it = 0
        avg_tron_minor = 0
        terminate = false

        while !terminate
            it += 1

            # Solve the branch problem.
            status = ExaTron.solveProblem(tron)
            x .= tron.x
            avg_tron_minor += tron.minor_iter

            # Check the termination condition.
            cviol1 = x[1] - (YffR[I]*x[5] + YftR[I]*x[7] + YftI[I]*x[8])
            cviol2 = x[2] - (-YffI[I]*x[5] - YftI[I]*x[7] + YftR[I]*x[8])
            cviol3 = x[3] - (YttR[I]*x[6] + YtfR[I]*x[7] - YtfI[I]*x[8])
            cviol4 = x[4] - (-YttI[I]*x[6] - YtfI[I]*x[7] - YtfR[I]*x[8])
            cviol5 = x[7]^2 + x[8]^2 - x[5]*x[6]
            cviol6 = x[9] - x[10] - atan(x[8], x[7])

            cnorm = max(abs(cviol1), abs(cviol2), abs(cviol3), abs(cviol4), abs(cviol5), abs(cviol6))

            if cnorm <= eta
                if cnorm <= 1e-6
                    terminate = true
                else
                    param[19,I] += mu*cviol1
                    param[20,I] += mu*cviol2
                    param[21,I] += mu*cviol3
                    param[22,I] += mu*cviol4
                    param[23,I] += mu*cviol5
                    param[31,I] += mu*cviol6

                    eta = eta / mu^0.9
                    omega  = omega / mu
                end
            else
                mu = min(mu_max, mu*10)
                eta = 1 / mu^0.1
                omega = 1 / mu
                param[24,I] = mu
            end

            if it >= max_auglag
                terminate = true
            end
        end

        avg_auglag_it += it
        avg_minor_it += (avg_tron_minor / it)
        u_curr[pij_idx] = x[1]
        u_curr[pij_idx+1] = x[2]
        u_curr[pij_idx+2] = x[3]
        u_curr[pij_idx+3] = x[4]
        u_curr[pij_idx+4] = x[5]
        u_curr[pij_idx+5] = x[6]
        wRIij[2*(I-1)+1] = x[7]
        wRIij[2*I] = x[8]
        u_curr[pij_idx+6] = x[9]
        u_curr[pij_idx+7] = x[10]
        param[24,I] = mu
    end

    return (avg_auglag_it / nline), (avg_minor_it / nline)
end


module DiskArrayTools
import DiskArrays: AbstractDiskArray, eachchunk, haschunks, Chunked,
estimate_chunksize, GridChunks, findints, readblock!, writeblock!
using Interpolations
using IterTools: imap
using Base.Iterators: product
using OffsetArrays: OffsetArray

#Define a new chunk grid type for irregular, but still block-shaped grid
struct IrregularGridChunks{N} <: AbstractArray{CartesianIndices{N,NTuple{N,UnitRange{Int}}},N}
    cs::NTuple{N,Vector{UnitRange{Int}}}
    s::NTuple{N,Int}
end
Base.size(cs::IrregularGridChunks) = cs.s
Base.size(cs::IrregularGridChunks,i) = cs.s[i]
function IrregularGridChunks(cs::Vector{UnitRange{Int}}...)
    s = length.(cs)
    IrregularGridChunks((cs...,),s)
end
Base.eltype(::IrregularGridChunks{N}) where N = CartesianIndices{N,NTuple{N,UnitRange{Int}}}
function Base.show(io::IO,g::IrregularGridChunks)
    stot = last.(g.cs)
    griddes = join(string.(g.s), "x")
    sizedes = join(string.(stot), "x")
    print(io,"Irregular chunk grid of size $griddes over DiskArray of size $stot")
end
function Base.getindex(cs::IrregularGridChunks{N},i::Vararg{Int, N}) where N
    rout = getindex.(cs.cs,i)
    CartesianIndices(rout)
end

export DiskArrayStack, diskstack, ConcatDiskArray, CFDiskArray
struct DiskArrayStack{T,N,M,NO}<:AbstractDiskArray{T,N}
    arrays::Array{M,NO}
end
function diskstack(a::Array{M,N}) where {N,M<:AbstractArray{T,NO}} where {T,NO}
    length(unique(size.(a))) == 1 || error("All arrays in the stack must have the same size")
    DiskArrayStack{T,N+NO,M,N}(a)
end
Base.size(r::DiskArrayStack) = (size(r.arrays[1])...,size(r.arrays)...)
haschunks(a::DiskArrayStack) = haschunks(a.arrays[1])
iscompressed(a::DiskArrayStack) = any(iscompressed,a.arrays)
function eachchunk(a::DiskArrayStack{<:Any,<:Any,<:Any,NO}) where NO
    iterold = eachchunk(a.arrays[1])
    cs = (iterold.chunksize...,ntuple(one,NO)...)
    co = (iterold.offset...,ntuple(zero,NO)...)
    GridChunks(a,cs,offset=co)
end

function readblock!(a::DiskArrayStack{<:Any,N,<:Any,NO},aout,i::AbstractVector...) where {N,NO}

  innerinds = i[1:(N-NO)]

  outerinds = i[(N-NO+1):N]
  innercolon = map(_->(:), innerinds)
  iiter = CartesianIndices(outerinds)
  inum  = CartesianIndices(size(iiter))
  foreach(zip(iiter,inum)) do (iouter,iret)
    arnow = a.arrays[iouter]
    if isa(arnow, AbstractDiskArray)
      readblock!(a.arrays[iouter],view(aout,innercolon...,iret.I...),innerinds...)
    else
      aout[innercolon...,iret.I...] = arnow[innerinds...]
    end
  end
  nothing
end

function writeblock!(a::DiskArrayStack{<:Any,N,<:Any,NO},v,i::AbstractVector...) where {N,NO}
  innerinds = i[1:(N-NO)]

  outerinds = i[(N-NO+1):N]
  innercolon = map(_->(:), innerinds)
  iiter = CartesianIndices(outerinds)
  inum  = CartesianIndices(size(iiter))
  foreach(zip(iiter,inum)) do (iouter,iret)
    arnow = a.arrays[iouter]
    if isa(arnow, AbstractDiskArray)
      writeblock!(a.arrays[iouter],view(v,innercolon...,iret.I...),innerinds...)
    else
      arnow[innerinds...] = v[innercolon...,iret.I...]
    end
  end
  nothing
end

function Base.view(a::DiskArrayStack{<:Any,N,<:Any,NO},i...) where {N,NO}
    iinner = i[1:(N-NO)]
    ashort = map(view(a.arrays,i[(N-NO+1):N]...)) do ai
        view(ai,iinner...)
    end
    if ndims(ashort)==0
      ashort[]
    else
      diskstack(ashort)
    end
end

abstract type ResampledDiskArray{T,N} <: AbstractDiskArray{T,N} end

function readblock!(a::ResampledDiskArray, aout, i::AbstractUnitRange...)
  parentranges = map((ir,ip)->ip[ir],i,a.newinds)
  rr = get_readinds(a,parentranges)
  atemp = a.a[rr...]
  atemp2 = OffsetArray(atemp,map(r->(first(r)-1),rr))
  resample_disk(a,aout,atemp2,parentranges)
end

struct InterpolatedDiskArray{T,N,A<:AbstractArray{T,N},I,O,BC,CS<:Union{Nothing,GridChunks{N}}} <: ResampledDiskArray{T,N}
    a::A
    newinds::I
    meth::O
    bc::BC
    chunksize::CS
end
fixin(i,s) = max(1,min(s,i))
function round_readinds(s,r,an)
  mi,ma = extrema(r)
  if mi == ma
    m = fixin(round(Int,mi),s)
    m:m
  else
    a = if an
      (round(Int,mi,RoundNearestTiesUp), -round(Int,-ma,RoundNearestTiesUp))
    else
      (floor(Int,mi),ceil(Int,ma))
    end
    fixin(a[1],s):fixin(a[2],s)
  end
end

function get_readinds(a::InterpolatedDiskArray,r)
  allnearest = all(i->isa(i,Union{Nothing,BSpline{Constant}}),a.meth)
  round_readinds.(size(a.a),r,allnearest)
end

allmeths(order::Tuple,newinds) = map(order,newinds) do o,ni
    ni === nothing ? NoInterp() : BSpline(o)
end
function allmeths(order, newinds)
    allorders = map(newinds) do _
        order
    end
    allmeths(allorders,newinds)
end
function InterpolatedDiskArray(a::AbstractArray,chunksize,newinds...; order=Linear(), bc=Flat())
    eltype(a) <: Union{Missing,Real,Complex} || error("Can only interpolate real or complex values")
    s = size(a)
    ni2 = map((s,i)->i===nothing ? (1:s) : i,s,newinds)
    me = allmeths(order,newinds)
    InterpolatedDiskArray(a,ni2,me,bc,chunksize)
end
Base.size(a::InterpolatedDiskArray) = map(length,a.newinds)
haschunks(a::InterpolatedDiskArray{<:Any,<:Any,<:Any,<:Any,<:Any,<:Any,<:GridChunks}) = true
haschunks(a::InterpolatedDiskArray{<:Any,<:Any,<:Any,<:Any,<:Any,<:Any,Nothing}) = false
eachchunk(a::InterpolatedDiskArray{<:Any,<:Any,<:Any,<:Any,<:Any,<:Any,<:GridChunks}) = a.chunksize

function resample_disk(a::InterpolatedDiskArray,aout,atemp,parentranges)
  meth = map((s,m)->s==1 ? NoInterp() : m,size(atemp),a.meth)
  if all(isequal(NoInterp()),meth)
      aout .= atemp.parent
  else
      fremap(aout,atemp,parentranges,meth,bc=a.bc)
  end
end

function fremap(xout,xin,newinds,ipmeth;bc = Flat())
  interp = extrapolate(interpolate(xin, ipmeth), bc)
  for (i,ic) in enumerate(Iterators.product(newinds...))
    xout[i]=interp(ic...)
  end
end

#Use of Sentinel missing value
struct CFDiskArray{T,N,ST,P<:AbstractArray{ST,N}} <: AbstractDiskArray{Union{T,Missing},N}
    a::P
    mv::ST
    add_offset::T
    scale_factor::T
end
function CFDiskArray(a::AbstractArray{T}, attr::Dict) where T
  mv = get(attr,"missing_value",get(attr,"_FillValue",typemax(T)))
  offs,sc = if haskey(attr,"add_offset") || haskey(attr,"scale_factor")
    offs = get(attr,"add_offset",zero(Float16))
    sc = get(attr,"scale_factor",one(Float16))
    promote(offs,sc)
  else
    zero(T), one(T)
  end
  CFDiskArray(a, T(mv), offs, sc)
end

Base.size(a::CFDiskArray, args...) = size(a.a, args...)
haschunks(a::CFDiskArray) = haschunks(a.a)
eachchunk(a::CFDiskArray) = eachchunk(a.a)
iscompressed(a::CFDiskArray) = iscompressed(a.a)

function readblock!(a::CFDiskArray, aout, r::AbstractVector...)
    mv = a.mv
    sc,offs = a.scale_factor, a.add_offset
    broadcast!(j->scaleoffs(j,mv,sc,offs), aout, view(a.a,r...))
    nothing
end
scaleoffs(x,mv,sc,offs) = x==mv ? missing : x*sc+offs
function writeblock!(a::CFDiskArray, v, r::AbstractVector...)
    mv = a.mv
    sc,offs = a.scale_factor, a.add_offset
    broadcast!(j->scaleoffsinv(j,mv,sc,offs), view(a.a,r...), v)
    nothing
end
scaleoffsinv(x,mv::Integer,sc,offs) = ismissing(x) ? mv : round(typeof(mv),(x-offs)/sc)
scaleoffsinv(x,mv,sc,offs) = ismissing(x) ? mv : ((x-offs)/sc)

struct ConcatDiskArray{T,N,P} <: AbstractDiskArray{T,N}
    parents::P
    startinds::NTuple{N,Vector{Int}}
    size::NTuple{N,Int}
end

function ConcatDiskArray(arrays::AbstractArray)
    #First do some consistency checks
    arrinfo = map(arrays) do a
        eltype(a), size(a)
    end
    et = first(arrinfo)[1]
    nd = length(first(arrinfo)[2])
    sa = size.(arrays)
    all(i->i[1]==et,arrinfo) || error("Arrays don't have the same element type")
    all(i->length(i[2])==nd, arrinfo) || error("Arrays don't have the same dimensions")
    if ndims(arrays)<nd
        arrays = reshape(arrays,size(arrays)...,ones(Int,nd-ndims(arrays))...)
    end
    function othersize(x,id)
        (x[1:id-1]..., x[id+1:end]...)
    end
    si = map(1:nd) do id
        otherdims = ((1:id-1)...,(id+1:nd)...)
        a = size.(arrays)
        a = reduce(a, dims=id, init=ntuple(zero,nd)) do i,j
                if all(iszero,i)
                    j
                elseif othersize(i,id)==othersize(j,id)
                    j
                else
                    error("Dimensions don't match")
                end
        end
        I = fill!(Array{Union{Int,Colon},1}(undef,nd),1)
        I[id] = Colon()
        ar = sa[I...]
        ari = map(i->i[id],ar)
        sl = sum(ari)
        r = cumsum(ari)
        pop!(pushfirst!(r,0))
        r.+1, sl
    end
    ConcatDiskArray{et,nd,typeof(arrays)}(arrays, (getindex.(si,1)...,),(getindex.(si,2)...,))
end

Base.size(a::ConcatDiskArray) = a.size
haschunks(::ConcatDiskArray) = Chunked()
function readblock!(a::ConcatDiskArray, aout, inds::AbstractUnitRange...)
    #Find affected blocks and indices in blocks
    blockinds  = map(inds, a.startinds, size(a.parents)) do i, si, s
        bi1 = max(searchsortedlast(si,first(i)),1)
        bi2 = min(searchsortedfirst(si,last(i)+1)-1,s)
        bi1:bi2
    end
    map(CartesianIndices(blockinds)) do cI
        myar = a.parents[cI]
        mysize = size(myar)
        array_range = map(cI.I, a.startinds, mysize, inds) do ii, si, ms, indstoread
            max(first(indstoread)-si[ii]+1,1) : min(last(indstoread)-si[ii]+1, ms)
        end
        outer_range = map(cI.I, a.startinds, array_range, inds) do ii, si, ar, indstoread
            (first(ar)+si[ii]-first(indstoread)):(last(ar)+si[ii]-first(indstoread))
        end
        aout[outer_range...] .= a.parents[cI][array_range...]
    end
end
function writeblock!(a::ConcatDiskArray, aout, inds::AbstractUnitRange...)
  error("No method yet for writing into a ConcatDiskArray")
end

function eachchunk_fallback(aconc::ConcatDiskArray)
  nested = imap(zip(aconc.parents, Iterators.product(aconc.startinds...))) do (par, si)
      imap(eachchunk(par)) do cI
          cI .+ CartesianIndex(si.-1)
      end
  end
  Iterators.flatten(nested)
end

#This function will work for DiskArrayStack as well, we only need a fallback_chunksize function
function eachchunk(aconc::ConcatDiskArray)
  fbchunks = eachchunk_fallback(aconc)
  nd = ndims(first(fbchunks))
  dout = ntuple(_->Dict{NTuple{nd-1,UnitRange{Int}},Set{UnitRange{Int}}}(),nd)
  foreach(fbchunks) do c
    map(enumerate(c.indices)) do (ind,r)
        other = (c.indices[1:ind-1]...,c.indices[ind+1:end]...)
        s = get!(dout[ind],other) do
            Set{UnitRange{Int}}()
        end
        push!(s,c.indices[ind])
    end
  end
  #First check if all
  uniqueranges = unique.(values.(dout)) # Here we give up
  if any(length.(uniqueranges) .> 1)
    return fbchunks
  end
  simplifyaxes = map(uniqueranges) do ur
    allr = first(ur)
    allr = sort!(collect(allr),by=first)
    l = length.(allr)
    if length(l)==1
        allr,(cs = l[1], offs = 0, l = l[1])
    elseif all(isequal(l[2]),l[2:end-1]) && l[end]<=l[2]
        allr,(cs = l[2], offs = l[2]-l[1], l = allr[end][end])
    else
        allr,nothing
    end
  end
  if all(!isnothing,getindex.(simplifyaxes,2))
    grid = getindex.(simplifyaxes,2)
    cs,offs,l = getproperty.(grid,:cs),getproperty.(grid,:offs),getproperty.(grid,:l)
    return GridChunks(l,cs, offset = offs)
  else
    return IrregularGridChunks(getindex.(simplifyaxes,1)...)
  end
end
end # module


# This file is a part of Julia. License is MIT: http://julialang.org/license

#Period types
value(x::Period) = x.value

# The default constructors for Periods work well in almost all cases
# P(x) = new((convert(Int64,x))
# The following definitions are for Period-specific safety
for period in (:Year, :Month, :Week, :Day, :Hour, :Minute, :Second, :Millisecond)
    period_str = string(period)
    accessor_str = lowercase(period_str)
    # Convenience method for show()
    @eval _units(x::$period) = " " * $accessor_str * (abs(value(x)) == 1 ? "" : "s")
    # periodisless
    @eval periodisless(x::$period,y::$period) = value(x) < value(y)
    # AbstractString parsing (mainly for IO code)
    @eval $period(x::AbstractString) = $period(Base.parse(Int64,x))
    # Period accessors
    typ_str = period in (:Hour, :Minute, :Second, :Millisecond) ? "DateTime" : "TimeType"
    description = typ_str == "TimeType" ? "`Date` or `DateTime`" : "`$typ_str`"
    reference = period == :Week ? " For details see [`$accessor_str(::$typ_str)`](:func:`$accessor_str`)." : ""
    @eval begin
        @doc """
            $($period_str)(dt::$($typ_str)) -> $($period_str)

        The $($accessor_str) part of a $($description) as a `$($period_str)`.$($reference)
        """ ->
        $period(dt::$(Symbol(typ_str))) = $period($(Symbol(accessor_str))(dt))

        @doc """
            $($period_str)(v)

        Construct a `$($period_str)` object with the given `v` value. Input must be
        losslessly convertible to an `Int64`.
        """ $period(v)
    end
end
# Now we're safe to define Period-Number conversions
# Anything an Int64 can convert to, a Period can convert to
Base.convert{T<:Number}(::Type{T},x::Period) = convert(T,value(x))
Base.convert{T<:Period}(::Type{T},x::Real) = T(x)

#Print/show/traits
Base.string{P<:Period}(x::P) = string(value(x),_units(x))
Base.show(io::IO,x::Period) = print(io,string(x))
Base.zero{P<:Period}(::Union{Type{P},P}) = P(0)
Base.one{P<:Period}(::Union{Type{P},P}) = P(1)
Base.typemin{P<:Period}(::Type{P}) = P(typemin(Int64))
Base.typemax{P<:Period}(::Type{P}) = P(typemax(Int64))

# Default values (as used by TimeTypes)
"""
    default(p::Period) -> Period

Returns a sensible "default" value for the input Period by returning `one(p)` for Year,
Month, and Day, and `zero(p)` for Hour, Minute, Second, and Millisecond.
"""
function default end

default{T<:DatePeriod}(p::Union{T,Type{T}}) = one(p)
default{T<:TimePeriod}(p::Union{T,Type{T}}) = zero(p)

(-){P<:Period}(x::P) = P(-value(x))
Base.isless{P<:Period}(x::P,y::P) = isless(value(x),value(y))
=={P<:Period}(x::P,y::P) = value(x) == value(y)

# Period Arithmetic, grouped by dimensionality:
import Base: div, fld, mod, rem, gcd, lcm, +, -, *, /, %, .+, .-, .*, .%
for op in (:+,:-,:lcm,:gcd)
    @eval ($op){P<:Period}(x::P,y::P) = P(($op)(value(x),value(y)))
end

for op in (:/,:div,:fld)
    @eval begin
        ($op){P<:Period}(x::P,y::P) = ($op)(value(x),value(y))
        ($op){P<:Period}(x::P,y::Real) = P(($op)(value(x),Int64(y)))
    end
end

for op in (:rem,:mod)
    @eval begin
        ($op){P<:Period}(x::P,y::P) = P(($op)(value(x),value(y)))
        ($op){P<:Period}(x::P,y::Real) = P(($op)(value(x),Int64(y)))
    end
end

/{P<:Period}(X::StridedArray{P}, y::P) = X ./ y
%{P<:Period}(X::StridedArray{P}, y::P) = X .% y
*{P<:Period}(x::P,y::Real) = P(value(x) * Int64(y))
*(y::Real,x::Period) = x * y
.*{P<:Period}(y::Real, X::StridedArray{P}) = X .* y
for (op,Ty,Tz) in ((:.*,Real,:P),
                   (:./,:P,Float64), (:./,Real,:P),
                   (:div,:P,Int64), (:div,Integer,:P),
                   (:.%,:P,:P),
                   (:mod,:P,:P))
    sop = string(op)
    op_ = sop[1] == '.' ? Symbol(sop[2:end]) : op
    @eval begin
        function ($op){P<:Period}(X::StridedArray{P},y::$Ty)
            Z = similar(X, $Tz)
            for (Idst, Isrc) in zip(eachindex(Z), eachindex(X))
                @inbounds Z[Idst] = ($op_)(X[Isrc],y)
            end
            return Z
        end
    end
end

# intfuncs
Base.gcdx{T<:Period}(a::T,b::T) = ((g,x,y)=gcdx(value(a),value(b)); return T(g),x,y)
Base.abs{T<:Period}(a::T) = T(abs(value(a)))

periodisless(::Period,::Year)        = true
periodisless(::Period,::Month)       = true
periodisless(::Year,::Month)         = false
periodisless(::Period,::Week)        = true
periodisless(::Year,::Week)          = false
periodisless(::Month,::Week)         = false
periodisless(::Period,::Day)         = true
periodisless(::Year,::Day)           = false
periodisless(::Month,::Day)          = false
periodisless(::Week,::Day)           = false
periodisless(::Period,::Hour)        = false
periodisless(::Minute,::Hour)        = true
periodisless(::Second,::Hour)        = true
periodisless(::Millisecond,::Hour)   = true
periodisless(::Period,::Minute)      = false
periodisless(::Second,::Minute)      = true
periodisless(::Millisecond,::Minute) = true
periodisless(::Period,::Second)      = false
periodisless(::Millisecond,::Second) = true
periodisless(::Period,::Millisecond) = false

# return (next coarser period, conversion factor):
coarserperiod{P<:Period}(::Type{P}) = (P,1)
coarserperiod(::Type{Millisecond}) = (Second,1000)
coarserperiod(::Type{Second}) = (Minute,60)
coarserperiod(::Type{Minute}) = (Hour,60)
coarserperiod(::Type{Hour}) = (Day,24)
coarserperiod(::Type{Day}) = (Week,7)
coarserperiod(::Type{Month}) = (Year,12)

# Stores multiple periods in greatest to least order by type, not values,
# canonicalized to eliminate zero periods, merge equal period types,
# and convert more-precise periods to less-precise periods when possible
"""
    CompoundPeriod

A `CompoundPeriod` is useful for expressing time periods that are not a fixed multiple of
smaller periods. For example, \"a year and a  day\" is not a fixed number of days, but can
be expressed using a `CompoundPeriod`. In fact, a `CompoundPeriod` is automatically
generated by addition of different period types, e.g. `Year(1) + Day(1)` produces a
`CompoundPeriod` result.
"""
type CompoundPeriod <: AbstractTime
    periods::Array{Period,1}
    function CompoundPeriod(p::Vector{Period})
        n = length(p)
        if n > 1
            sort!(p, rev=true, lt=periodisless)
            # canonicalize p by merging equal period types and removing zeros
            i = j = 1
            while j <= n
                k = j+1
                while k <= n
                    if typeof(p[j]) == typeof(p[k])
                        p[j] += p[k]
                        k += 1
                    else
                        break
                    end
                end
                if p[j] != zero(p[j])
                    p[i] = p[j]
                    i += 1
                end
                j = k
            end
            n = i - 1 # new length
        elseif n == 1 && value(p[1]) == 0
            p = Period[]
            n = 0
        end
        # canonicalize Periods by pushing "overflow" into a coarser period.
        if n > 0
            pc = sizehint!(Period[], n)
            P = typeof(p[n])
            v = value(p[n])
            i = n - 1
            while true
                Pc, f = coarserperiod(P)
                if i > 0 && typeof(p[i]) == P
                    v += value(p[i])
                    i -= 1
                end
                v0 = f == 1 ? v : rem(v, f)
                v0 != 0 && push!(pc, P(v0))
                if v != v0
                    P = Pc
                    v = div(v - v0, f)
                elseif i > 0
                    P = typeof(p[i])
                    v = value(p[i])
                    i -= 1
                else
                    break
                end
            end
            p = reverse!(pc)
            n = length(p)
        else
            return new(resize!(p, n))
        end
        # reduce the amount of mixed positive/negative Periods.
        if n > 0
            pc = sizehint!(Period[], n)
            i = n
            while i > 0
                j = i

                # Determine sign of the largest period in this group which
                # can be converted into via coarserperiod.
                last = Union{}
                current = typeof(p[i])
                while i > 0 && current != last
                    if typeof(p[i]) == current
                        i -= 1
                    end
                    last, current = current, coarserperiod(current)[1]
                end
                s = sign(value(p[i + 1]))

                # Adjust all the periods in the group based upon the
                # largest period sign.
                P = typeof(p[j])
                v = 0
                while j > i
                    Pc, f = coarserperiod(P)
                    if j > 0 && typeof(p[j]) == P
                        v += value(p[j])
                        j -= 1
                    end
                    v0 = f == 1 ? v : mod(v, f * s)
                    v0 != 0 && push!(pc, P(v0))
                    if v != v0
                        P = Pc
                        v = div(v - v0, f)
                    elseif j > 0
                        P = typeof(p[j])
                        v = 0
                    else
                        break
                    end
                end
            end
            p = reverse!(pc)
        end
        return new(p)
    end
end

"""
    CompoundPeriod(periods) -> CompoundPeriod

Construct a `CompoundPeriod` from a `Vector` of `Period`s. The constructor will
automatically simplify the periods into a canonical form according to the following rules:

* All `Period`s of the same type will be added together
* Any `Period` large enough be partially representable by a coarser `Period` will be broken
  into multiple `Period`s (eg. `Hour(30)` becomes `Day(1) + Hour(6)`)
* `Period`s with opposite signs will be combined when possible
  (eg. `Hour(1) - Day(1)` becomes `-Hour(23)`)

Due to the canonicalization, `CompoundPeriod` is also useful for converting time periods
into more human-comprehensible forms.

# Examples
```julia
julia> Dates.CompoundPeriod([Dates.Hour(12), Dates.Hour(13)])
1 day, 1 hour

julia> Dates.CompoundPeriod([Dates.Hour(-1), Dates.Minute(1)])
-59 minutes

julia> Dates.CompoundPeriod([Dates.Month(1), Dates.Week(-2)])
1 month, -2 weeks

julia> Dates.CompoundPeriod(Dates.Minute(50000)))
4 weeks, 6 days, 17 hours, 20 minutes
```
"""
CompoundPeriod{P<:Period}(p::Vector{P}) = CompoundPeriod(Array{Period}(p))

Base.convert(::Type{CompoundPeriod}, x::Period) = CompoundPeriod(Period[x])
function Base.string(x::CompoundPeriod)
    if isempty(x.periods)
        return "empty period"
    else
        s = ""
        for p in x.periods
            s *= ", " * string(p)
        end
        return s[3:end]
    end
end
Base.show(io::IO,x::CompoundPeriod) = print(io,string(x))

# E.g. Year(1) + Day(1)
(+)(x::Period,y::Period) = CompoundPeriod(Period[x,y])
(+)(x::CompoundPeriod,y::Period) = CompoundPeriod(vcat(x.periods,y))
(+)(y::Period,x::CompoundPeriod) = x + y
(+)(x::CompoundPeriod,y::CompoundPeriod) = CompoundPeriod(vcat(x.periods,y.periods))
# E.g. Year(1) - Month(1)
(-)(x::Period,y::Period) = CompoundPeriod(Period[x,-y])
(-)(x::CompoundPeriod,y::Period) = CompoundPeriod(vcat(x.periods,-y))
(-)(x::CompoundPeriod) = CompoundPeriod(-x.periods)
(-)(y::Union{Period,CompoundPeriod},x::CompoundPeriod) = (-x) + y

GeneralPeriod = Union{Period,CompoundPeriod}
(+)(x::GeneralPeriod) = x
(+){P<:GeneralPeriod}(x::StridedArray{P}) = x

for op in (:.+, :.-)
    op_ = Symbol(string(op)[2:end])
    @eval begin
        function ($op){P<:GeneralPeriod}(X::StridedArray{P},y::GeneralPeriod)
            Z = similar(X, CompoundPeriod)
            for (Idst, Isrc) in zip(eachindex(Z), eachindex(X))
                @inbounds Z[Idst] = ($op_)(X[Isrc],y)
            end
            return Z
        end
        ($op){P<:GeneralPeriod}(x::GeneralPeriod,Y::StridedArray{P}) = ($op)(Y,x) |> ($op_)
        ($op_){P<:GeneralPeriod}(x::GeneralPeriod,Y::StridedArray{P}) = ($op)(Y,x) |> ($op_)
        ($op_){P<:GeneralPeriod}(Y::StridedArray{P},x::GeneralPeriod) = ($op)(Y,x)
        ($op_){P<:GeneralPeriod, Q<:GeneralPeriod}(X::StridedArray{P}, Y::StridedArray{Q}) =
            reshape(CompoundPeriod[($op_)(x,y) for (x,y) in zip(X, Y)], promote_shape(size(X),size(Y)))
    end
end

(==)(x::CompoundPeriod, y::Period) = x == CompoundPeriod(y)
(==)(y::Period, x::CompoundPeriod) = x == y
(==)(x::CompoundPeriod, y::CompoundPeriod) = x.periods == y.periods

# Capture TimeType+-Period methods
(+)(a::TimeType,b::Period,c::Period) = (+)(a,b+c)
(-)(a::TimeType,b::Period,c::Period) = (-)(a,b-c)
(+)(a::TimeType,b::Period,c::Period,d::Period...) = (+)((+)(a,b+c),d...)
(-)(a::TimeType,b::Period,c::Period,d::Period...) = (-)((-)(a,b-c),d...)

function (+)(x::TimeType,y::CompoundPeriod)
    for p in y.periods
        x += p
    end
    return x
end
(+)(x::CompoundPeriod,y::TimeType) = y + x

function (-)(x::TimeType,y::CompoundPeriod)
    for p in y.periods
        x -= p
    end
    return x
end
(-)(x::CompoundPeriod,y::TimeType) = y - x

# Fixed-value Periods (periods corresponding to a well-defined time interval,
# as opposed to variable calendar intervals like Year).
typealias FixedPeriod Union{Week,Day,Hour,Minute,Second,Millisecond}

# like div but throw an error if remainder is nonzero
function divexact(x,y)
    q,r = divrem(x, y)
    r == 0 || throw(InexactError())
    return q
end

# FixedPeriod conversions and promotion rules
const fixedperiod_conversions = [(Week,7),(Day,24),(Hour,60),(Minute,60),(Second,1000),(Millisecond,1)]
for i = 1:length(fixedperiod_conversions)
    (T,n) = fixedperiod_conversions[i]
    N = 1
    for j = i-1:-1:1 # less-precise periods
        (Tc,nc) = fixedperiod_conversions[j]
        N *= nc
        vmax = typemax(Int64) ÷ N
        vmin = typemin(Int64) ÷ N
        @eval function Base.convert(::Type{$T}, x::$Tc)
            $vmin ≤ value(x) ≤ $vmax || throw(InexactError())
            return $T(value(x)*$N)
        end
    end
    N = n
    for j = i+1:length(fixedperiod_conversions) # more-precise periods
        (Tc,nc) = fixedperiod_conversions[j]
        @eval Base.convert(::Type{$T}, x::$Tc) = $T(divexact(value(x), $N))
        @eval Base.promote_rule(::Type{$T},::Type{$Tc}) = $Tc
        N *= nc
    end
end
# have to declare thusly so that diagonal dispatch above takes precedence:
(==){T<:FixedPeriod,S<:FixedPeriod}(x::T,y::S) = (==)(promote(x,y)...)
Base.isless{T<:FixedPeriod,S<:FixedPeriod}(x::T,y::S) = isless(promote(x,y)...)

# other periods with fixed conversions but which aren't fixed time periods
typealias OtherPeriod Union{Month,Year}
let vmax = typemax(Int64) ÷ 12, vmin = typemin(Int64) ÷ 12
    @eval function Base.convert(::Type{Month}, x::Year)
        $vmin ≤ value(x) ≤ $vmax || throw(InexactError())
        Month(value(x)*12)
    end
end
Base.convert(::Type{Year}, x::Month) = Year(divexact(value(x),12))
Base.promote_rule(::Type{Year}, ::Type{Month}) = Month
(==){T<:OtherPeriod,S<:OtherPeriod}(x::T,y::S) = (==)(promote(x,y)...)
Base.isless{T<:OtherPeriod,S<:OtherPeriod}(x::T,y::S) = isless(promote(x,y)...)

# truncating conversions to milliseconds and days:
toms(c::Millisecond) = value(c)
toms(c::Second)      = 1000*value(c)
toms(c::Minute)      = 60000*value(c)
toms(c::Hour)        = 3600000*value(c)
toms(c::Day)         = 86400000*value(c)
toms(c::Week)        = 604800000*value(c)
toms(c::Month)       = 86400000.0*30.436875*value(c)
toms(c::Year)        = 86400000.0*365.2425*value(c)
toms(c::CompoundPeriod) = isempty(c.periods)?0.0 : Float64(sum(toms,c.periods))
days(c::Millisecond) = div(value(c),86400000)
days(c::Second)      = div(value(c),86400)
days(c::Minute)      = div(value(c),1440)
days(c::Hour)        = div(value(c),24)
days(c::Day)         = value(c)
days(c::Week)        = 7*value(c)
days(c::Year)        = 365.2425*value(c)
days(c::Month)       = 30.436875*value(c)
days(c::CompoundPeriod) = isempty(c.periods)?0.0 : Float64(sum(days,c.periods))


# differentiated shock capturing functions

# main entry point
"""
  Main function for assembling the shock capturing terms into the Jacobian.

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * capture: an [`AbstractShockCapturing`](@ref)
   * assem: an [`AssembleElementData`](@ref)
"""
function applyShockCapturing_diff(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData, opts,
                             capture::AbstractVolumeShockCapturing,
                             assem::AssembleElementData)

  # nothing to do here, call the implementation
  sensor = getShockSensor(capture)
  calcShockCapturing_diff(mesh, sbp, eqn, opts, sensor, capture, assem)

  return nothing
end

function applyShockCapturing_diff(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData{Tsol, Tres}, opts,
                             capture::AbstractFaceShockCapturing,
                             assem::AssembleElementData) where {Tsol, Tres}

  sensor = getShockSensor(capture)
  _applyShockCapturing_diff(mesh, sbp, eqn, opts, sensor, capture, assem)

  return nothing
end

function _applyShockCapturing_diff(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData{Tsol, Tres}, opts,
                             sensor::AbstractShockSensor,
                             capture::AbstractFaceShockCapturing,
                             assem::AssembleElementData) where {Tsol, Tres}

  if mesh.commsize > 1 &&
    getParallelData(eqn.shared_data) != PARALLEL_DATA_ELEMENT

    error("shock capturing requires PARALLEL_DATA_ELEMENT")
  end

  assertReceivesWaited(eqn.shared_data)  # parallel communication for the regular
                                     # face integrals should already have
                                     # finished parallel communication

  # re-using the shockmesh from one iteration to the next makes allows
  # the algorithm to re-use existing arrays that depend on the number of
  # elements with the shock in it.
  shockmesh = eqn.params.shockmesh
  reset(shockmesh)
  for i=1:mesh.numEl
    q_i = ro_sview(eqn.q, :, :, i)
    coords_i = ro_sview(mesh.coords, :, :, i)
    dxidx_i = ro_sview(mesh.dxidx, :, :, :, i)
    jac_i = ro_sview(mesh.jac, :, i)

    val = isShockElement(eqn.params, sbp, sensor, q_i, i, coords_i, dxidx_i,
                            jac_i)
    if val
      # push to shockmesh
      push!(shockmesh, i)
    end
  end

  completeShockElements(mesh, shockmesh)

  allocateArrays(capture, mesh, shockmesh)

  # call shock capturing scheme
  calcShockCapturing_diff(mesh, sbp, eqn, opts, capture, shockmesh, assem)

  return nothing
end

#------------------------------------------------------------------------------
# revq

"""
  Reverse mode of `applyShockCapturing` wrt q

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * data: an [`AbstractShockCaputring`](@ref) object
"""
function applyShockCapturing_revq(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData, opts,
                             capture::AbstractVolumeShockCapturing)

  # unlike AbstractFaceShockCapturing, nothing to do here

  sensor = getShockSensor(capture)
  calcShockCapturing_revq(mesh, sbp, eqn, opts, sensor, capture)

  return nothing
end

# Face shock capturing version
function applyShockCapturing_revq(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData, opts,
                             capture::AbstractFaceShockCapturing)

  sensor = getShockSensor(capture)
  shockmesh = eqn.params.shockmesh
  setupShockCapturing(mesh, sbp, eqn, opts, sensor, capture, shockmesh)

  # call shock capturing scheme
  calcShockCapturing_revq(mesh, sbp, eqn, opts, capture, shockmesh)

  return nothing
end

#------------------------------------------------------------------------------
# revm

"""
  Reverse mode of `applyShockCapturing` wrt metrics

  **Inputs**

   * mesh
   * sbp
   * eqn
   * opts
   * data: an [`AbstractShockCaputring`](@ref) object
"""
function applyShockCapturing_revm(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData, opts,
                             capture::AbstractVolumeShockCapturing)

  # unlike AbstractFaceShockCapturing, nothing to do here

  sensor = getShockSensor(capture)
  initForRevm(sensor)
  calcShockCapturing_revm(mesh, sbp, eqn, opts, sensor, capture)
  finishRevm(mesh, sbp, eqn, opts, sensor)

  return nothing
end

function applyShockCapturing_revm(mesh::AbstractMesh, sbp::AbstractOperator,
                             eqn::EulerData, opts,
                             capture::AbstractFaceShockCapturing)

  # unlike AbstractFaceShockCapturing, nothing to do here

  sensor = getShockSensor(capture)
  shockmesh = eqn.params.shockmesh
  initForRevm(sensor)
  setupShockCapturing(mesh, sbp, eqn, opts, sensor, capture, shockmesh)
  calcShockCapturing_revm(mesh, sbp, eqn, opts, capture, shockmesh)
  finishRevm(mesh, sbp, eqn, opts, sensor)

  return nothing
end