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


################################################################################
#                   CORE DISPATCHVARIABLEREF METHOD EXTENSIONS
################################################################################
# Extend dispatch_variable_ref
function dispatch_variable_ref(model::InfiniteModel,
                               index::ReducedVariableIndex
                               )::ReducedVariableRef
    return ReducedVariableRef(model, index)
end

# Extend _add_data_object
function _add_data_object(model::InfiniteModel,
                          object::VariableData{<:ReducedVariable}
                          )::ReducedVariableIndex
    return MOIUC.add_item(model.reduced_vars, object)
end

# Extend _data_dictionary (type based)
function _data_dictionary(model::InfiniteModel, ::Type{ReducedVariable}
    )::MOIUC.CleverDict{ReducedVariableIndex, VariableData{ReducedVariable{GeneralVariableRef}}}
    return model.reduced_vars
end

# Extend _data_dictionary (reference based)
function _data_dictionary(vref::ReducedVariableRef
    )::MOIUC.CleverDict{ReducedVariableIndex, VariableData{ReducedVariable{GeneralVariableRef}}}
    return JuMP.owner_model(vref).reduced_vars
end

# Extend _data_object
function _data_object(vref::ReducedVariableRef
    )::VariableData{ReducedVariable{GeneralVariableRef}}
  object = get(_data_dictionary(vref), JuMP.index(vref), nothing)
  object === nothing && error("Invalid point variable reference, cannot find " *
                        "corresponding variable in the model. This is likely " *
                        "caused by using the reference of a deleted variable.")
  return object
end

"""
    internal_reduced_variable(vref::ReducedVariableRef,
                              key::Val{:my_ext_key})::ReducedVariable

Return the reduced variable object of `vref` assuming it is an internal variable
made during measure expansion within an optimizer model. This will apply to
optimizer model extensions that utilize `add_measure_variable` in combination
with `expand_measure`.
"""
function internal_reduced_variable end

# Extend _core_variable_object
function _core_variable_object(vref::ReducedVariableRef
    )::ReducedVariable{GeneralVariableRef}
    if !haskey(_data_dictionary(vref), JuMP.index(vref))
        model = JuMP.owner_model(vref)
        return internal_reduced_variable(vref, Val(optimizer_model_key(model)))
    else
        return _data_object(vref).variable
    end
end

# Extend _object_numbers
function _object_numbers(vref::ReducedVariableRef)::Vector{Int}
    return _core_variable_object(vref).object_nums
end

# Extend _parameter_numbers
function _parameter_numbers(vref::ReducedVariableRef)::Vector{Int}
    return _core_variable_object(vref).parameter_nums
end

################################################################################
#                             DEFINITION METHODS
################################################################################
"""
    JuMP.build_variable(_error::Function, ivref::GeneralVariableRef,
                        eval_supports::Dict{Int, Float64}; [check::Bool = true]
                        )::ReducedVariable{GeneralVariableRef}

Extend the `JuMP.build_variable` function to build a reduced infinite variable
based on the infinite variable/derivative/parameter function `ivref` with reduction support `eval_supports`.
Will check that input is appropriate if `check = true`. Errors if `ivref` is
not an infinite variable, `eval_supports` violate infinite parameter domains, or
if the support dimensions don't match the infinite parameter dimensions of `ivref`.
This is intended an internal method for use in evaluating measures.
"""
function JuMP.build_variable(_error::Function, ivref::GeneralVariableRef,
                             eval_supports::Dict{Int, Float64};
                             check::Bool = true
                             )::ReducedVariable{GeneralVariableRef}
    # check the inputs
    dvref = dispatch_variable_ref(ivref)
    if check && !(dvref isa Union{InfiniteVariableRef, DerivativeRef, ParameterFunctionRef})
         _error("Must specify an infinite variable/derivative/parameter function dependency.")
    elseif check && maximum(keys(eval_supports)) > length(parameter_list(dvref))
        _error("Support evaluation dictionary indices do not match the infinite " *
               "parameter dependencies of $(ivref).")
    end
    prefs = parameter_list(dvref)
    if check
        for (index, value) in eval_supports
            pref = prefs[index]
            if JuMP.has_lower_bound(pref) && !supports_in_set(value, infinite_set(pref))
                _error("Evaluation support violates infinite parameter domain(s).")
            end
        end
    end
    # get the parameter object numbers of the dependencies
    object_nums = Int[]
    for i in eachindex(prefs)
        if !haskey(eval_supports, i)
            union!(object_nums, _object_number(prefs[i]))
        end
    end
    # get the parameter numbers
    orig_nums = _parameter_numbers(ivref)
    param_nums = [orig_nums[i] for i in eachindex(orig_nums)
                  if !haskey(eval_supports, i)]
    # round the support values in accordance with the significant digits
    for (k, v) in eval_supports
        eval_supports[k] = round(v, sigdigits = significant_digits(prefs[k]))
    end
    # build the variable
    return ReducedVariable(ivref, eval_supports, param_nums, object_nums)
end

"""
    JuMP.add_variable(model::InfiniteModel, var::ReducedVariable,
                      [name::String = ""])::GeneralVariableRef

Extend the [`JuMP.add_variable`](@ref JuMP.add_variable(::JuMP.Model, ::JuMP.ScalarVariable, ::String))
function to accomodate `InfiniteOpt` reduced variable types. Adds `var` to the
infinite model `model` and returns a [`GeneralVariableRef`](@ref).
Primarily intended to be an internal function used in evaluating measures.
"""
function JuMP.add_variable(model::InfiniteModel, var::ReducedVariable,
                           name::String = "")::GeneralVariableRef
    ivref = dispatch_variable_ref(var.infinite_variable_ref)
    JuMP.check_belongs_to_model(ivref, model)
    data_object = VariableData(var, name)
    vindex = _add_data_object(model, data_object)
    push!(_reduced_variable_dependencies(ivref), vindex)
    gvref = GeneralVariableRef(model, vindex.value, typeof(vindex))
    return gvref
end

################################################################################
#                          PARAMETER REFERENCE METHODS
################################################################################
"""
    infinite_variable_ref(vref::ReducedVariableRef)::GeneralVariableRef

Return the infinite variable/derivative/parameter function reference associated
with the reduced infinite variable `vref`.

**Example**
```julia-repl
julia> infinite_variable_ref(vref)
g(t, x)
```
"""
function infinite_variable_ref(vref::ReducedVariableRef)::GeneralVariableRef
    return _core_variable_object(vref).infinite_variable_ref
end

"""
    eval_supports(vref::ReducedVariableRef)::Dict{Int, Float64}

Return the evaluation supports associated with the reduced infinite variable
`vref`.

**Example**
```julia-repl
julia> eval_supports(vref)
Dict{Int64,Float64} with 1 entry:
  1 => 0.5
```
"""
function eval_supports(vref::ReducedVariableRef)::Dict{Int, Float64}
    return _core_variable_object(vref).eval_supports
end

"""
    raw_parameter_refs(vref::ReducedVariableRef)::VectorTuple{GeneralVariableRef}

Return the raw [`VectorTuple`](@ref) of the parameter references that `vref`
depends on. This is primarily an internal method where
[`parameter_refs`](@ref parameter_refs(vref::ReducedVariableRef))
is intended as the preferred user function.
"""
function raw_parameter_refs(vref::ReducedVariableRef
                            )::VectorTuple{GeneralVariableRef}
    orig_prefs = raw_parameter_refs(infinite_variable_ref(vref))
    eval_supps = eval_supports(vref)
    delete_indices = [haskey(eval_supps, i) for i in eachindex(orig_prefs)]
    return deleteat!(copy(orig_prefs), delete_indices)
end

"""
    parameter_refs(vref::ReducedVariableRef)::Tuple

Return the infinite parameter references associated with the reduced infinite variable
`vref`. This is formatted as a `Tuple` of containing the parameter references as
they were inputted to define the untranscripted infinite variable except, the
evaluated parameters are excluded.

**Example**
```julia-repl
julia> parameter_refs(vref)
(t, [x[1], x[2]])
```
"""
function parameter_refs(vref::ReducedVariableRef)::Tuple
    return Tuple(raw_parameter_refs(vref))
end

"""
    parameter_list(vref::ReducedVariableRef)::Vector{GeneralVariableRef}

Return a vector of the parameter references that `vref` depends on. This is
primarily an internal method where [`parameter_refs`](@ref parameter_refs(vref::ReducedVariableRef))
is intended as the preferred user function.
"""
function parameter_list(vref::ReducedVariableRef)::Vector{GeneralVariableRef}
    orig_prefs = parameter_list(infinite_variable_ref(vref))
    eval_supps = eval_supports(vref)
    return [orig_prefs[i] for i in eachindex(orig_prefs) if !haskey(eval_supps, i)]
end

################################################################################
#                            VARIABLE DEPENDENCIES
################################################################################
# Extend _derivative_dependencies
function _derivative_dependencies(vref::ReducedVariableRef
                                  )::Vector{DerivativeIndex}
    return _data_object(vref).derivative_indices
end

"""
    used_by_derivative(vref::ReducedVariableRef)::Bool

Return a `Bool` indicating if `vref` is used by a derivative.

**Example**
```julia-repl
julia> used_by_derivative(vref)
true
```
"""
function used_by_derivative(vref::ReducedVariableRef)::Bool
    return !isempty(_derivative_dependencies(vref))
end

"""
    is_used(vref::ReducedVariableRef)::Bool

Return a `Bool` indicating if `vref` is used in the model.

**Example**
```julia-repl
julia> is_used(vref)
true
```
"""
function is_used(vref::ReducedVariableRef)::Bool
    if used_by_measure(vref) || used_by_constraint(vref)
        return true
    end
    for dindex in _derivative_dependencies(vref)
        if is_used(DerivativeRef(JuMP.owner_model(vref), dindex))
            return true
        end
    end
    return false
end

################################################################################
#                            VARIABLE INFO METHODS
################################################################################
"""
    JuMP.has_lower_bound(vref::ReducedVariableRef)::Bool

Extend [`JuMP.has_lower_bound`](@ref) to return a `Bool` whether the original
infinite variable of `vref` has a lower bound.

**Example**
```julia-repl
julia> has_lower_bound(vref)
true
```
"""
function JuMP.has_lower_bound(vref::ReducedVariableRef)::Bool
    return JuMP.has_lower_bound(infinite_variable_ref(vref))
end

"""
    JuMP.lower_bound(vref::ReducedVariableRef)::Float64

Extend [`JuMP.lower_bound`](@ref) to return the lower bound of the original
infinite variable of `vref`. Errors if `vref` doesn't have a lower bound.

**Example**
```julia-repl
julia> lower_bound(vref)
0.0
```
"""
function JuMP.lower_bound(vref::ReducedVariableRef)::Float64
    ivref = dispatch_variable_ref(infinite_variable_ref(vref))
    if !JuMP.has_lower_bound(ivref)
        error("Variable $(vref) does not have a lower bound.")
    end
    return JuMP.lower_bound(ivref)
end

# Extend to return the index of the lower bound constraint associated with the
# original infinite variable of `vref`.
function InfiniteOpt._lower_bound_index(vref::ReducedVariableRef)::ConstraintIndex
    ivref = dispatch_variable_ref(infinite_variable_ref(vref))
    if !JuMP.has_lower_bound(ivref)
        error("Variable $(vref) does not have a lower bound.")
    end
    return InfiniteOpt._lower_bound_index(ivref)
end

"""
    JuMP.LowerBoundRef(vref::ReducedVariableRef)::InfOptConstraintRef

Extend [`JuMP.LowerBoundRef`](@ref) to extract a constraint reference for the
lower bound of the original infinite variable of `vref`.

**Example**
```julia-repl
julia> cref = LowerBoundRef(vref)
var >= 0.0
```
"""
function JuMP.LowerBoundRef(vref::ReducedVariableRef)::InfOptConstraintRef
    return JuMP.LowerBoundRef(infinite_variable_ref(vref))
end

"""
    JuMP.has_upper_bound(vref::ReducedVariableRef)::Bool

Extend [`JuMP.has_upper_bound`](@ref) to return a `Bool` whether the original
infinite variable of `vref` has an upper bound.

**Example**
```julia-repl
julia> has_upper_bound(vref)
true
```
"""
function JuMP.has_upper_bound(vref::ReducedVariableRef)::Bool
    return JuMP.has_upper_bound(infinite_variable_ref(vref))
end

"""
    JuMP.upper_bound(vref::ReducedVariableRef)::Float64

Extend [`JuMP.upper_bound`](@ref) to return the upper bound of the original
infinite variable of `vref`. Errors if `vref` doesn't have a upper bound.

**Example**
```julia-repl
julia> upper_bound(vref)
0.0
```
"""
function JuMP.upper_bound(vref::ReducedVariableRef)::Float64
    ivref = dispatch_variable_ref(infinite_variable_ref(vref))
    if !JuMP.has_upper_bound(ivref)
        error("Variable $(vref) does not have a upper bound.")
    end
    return JuMP.upper_bound(ivref)
end

# Extend to return the index of the upper bound constraint associated with the
# original infinite variable of `vref`.
function InfiniteOpt._upper_bound_index(vref::ReducedVariableRef)::ConstraintIndex
    ivref = dispatch_variable_ref(infinite_variable_ref(vref))
    if !JuMP.has_upper_bound(ivref)
        error("Variable $(vref) does not have a upper bound.")
    end
    return InfiniteOpt._upper_bound_index(ivref)
end

"""
    JuMP.UpperBoundRef(vref::ReducedVariableRef)::InfOptConstraintRef

Extend [`JuMP.UpperBoundRef`](@ref) to extract a constraint reference for the
upper bound of the original infinite variable of `vref`.

**Example**
```julia-repl
julia> cref = UpperBoundRef(vref)
var <= 1.0
```
"""
function JuMP.UpperBoundRef(vref::ReducedVariableRef)::InfOptConstraintRef
    return JuMP.UpperBoundRef(infinite_variable_ref(vref))
end

"""
    JuMP.is_fixed(vref::ReducedVariableRef)::Bool

Extend [`JuMP.is_fixed`](@ref) to return `Bool` whether the original infinite
variable of `vref` is fixed.

**Example**
```julia-repl
julia> is_fixed(vref)
true
```
"""
function JuMP.is_fixed(vref::ReducedVariableRef)::Bool
    return JuMP.is_fixed(infinite_variable_ref(vref))
end

"""
    JuMP.fix_value(vref::ReducedVariableRef)::Float64

Extend [`JuMP.fix_value`](@ref) to return the fix value of the original infinite
variable of `vref`. Errors if variable is not fixed.

**Example**
```julia-repl
julia> fix_value(vref)
0.0
```
"""
function JuMP.fix_value(vref::ReducedVariableRef)::Float64
    ivref = dispatch_variable_ref(infinite_variable_ref(vref))
    if !JuMP.is_fixed(ivref)
        error("Variable $(vref) is not fixed.")
    end
    return JuMP.fix_value(ivref)
end

# Extend to return the index of the fix constraint associated with the original
# infinite variable of `vref`.
function InfiniteOpt._fix_index(vref::ReducedVariableRef)::ConstraintIndex
    ivref = dispatch_variable_ref(infinite_variable_ref(vref))
    if !JuMP.is_fixed(ivref)
        error("Variable $(vref) is not fixed.")
    end
    return InfiniteOpt._fix_index(ivref)
end

"""
    JuMP.FixRef(vref::ReducedVariableRef)::InfOptConstraintRef

Extend [`JuMP.FixRef`](@ref) to return the constraint reference of the fix
constraint associated with the original infinite variable of `vref`. Errors
`vref` is not fixed.

**Examples**
```julia-repl
julia> cref = FixRef(vref)
var == 1.0
```
"""
function JuMP.FixRef(vref::ReducedVariableRef)::InfOptConstraintRef
    return JuMP.FixRef(infinite_variable_ref(vref))
end

# Fallback dispatch
function JuMP.start_value(vref::ReducedVariableRef)
    return JuMP.start_value(infinite_variable_ref(vref))
end

"""
    start_value_function(vref::ReducedVariableRef)::Union{Nothing, Function}

Return the function that is used to generate the start values of `vref` for
particular support values. Returns `nothing` if no start behavior has been
specified.

**Example**
```julia-repl
julia> start_value_func(vref)
my_func
```
"""
function start_value_function(vref::ReducedVariableRef)::Union{Nothing, Function}
    return start_value_function(infinite_variable_ref(vref))
end

"""
    JuMP.is_binary(vref::ReducedVariableRef)::Bool

Extend [`JuMP.is_binary`](@ref) to return `Bool` whether the original infinite
variable of `vref` is binary.

**Example**
```julia-repl
julia> is_binary(vref)
true
```
"""
function JuMP.is_binary(vref::ReducedVariableRef)::Bool
    return JuMP.is_binary(infinite_variable_ref(vref))
end

# Extend to return the index of the binary constraint associated with the
# original infinite variable of `vref`.
function InfiniteOpt._binary_index(vref::ReducedVariableRef)::ConstraintIndex
    ivref = dispatch_variable_ref(infinite_variable_ref(vref))
    if !JuMP.is_binary(ivref)
        error("Variable $(vref) is not binary.")
    end
    return InfiniteOpt._binary_index(ivref)
end

"""
    JuMP.BinaryRef(vref::ReducedVariableRef)::InfOptConstraintRef

Extend [`JuMP.BinaryRef`](@ref) to return a constraint reference to the
constraint constrainting the original infinite variable of `vref` to be binary.
Errors if one does not exist.

**Example**
```julia-repl
julia> cref = BinaryRef(vref)
var binary
```
"""
function JuMP.BinaryRef(vref::ReducedVariableRef)::InfOptConstraintRef
    return JuMP.BinaryRef(infinite_variable_ref(vref))
end

"""
    JuMP.is_integer(vref::ReducedVariableRef)::Bool

Extend [`JuMP.is_integer`](@ref) to return `Bool` whether the original infinite
variable of `vref` is integer.

**Example**
```julia-repl
julia> is_integer(vref)
true
```
"""
function JuMP.is_integer(vref::ReducedVariableRef)::Bool
    return JuMP.is_integer(infinite_variable_ref(vref))
end

# Extend to return the index of the integer constraint associated with the
# original infinite variable of `vref`.
function InfiniteOpt._integer_index(vref::ReducedVariableRef)::ConstraintIndex
    ivref = dispatch_variable_ref(infinite_variable_ref(vref))
    if !JuMP.is_integer(ivref)
        error("Variable $(vref) is not an integer.")
    end
    return InfiniteOpt._integer_index(ivref)
end

"""
    JuMP.IntegerRef(vref::ReducedVariableRef)::InfOptConstraintRef

Extend [`JuMP.IntegerRef`](@ref) to return a constraint reference to the
constraint constrainting the original infinite variable of `vref` to be integer.
Errors if one does not exist.

**Example**
```julia-repl
julia> cref = IntegerRef(vref)
var integer
```
"""
function JuMP.IntegerRef(vref::ReducedVariableRef)::InfOptConstraintRef
    return JuMP.IntegerRef(infinite_variable_ref(vref))
end

################################################################################
#                                  DELETION
################################################################################
# Extend _delete_variable_dependencies (for use with JuMP.delete)
function _delete_variable_dependencies(vref::ReducedVariableRef)::Nothing
    # remove mapping to infinite variable
    ivref = dispatch_variable_ref(infinite_variable_ref(vref))
    filter!(e -> e != JuMP.index(vref), _reduced_variable_dependencies(ivref))
    # delete associated derivative variables and mapping
    model = JuMP.owner_model(vref)
    for index in _derivative_dependencies(vref)
        JuMP.delete(model, dispatch_variable_ref(model, index))
    end
    return
end


bstr = [path * "/SampleFiles/Bottle/B024*",
        path * "/SampleFiles/Bottle/B2021901700*",
        path * "/SampleFiles/Bottle/B203*",
        path * "/SampleFiles/Bottle/B20319115Rainfallmm"]

printstyled("  Bottle (UNAVCO strain data)\n", color=:light_green)
for i = 1:4
  S = verified_read_data("bottle", bstr[i], nx_new=14400, nx_add=14400, vl=true)
  fill_pbo!(S)
  if i == 4
    @test S.name[1] == "quarry203bwa2007"
    @test isapprox(S.fs[1], 1.0/(30*60))
    @test isapprox(S.loc[1].el, 814.4)
    @test S.units[1] == "mm"
  end
end

printstyled("    source logging\n", color=:light_green)
redirect_stdout(out) do
  S = SeisIO.read_bottle(bstr[4], 14400, 14400, true, true, 0)
  show_src(S, 1)
  show_src(S[1])
  show_src(S)
end


module MartaCT

using Reexport
import
    LinearAlgebra,
    IterTools,
    YAML,
    FFTW,
    Statistics,
    StatsBase,
    Distributions,
    IntervalSets

include("Utils.jl")
include("Applicative.jl")
include("Monads.jl")
include("TypeDict.jl")
include("CalibrationBase.jl")
include("Interpolation.jl")
include("Coordinates.jl")
include("Geometry.jl")
include("CTImages.jl")
include("FanBeam.jl")
include("AbstractAlgorithms.jl")
include("RadonAlgorithm.jl")
include("Simulations.jl")
include("CTIO.jl")
include("Info.jl")
include("TestImages.jl")
include("Calibration.jl")
include("CTScan.jl")
include("CTPlots.jl")

@reexport using
    .Utils,
    .Applicative,
    .Monads,
    .Info,
    .CalibrationBase,
    .Coordinates,
    .Geometry,
    .CTIO,
    .CTImages,
    .FanBeam,
    .AbstractAlgorithms,
    .RadonAlgorithm,
    .TestImages,
    .Calibration,
    .CTScan,
    .CTPlots

include("precompile_includer.jl")

end # module


module SGDOptim

using Compat
using Reexport
using ArrayViews
@reexport using EmpiricalRisks

import Base: start, next, done
import Base.LinAlg: axpy!

export
    # from ArrayViews
    views,

    # streams.jl
    SampleStream,
    GenericSampleSeq,
    sample_seq,
    minibatch_seq,

    # callback.jl
    simple_trace,
    gtcompare_trace,

    # sgd_std.jl
    sgd!,
    sgd

# source files

include("utils.jl")
include("streams.jl")
include("callback.jl")

include("sgdstd.jl")

end # module


using Mosek

escapename(n::String) =
    replace(n,r"(^#)","_#")

fmtcof(c::Float64) =
    if     c < 0.0
        if     c < -1.0 || c > -1.0
             "- $(-c)"
        else
            "-"
        end
    elseif c > 0.0
        if c < 1.0 || c > 1.0
            "+ $c"
        else
            "+"
        end
    else
        "0"
    end

write(t::Mosek.Task, filename::String) =
    open(filename,"w") do f
        write(t,f)
    end

function write(t::Mosek.Task, f::IO)
    numvar    = getnumvar(t)
    numbarvar = getnumbarvar(t)
    numcon    = getnumcon(t)
    numcone   = getnumcone(t)
    numanz    = getnumanz(t)

    varnames = String[ escapename(getvarname(t,i)) for i in 1:numvar]
    barvarnames = String[ escapename(getbarvarname(t,i)) for i in 1:numbarvar]
    connames = String[ escapename(getconname(t,i)) for i in 1:numcon]
    conenames = String[ escapename(getconename(t,i)) for i in 1:numcone]

    for i in 1:numvar
        if length(varnames[i]) == 0
            varnames[i] = "#x$i"
        end
    end

    for i in 1:numcon
        if length(connames[i]) == 0
            connames[i] = "#c$i"
        end
    end

    for i in 1:numcone
        if length(conenames[i]) == 0
            conenames[i] = "#k$i"
        end
    end

    for i in 1:numbarvar
        if length(barvarnames[i]) == 0
            barvarnames[i] = "#X̂$i"
        end
    end

    maxvarnamelen = max(map(n -> length(n),varnames)...)
    maxconnamelen = max(map(n -> length(n),connames)...)

    taskname = gettaskname(t)

    println(f,"Task $(escape_string(taskname))")

    _,barcidx = getbarcsparsity(t)

    _,baraidx = getbarasparsity(t)
    baraptrb = zeros(Int,numcon+1)
    for idx in baraidx
        (barai,baraj) = getbaraidxij(t,idx)
        baraptrb[barai+1] += 1
    end
    for i in 1:numcon
        baraptrb[i+1] += baraptrb[i]
    end
    baraptrb += 1

    baraptre = copy(baraptrb)
    baraidx_idx = Vector{Int}(length(baraidx))
    for idx in baraidx
        (barai,baraj) = getbaraidxij(t,idx)
        baraidx_idx[baraptre[barai]] = idx
        baraptre[barai] += 1
    end

    # Objective
    objsense = getobjsense(t)
    objname  = getobjname(t)

    if objsense == MSK_OBJECTIVE_SENSE_MAXIMIZE
        println(f,"Maximize $(escape_string(objname))")
    else
        println(f,"Minimze $(escape_string(objname))")
    end
    c = getc(t)
    print(f,"    ")
    for j in 1:numvar
        if c[j] < 0.0 || c[j] > 0.0
            print(f,"$(fmtcof(c[j])) $(varnames[j]) ")
        end
    end

    # Bar entries go here

    cfix = getcfix(t)
    if cfix < 0 || cfix > 0
        print(f,fmtcof(cfix))
    end
    println()

    # Constraints
    if numcon+numcone > 0
        println(f,"Subject to")

        for i in 1:numcon
            (nzi,subi,vali) = getarow(t,Int32(i))
            (bk,bl,bu) = getconbound(t,Int32(i))

            print(f,"    $(connames[i]): ")
            if bk == MSK_BK_LO || bk == MSK_BK_RA
                print(f,"$bl < ")
            end
            for k in 1:nzi
                print(f,"$(fmtcof(vali[k])) $(varnames[subi[k]]) ")
            end

            for k in baraptrb[i]:baraptrb[i+1]-1
                idx = baraidx_idx[k]
                (barai,baraj,num,sub,w) = getbaraidx(t,idx)
                if num == 1
                    print(f,"$(fmtcof(w[1])) #MX$(sub[1]) ⋅ $(barvarnames[baraj]) ")
                elseif num > 1
                    print(f,"(")
                    printf(f,"$(fmtcof(w[1])) #MX$(sub[1])")
                    print(f,") ⋅ $(barvarnames[baraj]) ")
                end
            end

            if bk == MSK_BK_UP || bk == MSK_BK_RA
                println(f,"< $bu")
            elseif bk == MSK_BK_FX
                println(f,"= $bu")
            end

        end

        for k in 1:numcone
            (ct,conepar,nummem,submem) = getcone(t,k)
            dom =
                if ct == MSK_CT_QUAD "Q"
                elseif ct == MSK_CT_RQUAD "Q_r"
                else "?"
                end

            print(f,"    $(conenames[k]): ($(varnames[submem[1]])")
            for i in 2:nummem
                print(f,",$(varnames[submem[i]])")
            end
            println(f,") ∈ $dom(nummem)")
        end
    end

    # Variable bounds
    if numvar > 0
        println(f,"Variables")
        for j in 1:numvar
            bk,bl,bu = getvarbound(t,j)

            if bk == MSK_BK_FX
                println(f,"    $(varnames[j]) = $bl")
            else
                blstr =
                    if     bk == MSK_BK_LO || bk == MSK_BK_RA "[$bl"
                    else   "]-Inf"
                    end
                bustr =
                    if     bk == MSK_BK_UP || bk == MSK_BK_RA "$bu]"
                    else   "+Inf["
                    end

                println(f,"    $(varnames[j]) ∈ $blstr;$bustr")
            end
        end

        for j in 1:numbarvar
            dim = getdimbarvarj(t,j)
            println(f,"    $(barvarnames[j]) ∈ S($dim)")
        end

    end

    # Matsto
    numsymmat = getnumsymmat(t)
    if numsymmat > 0
        println(f,"Symmetric matrixes")
        for i in 1:numsymmat
            dim,nz,tp = getsymmatinfo(t,i)
            (msubi,msubj,mvalij) = getsparsesymmat(t,i)
            print(f,"    #MX$i ($dim): [ " )
            if nz > 0
                print(f,"($(msubi[1]),$(msubj[1]),$(mvalij[1]))")
                for i in 2:nz
                    print(f,",($(msubi[i]),$(msubj[i]),$(mvalij[i]))")
                end
            end
            println(f," ]")
        end
    end
end

maketask() do t
    readdata(t,ARGS[1])
    write(t,STDOUT)
end


module PolarimetryCalibration

include("mueller.jl")

export müller_rec, müller_sky, lin_pol_source

end # module


# using Quads
using Test
using Random

import Quads.PokerCard
import Quads.PokerDeck
import Quads.Lookup
import Quads.Evaluator

@testset "Quads.jl" begin
    # test the decks
    @test length(PokerDeck.make_deck()) == 52
    @test length(PokerDeck.draw!(5, PokerDeck.make_deck())) == 5

    # some hand and eval tests
    hand_string = ["2h", "3h", "4h", "5h", "Ah"]
    @test length(PokerCard.newCard.(hand_string)) == 5
    hand_ints = PokerCard.newCard.(hand_string)
    @test Evaluator.class_to_string(Evaluator.get_rank_class(Evaluator.evaluate(hand_ints))) == "Straight Flush"

    # seven cards from the deck and eval
    hand_ints = PokerDeck.draw!(7, PokerDeck.make_deck())
    @test typeof(PokerCard.int_to_pretty_str.(hand_ints)) == Array{String,1}
    @test typeof(Evaluator.get_rank_class(Evaluator.evaluate(hand_ints))) == Int

    # some timing tests for 1000 random draws and evals
    t = @elapsed for i in 1:1000
        hand_ints = PokerDeck.draw!(7, PokerDeck.make_deck())
        Evaluator.get_rank_class(Evaluator.evaluate(hand_ints))
    end # elapsed time

    # 1k draws and evals in under 1s
    @assert t < 1

end # tests


using lmp_str

file_out = "test.data"

data_cell = genr_cell([5, 5, 5])
str = structure_si3n4()
data = genr(data_cell, str)
write_data(data, file_out)


using Turing, Random
using Turing: Selector, reconstruct, invlink, CACHERESET,
    SampleFromPrior, Sampler, runmodel!, SampleFromUniform
using Turing.RandomVariables
using Turing.RandomVariables: uid, _getidcs,
    set_retained_vns_del_by_spl!, is_flagged,
    set_flag!, unset_flag!, VarInfo, TypedVarInfo,
    getlogp, setlogp!, resetlogp!, acclogp!
using Distributions
using ForwardDiff: Dual
using Test

i, j, k = 1, 2, 3

dir = splitdir(splitdir(pathof(Turing))[1])[1]
include(dir*"/test/test_utils/AllUtils.jl")

@testset "RandomVariables.jl" begin
    @turing_testset "TypedVarInfo" begin
        @model gdemo(x, y) = begin
            s ~ InverseGamma(2,3)
            m ~ TruncatedNormal(0.0,sqrt(s),0.0,2.0)
            x ~ Normal(m, sqrt(s))
            y ~ Normal(m, sqrt(s))
        end
        model = gdemo(1.0, 2.0)

        vi = VarInfo()
        model(vi, SampleFromUniform())
        tvi = TypedVarInfo(vi)

        meta = vi.metadata
        for f in fieldnames(typeof(tvi.metadata))
            fmeta = getfield(tvi.metadata, f)
            for vn in fmeta.vns
                @test tvi[vn] == vi[vn]
                ind = meta.idcs[vn]
                tind = fmeta.idcs[vn]
                @test meta.dists[ind] == fmeta.dists[tind]
                @test meta.orders[ind] == fmeta.orders[tind]
                @test meta.gids[ind] == fmeta.gids[tind]
                for flag in keys(meta.flags)
                    @test meta.flags[flag][ind] == fmeta.flags[flag][tind]
                end
                range = meta.ranges[ind]
                trange = fmeta.ranges[tind]
                @test all(meta.vals[range] .== fmeta.vals[trange])
            end
        end
    end
    @turing_testset "Base" begin
        # Test Base functions:
        #   string, Symbol, ==, hash, in, keys, haskey, isempty, push!, empty!,
        #   getindex, setindex!, getproperty, setproperty!
        csym = gensym()
        vn1 = VarName(csym, :x, "[1][2]", 1)
        @test string(vn1) == "{$csym,x[1][2]}:1"
        @test string(vn1, all=false) == "x[1][2]"
        @test Symbol(vn1) == Symbol("x[1][2]")

        vn2 = VarName(csym, :x, "[1][2]", 1)
        @test vn2 == vn1
        @test hash(vn2) == hash(vn1)
        @test in(vn1, (:x,))

        function test_base!(vi)
            empty!(vi)
            @test vi.logp == 0
            @test vi.num_produce == 0

            vn = VarName(gensym(), :x, "", 1)
            dist = Normal(0, 1)
            r = rand(dist)
            gid = Selector()

            @test isempty(vi)
            @test ~haskey(vi, vn)
            push!(vi, vn, r, dist, gid)
            @test ~isempty(vi)
            @test haskey(vi, vn)

            @test length(vi[vn]) == 1
            @test length(vi[SampleFromPrior()]) == 1

            @test vi[vn] == r
            @test vi[SampleFromPrior()][1] == r
            vi[vn] = [2*r]
            @test vi[vn] == 2*r
            @test vi[SampleFromPrior()][1] == 2*r
            vi[SampleFromPrior()] = [3*r]
            @test vi[vn] == 3*r
            @test vi[SampleFromPrior()][1] == 3*r

            empty!(vi)
            @test isempty(vi)
            push!(vi, vn, r, dist, gid)

            function test_in()
                space = (:x, :y, :(z[1]))
                vn1 = genvn(:x)
                vn2 = genvn(:y)
                vn3 = genvn(:(x[1]))
                vn4 = genvn(:(z[1][1]))
                vn5 = genvn(:(z[2]))
                vn6 = genvn(:z)

                @test in(vn1, space)
                @test in(vn2, space)
                @test in(vn3, space)
                @test in(vn4, space)
                @test ~in(vn5, space)
                @test ~in(vn6, space)
            end
            test_in()
        end
        vi = VarInfo()
        test_base!(vi)
        test_base!(empty!(TypedVarInfo(vi)))
    end
    @turing_testset "runmodel!" begin
        # Test that eval_num is incremented when calling runmodel!
        @model testmodel() = begin
            x ~ Normal()
        end
        alg = HMC(0.1, 5)
        spl = Sampler(alg, testmodel())
        vi = VarInfo()
        m = testmodel()
        m(vi)
        runmodel!(m, vi, spl)
        @test spl.state.eval_num == 1
        runmodel!(m, vi, spl)
        @test spl.state.eval_num == 2
    end
    @turing_testset "flags" begin
        # Test flag setting:
        #    is_flagged, set_flag!, unset_flag!
        function test_varinfo!(vi)
            vn_x = VarName(gensym(), :x, "", 1)
            dist = Normal(0, 1)
            r = rand(dist)
            gid = Selector()

            push!(vi, vn_x, r, dist, gid)

            # del is set by default
            @test is_flagged(vi, vn_x, "del") == false

            set_flag!(vi, vn_x, "del")
            @test is_flagged(vi, vn_x, "del") == true

            unset_flag!(vi, vn_x, "del")
            @test is_flagged(vi, vn_x, "del") == false
        end
        vi = VarInfo()
        test_varinfo!(vi)
        test_varinfo!(empty!(TypedVarInfo(vi)))
    end
    @turing_testset "link!" begin
        # Test linking spl and vi:
        #    link!, invlink!, istrans
        @model gdemo(x, y) = begin
            s ~ InverseGamma(2,3)
            m ~ Uniform(0, 2)
            x ~ Normal(m, sqrt(s))
            y ~ Normal(m, sqrt(s))
        end
        model = gdemo(1.0, 2.0)

        vi = VarInfo()
        model(vi, SampleFromUniform())

        @test all(x -> ~istrans(vi, x), vi.vns)
        alg = HMC(0.1, 5)
        spl = Sampler(alg, model)
        v = copy(vi.vals)
        link!(vi, spl)
        @test all(x -> istrans(vi, x), vi.vns)
        invlink!(vi, spl)
        @test all(x -> ~istrans(vi, x), vi.vns)
        @test vi.vals == v

        vi = TypedVarInfo(vi)
        alg = HMC(0.1, 5)
        spl = Sampler(alg, model)
        @test all(x -> ~istrans(vi, x), vi.metadata.s.vns)
        @test all(x -> ~istrans(vi, x), vi.metadata.m.vns)
        v_s = copy(vi.metadata.s.vals)
        v_m = copy(vi.metadata.m.vals)
        link!(vi, spl)
        @test all(x -> istrans(vi, x), vi.metadata.s.vns)
        @test all(x -> istrans(vi, x), vi.metadata.m.vns)
        invlink!(vi, spl)
        @test all(x -> ~istrans(vi, x), vi.metadata.s.vns)
        @test all(x -> ~istrans(vi, x), vi.metadata.m.vns)
        @test vi.metadata.s.vals == v_s
        @test vi.metadata.m.vals == v_m
    end
    @turing_testset "setgid!" begin
        vi = VarInfo()
        vn = VarName(gensym(), :x, "", 1)
        dist = Normal(0, 1)
        r = rand(dist)
        gid1 = Selector()
        gid2 = Selector(2, :HMC)

        push!(vi, vn, r, dist, gid1)
        @test vi.gids[vi.idcs[vn]] == Set([gid1])
        setgid!(vi, gid2, vn)
        @test vi.gids[vi.idcs[vn]] == Set([gid1, gid2])

        vi = empty!(TypedVarInfo(vi))
        push!(vi, vn, r, dist, gid1)
        @test vi.metadata.x.gids[vi.metadata.x.idcs[vn]] == Set([gid1])
        setgid!(vi, gid2, vn)
        @test vi.metadata.x.gids[vi.metadata.x.idcs[vn]] == Set([gid1, gid2])
    end
    @testset "orders" begin
        function randr(vi::VarInfo, vn::VarName, dist::Distribution, spl::Turing.Sampler)
            if ~haskey(vi, vn)
                r = rand(dist)
                Turing.push!(vi, vn, r, dist, spl)
                r
            elseif is_flagged(vi, vn, "del")
                unset_flag!(vi, vn, "del")
                r = rand(dist)
                vi[vn] = Turing.vectorize(dist, r)
                Turing.setorder!(vi, vn, vi.num_produce)
                r
            else
                Turing.updategid!(vi, vn, spl)
                vi[vn]
            end
        end

        csym = gensym() # unique per model
        vn_z1 = VarName(csym, :z, "[1]", 1)
        vn_z2 = VarName(csym, :z, "[2]", 1)
        vn_z3 = VarName(csym, :z, "[3]", 1)
        vn_z4 = VarName(csym, :z, "[4]", 1)
        vn_a1 = VarName(csym, :a, "[1]", 1)
        vn_a2 = VarName(csym, :a, "[2]", 1)
        vn_b = VarName(csym, :b, "", 1)

        vi = VarInfo()
        dists = [Categorical([0.7, 0.3]), Normal()]

        spl1 = Turing.Sampler(PG(5), empty_model())
        spl2 = Turing.Sampler(PG(5), empty_model())

        # First iteration, variables are added to vi
        # variables samples in order: z1,a1,z2,a2,z3
        vi.num_produce += 1
        randr(vi, vn_z1, dists[1], spl1)
        randr(vi, vn_a1, dists[2], spl1)
        vi.num_produce += 1
        randr(vi, vn_b, dists[2], spl2)
        randr(vi, vn_z2, dists[1], spl1)
        randr(vi, vn_a2, dists[2], spl1)
        vi.num_produce += 1
        randr(vi, vn_z3, dists[1], spl1)
        @test vi.orders == [1, 1, 2, 2, 2, 3]
        @test vi.num_produce == 3

        vi.num_produce = 0
        set_retained_vns_del_by_spl!(vi, spl1)
        @test is_flagged(vi, vn_z1, "del")
        @test is_flagged(vi, vn_a1, "del")
        @test is_flagged(vi, vn_z2, "del")
        @test is_flagged(vi, vn_a2, "del")
        @test is_flagged(vi, vn_z3, "del")

        vi.num_produce += 1
        randr(vi, vn_z1, dists[1], spl1)
        randr(vi, vn_a1, dists[2], spl1)
        vi.num_produce += 1
        randr(vi, vn_z2, dists[1], spl1)
        vi.num_produce += 1
        randr(vi, vn_z3, dists[1], spl1)
        randr(vi, vn_a2, dists[2], spl1)
        @test vi.orders == [1, 1, 2, 2, 3, 3]
        @test vi.num_produce == 3

        vi = empty!(TypedVarInfo(vi))
        # First iteration, variables are added to vi
        # variables samples in order: z1,a1,z2,a2,z3
        vi.num_produce += 1
        randr(vi, vn_z1, dists[1], spl1)
        randr(vi, vn_a1, dists[2], spl1)
        vi.num_produce += 1
        randr(vi, vn_b, dists[2], spl2)
        randr(vi, vn_z2, dists[1], spl1)
        randr(vi, vn_a2, dists[2], spl1)
        vi.num_produce += 1
        randr(vi, vn_z3, dists[1], spl1)
        @test vi.metadata.z.orders == [1, 2, 3]
        @test vi.metadata.a.orders == [1, 2]
        @test vi.metadata.b.orders == [2]
        @test vi.num_produce == 3

        vi.num_produce = 0
        set_retained_vns_del_by_spl!(vi, spl1)
        @test is_flagged(vi, vn_z1, "del")
        @test is_flagged(vi, vn_a1, "del")
        @test is_flagged(vi, vn_z2, "del")
        @test is_flagged(vi, vn_a2, "del")
        @test is_flagged(vi, vn_z3, "del")

        vi.num_produce += 1
        randr(vi, vn_z1, dists[1], spl1)
        randr(vi, vn_a1, dists[2], spl1)
        vi.num_produce += 1
        randr(vi, vn_z2, dists[1], spl1)
        vi.num_produce += 1
        randr(vi, vn_z3, dists[1], spl1)
        randr(vi, vn_a2, dists[2], spl1)
        @test vi.metadata.z.orders == [1, 2, 3]
        @test vi.metadata.a.orders == [1, 3]
        @test vi.metadata.b.orders == [2]
        @test vi.num_produce == 3
    end
    @turing_testset "replay" begin
        # Generate synthesised data
        xs = rand(Normal(0.5, 1), 100)

        # Define model
        @model priorsinarray(xs, ::Type{T}=Float64) where {T} = begin
            priors = Vector{T}(undef, 2)
            priors[1] ~ InverseGamma(2, 3)
            priors[2] ~ Normal(0, sqrt(priors[1]))
            for i = 1:length(xs)
                xs[i] ~ Normal(priors[2], sqrt(priors[1]))
            end
            priors
        end

        # Sampling
        chain = sample(priorsinarray(xs), HMC(0.01, 10), 10)
    end
    @turing_testset "varname" begin
        csym = gensym()
        vn1 = VarName(csym, :x, "[1]", 1)
        @test vn1 == VarName{:x}(csym, "[1]", 1)
        @test vn1 == VarName{:x}(csym, "[1]")
        @test vn1 == VarName([csym, :x], "[1]")
        vn2 = VarName(csym, :x, "[2]", 1)
        @test vn2 == VarName(vn1, "[2]")

        # Symbol
        v_sym = string(:x)
        @test v_sym == "x"

        # Array
        v_arr = eval(varname(:(x[i]))[1])
        @test v_arr == "[1]"

        # Matrix
        v_mat = eval(varname(:(x[i,j]))[1])
        @test v_mat== "[1,2]"

        v_mat = eval(varname(:(x[i,j,k]))[1])
        @test v_mat== "[1,2,3]"

        v_mat = eval(varname(:((x[1,2][1+5][45][3][i])))[1])
        @test v_mat == "[1,2][6][45][3][1]"

        @model mat_name_test() = begin
            p = Array{Any}(undef, 2, 2)
            for i in 1:2, j in 1:2
                p[i,j] ~ Normal(0, 1)
            end
            p
        end
        chain = sample(mat_name_test(), HMC(0.2, 4), 1000)
        check_numerical(chain, ["p[1, 1]"], [0], atol = 0.25)

        # Multi array
        v_arrarr = eval(varname(:(x[i][j]))[1])
        @test v_arrarr == "[1][2]"

        @model marr_name_test() = begin
            p = Array{Array{Any}}(undef, 2)
            p[1] = Array{Any}(undef, 2)
            p[2] = Array{Any}(undef, 2)
            for i in 1:2, j in 1:2
                p[i][j] ~ Normal(0, 1)
            end
            p
        end

        chain = sample(marr_name_test(), HMC(0.2, 4), 1000)
        check_numerical(chain, ["p[1][1]"], [0], atol = 0.25)
    end
    @turing_testset "varinfo" begin
        dists = [Normal(0, 1), MvNormal([0; 0], [1.0 0; 0 1.0]), Wishart(7, [1 0.5; 0.5 1])]
        function test_varinfo!(vi)
            @test getlogp(vi) == 0
            setlogp!(vi, 1)
            @test getlogp(vi) == 1
            acclogp!(vi, 1)
            @test getlogp(vi) == 2
            resetlogp!(vi)
            @test getlogp(vi) == 0

            spl2 = Turing.Sampler(PG(5, :w, :u), empty_model())
            vn_w = VarName(gensym(), :w, "", 1)
            randr(vi, vn_w, dists[1], spl2, true)

            vn_x = VarName(gensym(), :x, "", 1)
            vn_y = VarName(gensym(), :y, "", 1)
            vn_z = VarName(gensym(), :z, "", 1)
            vns = [vn_x, vn_y, vn_z]

            spl1 = Turing.Sampler(PG(5, :x, :y, :z), empty_model())
            for i = 1:3
                r = randr(vi, vns[i], dists[i], spl1, false)
                val = vi[vns[i]]
                @test sum(val - r) <= 1e-9
            end

            idcs = _getidcs(vi, spl1)
            if idcs isa NamedTuple
                @test sum(length(getfield(idcs, f)) for f in fieldnames(typeof(idcs))) == 3
            else
                @test length(idcs) == 3
            end
            @test length(vi[spl1]) == 7

            idcs = _getidcs(vi, spl2)
            if idcs isa NamedTuple
                @test sum(length(getfield(idcs, f)) for f in fieldnames(typeof(idcs))) == 1
            else
                @test length(idcs) == 1
            end
            @test length(vi[spl2]) == 1

            vn_u = VarName(gensym(), :u, "", 1)
            randr(vi, vn_u, dists[1], spl2, true)

            idcs = _getidcs(vi, spl2)
            if idcs isa NamedTuple
                @test sum(length(getfield(idcs, f)) for f in fieldnames(typeof(idcs))) == 2
            else
                @test length(idcs) == 2
            end
            @test length(vi[spl2]) == 2
        end
        vi = VarInfo()
        test_varinfo!(vi)
        test_varinfo!(empty!(TypedVarInfo(vi)))

        @model igtest() = begin
            x ~ InverseGamma(2,3)
            y ~ InverseGamma(2,3)
            z ~ InverseGamma(2,3)
            w ~ InverseGamma(2,3)
            u ~ InverseGamma(2,3)
        end

        # Test the update of group IDs
        g_demo_f = igtest()

        # This test section no longer seems as applicable, considering the
        # user will never end up using an UntypedVarInfo. The `VarInfo`
        # Varible is also not passed around in the same way as it used to be.
        g = Turing.Sampler(Gibbs(PG(10, :x, :y, :z), HMC(0.4, 8, :w, :u)), g_demo_f)
        pg, hmc = g.state.samplers
        vi = VarInfo()
        g_demo_f(vi, SampleFromPrior())
        Turing.Inference.step!(Random.GLOBAL_RNG, g_demo_f, pg, 1)
        vi1 = pg.state.vi
        @test mapreduce(x -> x.gids, vcat, vi1.metadata) ==
            [Set([pg.selector]), Set([pg.selector]), Set([pg.selector]), Set{Selector}(), Set{Selector}()]

        g_demo_f(vi1, hmc)
        @test mapreduce(x -> x.gids, vcat, vi1.metadata) ==
            [Set([pg.selector]), Set([pg.selector]), Set([pg.selector]), Set([hmc.selector]), Set([hmc.selector])]

        g = Turing.Sampler(Gibbs(PG(10, :x, :y, :z), HMC(0.4, 8, :w, :u)), g_demo_f)
        pg, hmc = g.state.samplers
        vi = empty!(TypedVarInfo(vi))
        g_demo_f(vi, SampleFromPrior())
        pg.state.vi = vi
        Turing.Inference.step!(Random.GLOBAL_RNG, g_demo_f, pg, 1)
        vi = pg.state.vi
        g_demo_f(vi, hmc)
        @test vi.metadata.x.gids[1] == Set([pg.selector])
        @test vi.metadata.y.gids[1] == Set([pg.selector])
        @test vi.metadata.z.gids[1] == Set([pg.selector])
        @test vi.metadata.w.gids[1] == Set([hmc.selector])
        @test vi.metadata.u.gids[1] == Set([hmc.selector])
    end
end


@inline function check_simulation_stop!(sol, t::Float64, Y, YP, run::AbstractRun, p, bounds, opts::options_simulation_immutable{T};
    ϵ::Float64 = t < 1.0 ? opts.reltol : 0.0,
    ) where T<:Function

    if t ≥ run.tf
        run.info.flag = 0
        run.info.exit_reason = "Final time reached"
        return nothing
    end

    # continue checking the bounds or return after only evaluating the time
    !opts.check_bounds && (return nothing)

    I = calc_I(Y,p)

    check_stop_I(         p, run, sol, Y, YP, bounds, ϵ, I)
    check_stop_V(         p, run, sol, Y, YP, bounds, ϵ, I)
    check_stop_SOC(       p, run, sol, Y, YP, bounds, ϵ, I)
    check_stop_T(         p, run, sol, Y, YP, bounds, ϵ, I)
    check_stop_c_s_surf(  p, run, sol, Y, YP, bounds, ϵ, I)
    check_stop_c_e(       p, run, sol, Y, YP, bounds, ϵ, I)
    check_stop_η_plating( p, run, sol, Y, YP, bounds, ϵ, I)
    check_stop_dfilm(     p, run, sol, Y, YP, bounds, ϵ, I)
    opts.stop_function(   p, run, sol, Y, YP, bounds, ϵ, I)

    return nothing
end

@inline check_stop_I(p::R4, run::R3, sol, Y::R2, YP::R2, bounds::R5, ϵ::Float64, I::Float64
    ) where {R2<:Vector{Float64}, R3<:AbstractRun{method_I},R4<:model,R5<:boundary_stop_conditions_immutable} = nothing
@inline function check_stop_I(p::R4, run::R3, sol, Y::R2, YP::R2, bounds::R5, ϵ::Float64, I::Float64
    ) where {R2<:Vector{Float64}, method, R3<:AbstractRun{method},R4<:model,R5<:boundary_stop_conditions_immutable}

    if (I - bounds.I_max > ϵ)
        t_frac = (bounds.prev.I - bounds.I_max)/(bounds.prev.I - I)
        if t_frac < bounds.prev.t_final_interp_frac
            bounds.prev.t_final_interp_frac = t_frac
            run.info.flag = 7
            run.info.exit_reason = "Above maximum permitted C-rate"
        end
    elseif (bounds.I_min - I > ϵ)
        t_frac = (bounds.prev.I - bounds.I_min)/(bounds.prev.I - I)
        if t_frac < bounds.prev.t_final_interp_frac
            bounds.prev.t_final_interp_frac = t_frac
            run.info.flag = 8
            run.info.exit_reason = "Below minimum permitted C-rate"
        end
    end
    bounds.prev.I = I

    return nothing
end

@inline check_stop_V(p::R4, run::R3, sol, Y::R2, YP::R2, bounds::R5, ϵ::Float64, I::Float64
) where {R2<:Vector{Float64}, R3<:AbstractRun{method_V},R4<:model,R5<:boundary_stop_conditions_immutable} = nothing
@inline function check_stop_V(p::R4, run::R3, sol, Y::R2, YP::R2, bounds::R5, ϵ::Float64, I::Float64
    ) where {R2<:Vector{Float64}, method, R3<:AbstractRun{method},R4<:model,R5<:boundary_stop_conditions_immutable}

    V = calc_V(Y,p)
    if (bounds.V_min - V > ϵ) && I < 0
        t_frac = (bounds.prev.V - bounds.V_min)/(bounds.prev.V - V)
        if t_frac < bounds.prev.t_final_interp_frac
            bounds.prev.t_final_interp_frac = t_frac
            run.info.flag = 1
            run.info.exit_reason = "Below minimum voltage limit"
        end
    elseif (V - bounds.V_max > ϵ) && I > 0
        t_frac = (bounds.prev.V - bounds.V_max)/(bounds.prev.V - V)
        if t_frac < bounds.prev.t_final_interp_frac
            bounds.prev.t_final_interp_frac = t_frac
            run.info.flag = 2
            run.info.exit_reason = "Above maximum voltage limit"
        end
    end
    bounds.prev.V = V

    return nothing
end

@inline function check_stop_SOC(p::R4, run::R3, sol::solution, Y::R2, YP::R2, bounds::R5, ϵ::Float64, I::Float64
    ) where {R2<:Vector{Float64}, R3<:AbstractRun,R4<:model,R5<:boundary_stop_conditions_immutable}

    SOC = @inbounds sol.SOC[end]
    if (bounds.SOC_min - SOC > ϵ) && I < 0
        t_frac = (bounds.prev.SOC - bounds.SOC_min)/(bounds.prev.SOC - SOC)
        if t_frac < bounds.prev.t_final_interp_frac
            bounds.prev.t_final_interp_frac = t_frac
            run.info.flag = 3
            run.info.exit_reason = "Below minimum SOC limit"
        end
    elseif (SOC - bounds.SOC_max > ϵ) && I > 0
        t_frac = (bounds.prev.SOC - bounds.SOC_max)/(bounds.prev.SOC - SOC)
        if t_frac < bounds.prev.t_final_interp_frac
            bounds.prev.t_final_interp_frac = t_frac
            run.info.flag = 4
            run.info.exit_reason = "Above maximum SOC limit"
        end
    end
    bounds.prev.SOC = SOC

    return SOC
end

@inline check_stop_T(p::model_temp{false}, run, sol, Y, YP, bounds, ϵ, I) = nothing
@inline function check_stop_T(p::R4, run::R3, sol, Y::R2, YP::R2, bounds::R5, ϵ::Float64, I::Float64
    ) where {R2<:Vector{Float64}, R3<:AbstractRun,R4<:model_temp{true},R5<:boundary_stop_conditions_immutable}

    if !isnan(bounds.T_max)
        T = temperature_weighting(calc_T(Y,p),p)
        if T - bounds.T_max > ϵ
            t_frac = (bounds.prev.T - bounds.T_max)/(bounds.prev.T - T)
            if t_frac < bounds.prev.t_final_interp_frac
                bounds.prev.t_final_interp_frac = t_frac
                run.info.flag = 5
                run.info.exit_reason = "Above maximum permitted temperature"
            end
        end
        bounds.prev.T = T
    end

    return nothing
end

@inline function c_s_n_maximum(Y::Vector{Float64},p::model_solid_diff{:Fickian})
    c_s_n_max = -Inf
    @inbounds for i in 1:p.N.n
        @inbounds c_s_n_max = max(c_s_n_max, Y[p.ind.c_s_avg.n[p.N.r_n*i]])
    end
    return c_s_n_max
end
@inline function c_s_n_maximum(Y::Vector{Float64},p::Union{model_solid_diff{:polynomial},model_solid_diff{:quadratic}})
    c_s_n_max = -Inf
    @inbounds for ind in p.ind.c_s_avg.n
        @inbounds c_s_n_max = max(c_s_n_max, Y[ind])
    end
    return c_s_n_max
end

@inline function check_stop_c_s_surf(p::R4, run::R3, sol, Y::R2, YP::R2, bounds::R5, ϵ::Float64, I::Float64
    ) where {R2<:Vector{Float64}, R3<:AbstractRun,R4<:model,R5<:boundary_stop_conditions_immutable}

    if !isnan(bounds.c_s_n_max)
        c_s_n_max = c_s_n_maximum(Y,p)

        if I > 0
            if c_s_n_max - bounds.c_s_n_max*p.θ[:c_max_n] > ϵ
                t_frac = (bounds.prev.c_s_n - bounds.c_s_n_max*(p.θ[:c_max_n]))/(bounds.prev.c_s_n - c_s_n_max)
                if t_frac < bounds.prev.t_final_interp_frac
                    bounds.prev.t_final_interp_frac = t_frac
                    run.info.flag = 6
                    run.info.exit_reason = "Above c_s_n saturation threshold"
                end
            end
        end
        bounds.prev.c_s_n = c_s_n_max
    end

    return nothing
end

@inline function check_stop_c_e(p::R4, run::R3, sol, Y::R2, YP::R2, bounds::R5, ϵ::Float64, I::Float64
    ) where {R2<:Vector{Float64}, R3<:AbstractRun, R4<:model, R5<:boundary_stop_conditions_immutable}

    if !isnan(bounds.c_e_min)
        c_e_min = +Inf
        @inbounds for ind in p.ind.c_e
            @inbounds c_e_min = min(c_e_min, Y[ind])
        end
        if bounds.c_e_min - c_e_min > ϵ
            t_frac = (bounds.prev.c_e_min - bounds.c_e_min)/(bounds.prev.c_e_min - c_e_min)
            if t_frac < bounds.prev.t_final_interp_frac
                bounds.prev.t_final_interp_frac = t_frac
                run.info.flag = 9
                run.info.exit_reason = "Below minimum permitted c_e"
            end
        end
        bounds.prev.c_e_min = c_e_min
    end

    return nothing
end

@inline function check_stop_η_plating(p::R4, run::R3, sol, Y::R2, YP::R2, bounds::R5, ϵ::Float64, I::Float64
    ) where {R2<:Vector{Float64}, R3<:AbstractRun, R4<:model, R5<:boundary_stop_conditions_immutable}

    η_plating = calc_η_plating(Y,p)
    if !isnan(bounds.η_plating_min) && bounds.η_plating_min - η_plating > ϵ
        t_frac = (bounds.prev.η_plating - bounds.η_plating_min)/(bounds.prev.η_plating - η_plating)
        if t_frac < bounds.prev.t_final_interp_frac
            bounds.prev.t_final_interp_frac = t_frac
            run.info.flag = 9
            run.info.exit_reason = "Below minimum permitted η_plating"
        end
    end
    bounds.prev.η_plating = η_plating

    return nothing
end

@inline check_stop_dfilm(::model_age{false}, run, sol, Y, YP, bounds, ϵ, I) = nothing
@inline function check_stop_dfilm(p::R4, run::R3, sol, Y::R2, YP::R2, bounds::R5, ϵ::Float64, I::Float64
    ) where {R2<:Vector{Float64}, R3<:AbstractRun, R4<:model_age{:SEI}, R5<:boundary_stop_conditions_immutable}

    dfilm_max = -Inf
    @inbounds for i in 1:p.N.n
        @inbounds dfilm_max = max(dfilm_max, YP[p.ind.film[i]])
    end

    if !isnan(bounds.dfilm_max) && dfilm_max - bounds.dfilm_max > ϵ
        t_frac = (bounds.prev.dfilm - bounds.dfilm_max)/(bounds.prev.dfilm - dfilm_max)
        if t_frac < bounds.prev.t_final_interp_frac
            bounds.prev.t_final_interp_frac = t_frac
            run.info.flag = 10
            run.info.exit_reason = "Above maximum film growth rate"
        end
    end
    bounds.prev.dfilm = dfilm_max

    return nothing
end

@inline function check_solve(run::run_constant, sol::R1, int::R2, p, bounds, opts::R5, funcs, keep_Y::Bool, iter::Int64, Y::Vector{Float64}, t::Float64) where {R1<:solution,R2<:Sundials.IDAIntegrator,R5<:AbstractOptionsModel}
    if t == int.tprev
        # Sometimes the initial step at t = 0 can be too large. This reduces the step size
        if t == 0.0
            if iter == 2
                Sundials.IDASetInitStep(int.mem,opts.reltol)
            else
                error("Model failed to converge at t = $(t). Try tightening the absolute and relative tolerances.")
            end
        else
            error("Model failed to converge at t = $(t)")
        end
    elseif iter == opts.maxiters
        error("Reached max iterations of $(opts.maxiters) at t = $(t)")
    elseif within_bounds(run)
        # update Y only after checking the stop conditions. this is done to store a copy of the
        # previous sol run in case any back-interpolation is needed
        set_var!(sol.Y, keep_Y ? copy(Y) : Y, keep_Y)
    else # no errors and run.info.flag ≠ -1
        return false
    end

    return true
end

@inline function check_solve(run::run_function, sol::R1, int::R2, p::model, bounds::boundary_stop_conditions_immutable, opts::R5, funcs, keep_Y::Bool, iter::Int64, Y::Vector{Float64}, t::Float64) where {R1<:solution,R2<:Sundials.IDAIntegrator,R5<:AbstractOptionsModel}
    if iter == opts.maxiters
        error("Reached max iterations of $(opts.maxiters) at t = $(int.t)")
    elseif within_bounds(run)
        # update Y only after checking the stop conditions. this is done to store a copy of the
        # previous sol run in case any back-interpolation is needed
        set_var!(sol.Y, keep_Y ? copy(Y) : Y, keep_Y)

        # check to see if the run needs to be reinitialized
        if t - int.tprev < 1e-3opts.reltol
            check_reinitialization!(sol, int, run, p, bounds, opts, funcs)
        end

        return true
    else # no errors and run.info.flag ≠ -1
        return false
    end
end

@inline function check_reinitialization!(sol::R1, int::R2, run::R3, p::R4, bounds::R5, opts::R6, funcs) where {R1<:solution, R2<:Sundials.IDAIntegrator, R3<:AbstractRun,R4<:model,R5<:boundary_stop_conditions_immutable,R6<:AbstractOptionsModel}
    """
    Checking the current function for discontinuities.
    If there is a significant change in current after a step size of dt = reltol,
    then rerun Newton's method and reinitialize at this new time step.
    """

    Y = int.u.v
    YP = int.du.v
    # take a step of Δt = the relative tolerance
    t_new = int.t + opts.reltol

    value_old = value(run)
    value_new = run.func(t_new,Y,YP,p)

    # if the function values at t vs. t + Δt are very different (i.e., there is a discontinuity)
    # then reinitialize the DAE at t + Δt
    if !≈(value_old, value_new, atol=opts.abstol, rtol=opts.reltol)
        initialize_states!(p,Y,YP,run,opts,funcs,(@inbounds sol.SOC[end]); t=t_new)

        Sundials.IDAReInit(int.mem, t_new, Y, YP)
    end
    return nothing
end

@inline function check_errors_parameters_runtime(p::R1,opts::R2) where {R1<:model,R2<:AbstractOptionsModel}
    ϵ_sp, ϵ_sn = active_material(p)

    if ( ϵ_sp > 1 ) error("ϵ_p + ϵ_fp must be ∈ [0, 1)") end
    if ( ϵ_sn > 1 ) error("ϵ_n + ϵ_fn must be ∈ [0, 1)") end
    if ( p.θ[:θ_max_p] > p.θ[:θ_min_p] ) error("θ_max_p must be < θ_min_p") end
    if ( p.θ[:θ_min_n] > p.θ[:θ_max_n] ) error("θ_min_n must be < θ_max_n") end

    return nothing
end

function check_errors_initial(θ, numerics, N)
    if !(numerics.jacobian ∈ (:symbolic, :AD))
        error("`jacobian` can either be :symbolic or :AD")
    end

    return nothing
end

check_is_hold(x::Symbol,sol::solution) = (x == :hold) && (!isempty(sol) ? true : error("Cannot use `:hold` without a previous sol."))
check_is_hold(::Any,::solution) = false


module infer

function identify(causal_model)
end

function estimate(statstical_model)
end

end


# From AnticNumberField into something

mutable struct MapDataFromAnticNumberField{T}
  prim_image::T
  isid::Bool

  function MapDataFromAnticNumberField{T}(x::Bool) where T
    @assert x
    z = new{T}()
    z.isid = true
    return z
  end

  function MapDataFromAnticNumberField{T}(x::T) where T
    z = new{T}(x, false)
    return z
  end
end

function image(f::MapDataFromAnticNumberField, y::nf_elem)
  f.isid && return y
  z = parent(defining_polynomial(parent(y)))(y)
  return evaluate(z, f.prim_image)
end

function (f::MapDataFromAnticNumberField)(y::nf_elem)
  return image(f, y)
end

map_data_type(K::AnticNumberField, L) = MapDataFromAnticNumberField{elem_type(L)}

map_data(K::AnticNumberField, L, y) = MapDataFromAnticNumberField{typeof(y)}(y)

map_data(K::AnticNumberField, ::Bool) = MapDataFromAnticNumberField{nf_elem}(true)

map_data(K::AnticNumberField, L, ::Bool) = MapDataFromAnticNumberField{elem_type(L)}(true)

# From NfRel into something
mutable struct MapDataFromNfRel{T, S}
  prim_image::T
  base_field_map_data::S
  isid::Bool

  function MapDataFromNfRel{T, S}(x::T, y::S) where {T, S}
    z = new{T, S}(x, y, false)
    return z
  end

  function MapDataFromNfRel{T, S}(x::Bool) where {T, S}
    @assert x
    z = new{T, S}()
    z.isid = true
    return z
  end
end

function image(f::MapDataFromNfRel, y)
  f.isid && return y
  z = change_base_ring(y.data, f.base_field_map_data)
  return evaluate(z, f.prim_image)
end

function (f::MapDataFromNfRel)(y)
  return image(f, y)
end

map_data_type(K::NfRel, L) = MapDataFromNfRel{elem_type(L), map_data_type(base_field(K), L)}

map_data_type(K::NfRel, L, base_field_map_data::S) where {S} = MapDataFromNfRel{elem_type(L), S}

map_data(K::NfRel, L, y, g) = MapDataFromNfRel{typeof(y), typeof(g.image_data)}(y, g.image_data)

map_data(K::NfRel, L, y) = MapDataFromNfRel{typeof(y), map_data_type(base_field(K), L)}(y, map_data_type(base_field(K), L)(true))

map_data(K::NfRel, L, ::Bool) = MapDataFromNfRel{elem_type(L), map_data_type(base_field(K), L)}(true)

function map_data(K::NfRel, y)
  z = MapDataFromNfRel{typeof(y), map_data_type(base_field(K), parent(y))}
  return z(y, map_data(base_field(K), K, true))
end

# From NfAbsNS into something
mutable struct MapDataFromNfAbsNS{T}
  images::T
  isid::Bool

  function MapDataFromNfAbsNS{T}(x::T) where {T}
    z = new{T}(x, false)
    return z
  end

  function MapDataFromNfAbsNS{T}(x::Bool) where {T}
    @assert x
    z = new{T}()
    z.isid = true
    return z
  end
end

function image(f::MapDataFromNfAbsNS, y)
  f.isid && return y
  return evaluate(y.data, f.images)
end

function (f::MapDataFromNfAbsNS)(y)
  return image(f, y)
end

map_data_type(K::NfAbsNS, L) = MapDataFromNfAbsNS{Vector{elem_type(L)}}

map_data(K::NfAbsNS, L, y) = MapDataFromNfAbsNS{typeof(y)}(y)

map_data(K::NfAbsNS, L, ::Bool) = MapDataFromNfAbsNS{Vector{elem_type(L)}}(true)

# From NfRelNS into something
mutable struct MapDataFromNfRelNS{T, S}
  images::T
  base_field_map_data::S
  isid::Bool

  function MapDataFromNfRelNS{T, S}(x::T, y::S) where {T, S}
    z = new{T, S}(x, y, false)
    return z
  end

  function MapDataFromNfRelNS{T, S}(x::Bool) where {T, S}
    @assert x
    z = new{T, S}()
    z.isid = true
    return z
  end
end

function image(f::MapDataFromNfRelNS, y)
  z = change_base_ring(y.data, f.base_field_map_data)
  return evaluate(z, f.images)
end

function (f::MapDataFromNfRelNS)(y)
  return image(f, y)
end

map_data_type(K::NfRelNS, L) = MapDataFromNfRelNS{Vector{elem_type(L)}, map_data_type(base_field(K), L)}

map_data_type(K::NfRelNS, L, base_field_map_data::S) where {S} = MapDataFromNfRelNS{Vector{elem_type(L)}, S}

# Here g is a NumFieldMor
map_data(K::NfRelNS, L, y, g) = MapDataFromNfRelNS{typeof(y), typeof(g.image_data)}(y, g.image_data)

map_data(K::NfRelNS, L, y) = MapDataFromNfRelNS{typeof(y), map_data_type(base_field(K), L)}(y, map_data_type(base_field(K), L)(true))

#map_data(K::NfRelNS, ::Bool) = MapDataFromNfRelNS{Vector{elem_type(K)}, map_data_type(base_ring(K), K)}(true)

map_data(K::NfRelNS, L, ::Bool) = MapDataFromNfRelNS{Vector{elem_type(L)}, map_data_type(base_ring(K), L)}(true)

mutable struct NumFieldMor{S, T, U, V}
  header::MapHeader{S, T}
  image_data::U
  preimage_data::V

  function NumFieldMor{S, T, U, V}() where {S, T, U, V}
    z = new{S, T, U, V}()
    return z
  end

  NumFieldMor{S, T, U, V}(h::MapHeader{S, T}, i::U, p::V) where {S, T, U, V} = new{S, T, U, V}(f, i, p)
end

function _hom(K, L, x...; preimage = nothing)
  header = MapHeader(K, L)
  image_data = map_data(K, L, x...)
  if preimage !== nothing
    preimage_data = map_data(L, preimage...)
    return NumFieldMor(header, image_data, preimage_data)
  end
  z = NumFieldMor{typeof(K), typeof(L), typeof(image_data), map_data_type(L, K)}()
  z.header = header
  z.image_data = image_data
  return z
end

function image(f::NumFieldMor, x)
  return image(f.image_data, x)
end

function preimage(f::NumFieldMor, x)
  return image(f.preimage_data, x)
end

function (f::NumFieldMor)(x)
  return image(f, x)
end

function image_primitive_element(f::NumFieldMor{AnticNumberField})
  return f.image_data.prim_img
end

function image_primitive_element(f::NumFieldMor{<:NfRel})
  return f.image_data.prim_img
end


function TLEPropagatorBuilder(arg0::TLE, arg1::PositionAngle, arg2::jdouble)
    return TLEPropagatorBuilder((TLE, PositionAngle, jdouble), arg0, arg1, arg2)
end

function TLEPropagatorBuilder(arg0::TLE, arg1::PositionAngle, arg2::jdouble, arg3::DataContext)
    return TLEPropagatorBuilder((TLE, PositionAngle, jdouble, DataContext), arg0, arg1, arg2, arg3)
end

function build_propagator(obj::TLEPropagatorBuilder, arg0::Vector{jdouble})
    return jcall(obj, "buildPropagator", Propagator, (Vector{jdouble},), arg0)
end


## Preloads
using Statistics
using FFTW
using Plots
using BenchmarkTools
using Profile
using LinearAlgebra
using Measures

#internal AREPO units in cgs
arepoLength = 3.0856e20
arepoMass = 1.991e33
arepoVel = 1.0e5

#AREPO conversion factors
arepoTime = arepoLength/arepoVel
arepoDensity = arepoMass/arepoLength/arepoLength/arepoLength
arepoEnergy= arepoMass*arepoVel*arepoVel
arepoColumnDensity = arepoMass/arepoLength/arepoLength

#Helium abundance
ABHE = 0.1

#Mass of proton
mp=1.67e-24

#Grid size
nx,ny,nz=500,500,500

#Trial Read in
y =  Array{Float32}(undef, 125000000)

test = read("/Users/aksaydjari/Dropbox/GradSchool_AKS/Doug/ExtData/CloudFactory/density_grid_x1480y1390z1395",y)


# convenience macro
macro wrap_h(signature)
  Meta.isexpr(signature, :call) ||
    throw(ArgumentError("signature has to be a function call expression"))

  name = signature.args[1]
  args = signature.args[2:end]
  args_wo_h = [arg for arg in args if arg !== :h]

  quote
    if f.$name === nothing
      nothing
    else
      if isinplace(f)
        let _f = f.$name, h = h
          ($(args_wo_h...),) -> _f($(args...))
        end
      else
        let _f = f.$name, h = h
          ($(args_wo_h[2:end]...),) -> _f($(args[2:end]...))
        end
      end
    end
  end |> esc
end

struct ODEFunctionWrapper{iip,F,H,TMM,Ta,Tt,TJ,JP,SP,TW,TWt,TPJ,S,TCV} <: DiffEqBase.AbstractODEFunction{iip}
  f::F
  h::H
  mass_matrix::TMM
  analytic::Ta
  tgrad::Tt
  jac::TJ
  jac_prototype::JP
  sparsity::SP
  Wfact::TW
  Wfact_t::TWt
  paramjac::TPJ
  syms::S
  colorvec::TCV
end

function ODEFunctionWrapper(f::DiffEqBase.AbstractDDEFunction, h)
  # wrap functions
  jac = @wrap_h jac(J, u, h, p, t)
  Wfact = @wrap_h Wfact(W, u, h, p, dtgamma, t)
  Wfact_t = @wrap_h Wfact_t(W, u, h, p, dtgamma, t)

  ODEFunctionWrapper{isinplace(f),typeof(f.f),typeof(h),typeof(f.mass_matrix),
                     typeof(f.analytic),typeof(f.tgrad),typeof(jac),
                     typeof(f.jac_prototype),typeof(f.sparsity),
                     typeof(Wfact),typeof(Wfact_t),
                     typeof(f.paramjac),typeof(f.syms),typeof(f.colorvec)}(
                       f.f, h, f.mass_matrix, f.analytic, f.tgrad, jac,
                       f.jac_prototype, f.sparsity, Wfact, Wfact_t, f.paramjac,
                       f.syms, f.colorvec)
end

(f::ODEFunctionWrapper{true})(du, u, p, t) = f.f(du, u, f.h, p, t)
(f::ODEFunctionWrapper{false})(u, p, t) = f.f(u, f.h, p, t)