text:
"""
    morsefreq(c::CWT{W,T,Morse,N}) where {W,T,N}

Measures of frequency for generalized Morse wavelet. [with F. Rekibi]
The output returns the modal or peak.
For `be=0`, the wavelet becomes an analytic lowpass filter.
Reference: Lilly and Olhede (2009).  Higher-order properties of analytic wavelets.
IEEE Trans. Sig. Proc., 57 (1), 146--160.
"""
function morsefreq(c::CWT{W,T,Morse,N}) where {W,T,N}
    ga = c.waveType.ga
    be = c.waveType.be

    fm = @. exp((log(be) - log(ga)) / ga)

    if sum(be .== 0) != 0 && isempty(fm)
        fm = (log(2))^(1 / ga)
    elseif sum(be .== 0) != 0 && !isempty(fm)
        fm[be.==0] = (log(2))^(1 / ga[be.==0])
    end

    return fm
end

"""
    getNWavelets(n1,c) -> nOctaves, totalWavelets, sRanges, sWidth

Utility for understanding the spacing of the wavelets. `sRanges` is a list of
the s values used in each octave. `sWidth`` is a list of the corresponding
variance adjustments.
"""
function getNWavelets(n1, c::CWT)
    nOctaves = getNOctaves(n1, c)
    n, nSpace = setn(n1, c)
    isAve = !(typeof(c.averagingType) <: NoAve)
    if nOctaves ≤ c.averagingLength + getMinScaling(c) # there isn't enough space for anything but averaging.
        totalWavelets = isAve
        sRanges = [1.0]
        sWidth = [1.0]
        return nOctaves, totalWavelets, sRanges, sWidth
    end
    sRange = 2 .^ (polySpacing(nOctaves, c))
    totalWavelets = round(Int, length(sRange) + isAve)
    sWidth = varianceAdjust(c, totalWavelets, nOctaves)
    return nOctaves, totalWavelets, sRange, sWidth
end

"""
Different wavelet familes need to end at a different number of octaves because they have different tail behavior.
"""
getNOctaves(n1, c::CWT{W,T,M,N}) where {W,T,N,M} = log2(n1 >> 1 + 1) + c.extraOctaves
# choose the number of octaves so the last mean, which is at s*σ[1]
# is 3 standard devations away from the end
getNOctaves(n1, c::CWT{W,T,Morlet,N}) where {W,T,N} = log2((n1 >> 1 + 1) / (c.σ[1] + 3)) + c.extraOctaves
getNOctaves(n1, c::CWT{W,T,<:Paul,N}) where {W,T,N} = log2((n1 >> 1 + 1) / (2c.α + 5)) + c.extraOctaves
# choose the number of octaves so the last mean is 5 standard deviations from the end
function getNOctaves(n1, c::CWT{W,T,<:Dog,N}) where {W,T,N}
    μ = getMean(c)
    σ = getStd(c)
    log2(n1 >> 1 / (μ + 5σ)) + c.extraOctaves
end
# choose the number of octaves so the smallest support is twice the qmf
getNOctaves(n1, c::CWT{W,T,<:ContOrtho,N}) where {W,T,N} = log2(n1) - 2 - log2(length(qmf(c.waveType))) + c.extraOctaves
getNOctaves(n1, c::CWT{W,T,Morse,N}) where {W,T,N} = log2((n1 >> 1 + 1) / (morsefreq(c) + 1)) + c.extraOctaves
# getNOctaves(n1,c::CWT{W,T, Morse, N}) where {W, T, N} = 4 + c.extraOctaves

"""
As with the last octave, different wavelet families have different space decay rates, and in the case of symmetric or zero padding we don't want wavelets that bleed across the boundary from the center.
"""
getMinScaling(c::CWT{W,T,M,N}) where {W,T,N,M} = 0 # by default all scales are allowed (all of the orthogonal transforms)
getMinScaling(c::CWT{W,T,<:Morlet,N}) where {W,T,N} = 1 / (c.β)^0.8 # morlet is slightly too large at the boundaries by default
getMinScaling(c::CWT{W,T,<:Paul,N}) where {W,T,N} = 2 / (2c.α + 1) # Paul presents some difficulties, as the decay changes quickly (like 1/t^(α+1))
getMinScaling(c::CWT{W,T,<:Dog,N}) where {W,T,N} = 2 # like morlet, the decay for Dog is exponential and consistent across derivatives
# getMinScaling(c::CWT{W,T,<:Morse,N}) where {W,T,N,M} = log2(2 * morsefreq(c))
getMinScaling(c::CWT{W,T,<:Morse,N}) where {W,T,N,M} = 5 * morsefreq(c)

function varianceAdjust(this::CWT{W,T,M,N}, totalWavelets, nOct) where {W,T,N,M}
    # increases the width of the wavelets by σ[i] = (1+a(total-i)ᴾ)σₛ
    # which is a polynomial of order p going from 1 at the highest frequency to
    # sqrt(p) at the lowest
    β = this.β
    x = exp(nOct) / (100 + exp(nOct)) # the fewer octaves there are, the smaller the width adjustment we need
    a = (β^x - 1) / (totalWavelets - 1)^β
    sWidth = 1 .+ a .* (totalWavelets .- (1:totalWavelets)) .^ β
    if any(isnan.(sWidth))
        return [1]
    end
    return sWidth
end

function polySpacing(nOct, c)
    a = max(getMinScaling(c) + c.averagingLength, 0)
    O = nOct
    β = c.β
    Q = c.Q
    if O ≤ a
        # averagingLength and the min scaling are too large for anything to be done
        return [1.0]
    elseif 0 < O - a ≤ 1
        # there's only one octave, just return the linear rate of Q
        return range(a, O, length = 1 + round(Int, Q))
    end
    # x is the index, y is the scale
    # y= b*x^(1/β), solve for a and b with
    # (x₀,y₀)=(1, aveLength + minScale)
    # dy/dx = 1/s, so the Quality factor gives the slope at the last frequency
    b = (β / Q)^(1 / β) * (O)^((β - 1) / β)
    # the point x so that the second condition holds
    lastWavelet = Q * (O) / β
    # the point so that the first wavelet is at a
    firstWavelet = (a / b)^β
    # step size so that there are actually Q wavelets in the last octave
    startOfLastOctave = ((nOct - 1) / b)^β
    stepSize = (lastWavelet - startOfLastOctave) / Q
    samplePoints = range(firstWavelet, stop = lastWavelet,
        step = stepSize)#, length = round(Int, O * Q^(1/p)))
    return b .* (samplePoints) .^ (1 / β)
end

# a utility to just get the start, stop, and step size used in polySpacing. Only used for explanatory purposes
function genSamplePoints(nOct, c)
    a = getMinScaling(c) + c.averagingLength
    O = nOct
    β = c.β
    Q = c.Q
    if O ≤ a
        # averagingLength and the min scaling are too large for anything to be done
        return [1.0]
    elseif 0 < O - a ≤ 1
        # there's only one octave, just return the linear rate of Q
        return range(a, O, length = 1 + round(Int, Q))
    end
    # x is the index, y is the scale
    # y= b*x^(1/β), solve for a and b with
    # (x₀,y₀)=(1, aveLength + minScale)
    # dy/dx = 1/s, so the Quality factor gives the slope at the last frequency
    b = (β / Q)^(1 / β) * (O)^((β - 1) / β)
    # the point x so that the second condition holds
    lastWavelet = Q * (O) / β
    # the point so that the first wavelet is at a
    firstWavelet = (a / b)^β
    # step size so that there are actually Q wavelets in the last octave
    startOfLastOctave = ((nOct - 1) / b)^β
    stepSize = (lastWavelet - startOfLastOctave) / Q
    firstWavelet, lastWavelet, stepSize
end

"Adjust the length of the storage based on the boundary conditions."
function setn(n1, c)
    if boundaryType(c) <: ZPBoundary
        base2 = ceil(Int, log2(n1))   # power of 2 nearest to n1
        nSpace = 2^(base2)
        n = nSpace >> 1 + 1
    elseif boundaryType(c) <: SymBoundary
        # n1+1 rather than just n1 because this is going to be used in an rfft
        # for real data
        n = n1 + 1
        nSpace = 2 * n1
    elseif boundaryType(c) <: PerBoundary
        n = n1 >> 1 + 1
        nSpace = n1
    else
        error("No such boundary as $(boundaryType(c))")
    end
    return n, nSpace
end

# computing averaging function utils

# this makes sure the averaging function is real and positive
adjust(c::CWT) = 1
adjust(c::CWT{W,T,<:Dog}) where {W,T} = 1 / (im)^(c.α)

function locationShift(c::CWT{W,T,<:Morlet,N}, s, ω, sWidth) where {W,T,N}
    s0 = 3 * c.σ[1] * s * sWidth
    ω_shift = ω .+ c.σ[1] * s0
    return (s0, ω_shift)
end

function locationShift(c::CWT{W,T,<:Dog,N}, s, ω, sWidth) where {W,T,N}
    μNoS = getMean(c)
    μLast = μNoS * (2s)
    s0 = μLast / 2 / getStd(c)
    μ = μNoS * s0
    ω_shift = ω .+ μ
    return (s0, ω_shift)
end

function locationShift(c::CWT{W,T,<:Morse,N}, s, ω, sWidth) where {W,T,N}
    # s0 = c.waveType.cf * s * sWidth
    s0 = morsefreq(c) * s * sWidth
    # ω_shift = ω .+ c.waveType.cf * s0
    ω_shift = ω .+ s0
    return (s0, ω_shift)
end

"Get the mean of the mother wavelet where s=1."
function getMean(c::CWT{W,T,<:Dog}, s = 1) where {W,T}
    Gm1 = gamma((c.α + 1) / 2)
    Gm2 = gamma((c.α + 2) / 2)
    Gm2 / Gm1 * sqrt(2) * s^2
end
getMean(c::CWT{W,T,<:Paul}, s = 1) where {W,T} = (c.α + 1) * s
function getMean(c::CWT{W,T,<:Morlet}, s = 1) where {W,T}
    return s * c.σ[1]
end
function getMean(c::CWT{W,T,<:Morse}, s = 1) where {W,T}
    #return s*c.waveType.cf
    return s * morsefreq(c)
end

"""
    getStd(c::CWT{W, T, <:Dog}, s=1) where {W,T}

Get the standard deviation of the mother wavelet.
"""
getStd(c::CWT{W,T,<:Dog}, s = 1) where {W,T} = sqrt(c.α + 1 - getMean(c)^2) * s^(3 / 2)
getStd(c::CWT{W,T,<:Paul}, s = 1) where {W,T} = sqrt((c.α + 2) * (c.α + 1)) * s
getStd(c::CWT{W,T,<:Morlet}, s = 1) where {W,T} = (s * c.β^0.8) / sqrt(2)
getStd(c::CWT{W,T,<:Morse}, s = 1) where {W,T} = (s * c.β^0.8) / sqrt(2)

function locationShift(c::CWT{W,T,<:Paul,N}, s, ω, sWidth) where {W,T,N}
    s0 = s * sqrt(c.α + 1)
    ω_shift = ω .+ (c.α .+ 1) * s0
    return (s0, ω_shift)
end

function getUpperBound(c::CWT{W,T,<:Morlet,N}, s) where {W,T,N}
    return c.σ[1] * s
end

function getUpperBound(c::CWT{W,T,<:Dog,N}, s) where {W,T,N}
    return sqrt(2) * s * gamma((c.α + 2) / 2) / gamma((c.α + 1) / 2)
end

function getUpperBound(c::CWT{W,T,<:Paul,N}, s) where {W,T,N}
    return (c.α + 1) * s
end

"""
    getMeanFreq(Ŵ, fsample=2000) -> arrayOfFreqs

Calculate each of the mean frequencies of a collection of analytic or real wavelets `Ŵ`.
The default sampling rate `fsample=2kHz`, so the maximum frequency is 1kHz.
"""
function getMeanFreq(Ŵ, fsample = 2000)
    eachNorm = [norm(w, 1) for w in eachslice(Ŵ, dims = ndims(Ŵ))]'
    freqs = range(0, fsample / 2, length = size(Ŵ, 1))
    dropdims(sum(freqs .* abs.(Ŵ), dims = 1) ./ eachNorm, dims = 1)
end

function getMeanFreq(n1, cw::CWT, fsample = 2000)
    Ŵ, ω = computeWavelets(n1, cw)
    eachNorm = [norm(w, 1) for w in eachslice(Ŵ, dims = ndims(Ŵ))]'
    freqs = range(0, fsample / 2, length = size(Ŵ, 1))
    dropdims(sum(freqs .* abs.(Ŵ), dims = 1) ./ eachNorm, dims = 1)
end

# orthogonal wavelet utils
function getContWaveFromOrtho(c, N)
    depth, ψLen = calcDepth(c, N)
    φ, ψ, t = makewavelet(c.waveType.o, depth)
    return φ, -ψ, ψLen
end

"""
    calcDepth(w,N) -> nIters, sigLength

Given a `CWT` with an orthogonal wavelet type, calculate the depth `nIters`
necessary to get a resulting wavelet longer than `N`. `sigLength` is the
resulting length.
"""
function calcDepth(w, N)
    origLen = length(qmf(w.waveType))
    netLen = origLen
    nIters = 0
    if N < netLen
        return nIters
    end
    while N > netLen - 2^nIters + 1
        netLen = nextLength(netLen, origLen)
        nIters += 1
    end
    return (nIters, netLen - 2^nIters + 1)
end
# the actual length calculator, as this is recursive
nextLength(curr, orig) = curr * 2 + orig - 1 # upsample then convolve
# pad a vector `v` to have length `N`. Or chop off from both sides if its too long
function padTo(v, N)
    nDiff = N - length(v)
    if nDiff > 0 # there should be more entries than there are
        cat(v, zeros(eltype(v), max(N - length(v), 0)), dims = 1)
    elseif nDiff == 0
        return v
    else # there are too many, so chop off half the difference from each side
        inds = ceil(Int, 1 - nDiff / 2):length(v)+ceil(Int, nDiff / 2)
        return v[inds]
    end
end

# create interpolater for the orthogonal cases
genInterp(ψ) = interpolate(ψ, BSpline(Quadratic(Reflect(OnGrid()))))

"""
    computeDualWeights(Ŵ) -> β

Compute the weight given to each wavelet so that in the Fourier domain, the sum across
wavelets is as close to 1 at every frequency below the peak of the last wavelet
(that is `β' .*abs.(Ŵ[1:lastFreq,:]) ≈ ones(lastFreq)` in an ℓ^2 sense).
"""
function computeDualWeights(Ŵ, wav)
    @views lastReasonableFreq = computeLastFreq(Ŵ[:, end], wav)
    Wdag = pinv(Ŵ[1:lastReasonableFreq, :])
    β = conj.(Wdag * ones(size(Wdag, 2)))'
    return β
end
# function computeDualWeights(Ŵ, wav::CWT{W,T,<:ContOrtho,N}, n,naive=true) where {W,T,N}
#     β = computeNaiveDualWeights(Ŵ, wav, n)
#     return β
# end
computeLastFreq(ŵ, wav) = length(ŵ) - 3 # we have constructed the wavelets so the last frequency is <1%, so trying to get it to 1 will fail
computeLastFreq(ŵ, wav::CWT{W,T,<:Morlet}) where {W,T} = findlast(abs.(ŵ) .> 1 / 4 * maximum(abs.(ŵ))) # the gaussian decay makes this way too large by the end

"""
    computeNaiveDualWeights(Ŵ, wav, n1)

Compute the dual weights using the scaling amounts and the normalization power p. Primarily
used when the least squares version is poorly constructed. When the least squares version
doesn't perform well, this is also likely to have poor reconstruction, but it won't give
extremely negative or oscillatory weights like the least squares version.
"""
function computeNaiveDualWeights(Ŵ, wav, n1)
    _, _, sRange, _ = ContinuousWavelets.getNWavelets(n1, wav)
    isAve = !(wav.averagingType isa NoAve)
    p = wav.p
    # every wavelet is multiplied by 1/s^(1/p), so we need to multiply by that to normalize in frequency. The extra +1 is to convert from uniform in scale to uniform in frequency
    sToTheP = sRange' .^ (1 / p + 1)
    if isAve && size(Ŵ, 2) > 1
        # averaging needs to have maximum value of 1
        aveMultiplier = 1 / norm(Ŵ[:, 1], Inf)
        dual = sum(Ŵ[:, 2:end] .* sToTheP, dims = 2)
        renormalize = norm(dual, Inf)
        dual ./= renormalize
        #dual += Ŵ[:,1] .* sToTheP[1]
        β = cat(aveMultiplier, 2 .* sToTheP ./ renormalize..., dims = 2)
    elseif size(Ŵ, 2) > 1
        dual = sum(Ŵ .* sToTheP, dims = 2)
        β = sToTheP ./ norm(dual, Inf)
        β[2:end] .* 2
    else
        # there's only the averaging
        β = [1 / norm(Ŵ[:, 1], Inf)]
    end
    return β
end

"""
    getDualCoverage(n,cWav, invType) -> dualCover, dualNorm

Get the sum of the weights and its deviation from 1.
"""
function getDualCoverage(n, cWav, invType)
    Ŵ = computeWavelets(n, cWav)[1]
    if invType isa DualFrames
        dualCover = sum(conj.(Ŵ) .* explicitConstruction(Ŵ), dims = 2)
        return dualCover, norm(dualCover .- 1)
    elseif invType isa PenroseDelta
        β = computeDualWeights(Ŵ, cWav)
    elseif invType isa NaiveDelta
        β = computeNaiveDualWeights(Ŵ, cWav, n)
    end
    dualCover = sum(conj.(β) .* Ŵ, dims = 2)
    dualCover, norm(dualCover .- 1)
end

function dualDeviance(n, cWav, naive)
    getDualCoverage(n, cWav, naive)
end

"""
    explicitConstruction(Ŵ)

Explicitly construct the canonical dual frame elements for a translation invariant frame.
This is likely to end poorly due to the low representation at high frequencies.
"""
function explicitConstruction(Ŵ)
    Ŵdual = conj.(Ŵ ./ [norm(ŵ)^2 for ŵ in eachslice(Ŵ, dims = 1)])
end

"""
    caveats(n1, c::CWT{B,CT,W}; coiTolerance = exp(-2), fsample = 2000) -> sRange, meanFreqs, coi

Given a length `n1` and a CWT struct `c`, returns the scales `sRange` used, the mean frequencies of the wavelets `meanFreqs`, and the cone of influence `coi` for each wavelet.
Returns the period, the scales, and the cone of influence for the given wavelet transform.
If you have sampling information, you will need to scale the vector scale appropriately by
1/δt, and the actual transform by δt^(1/p).
"""
function caveats(n1, c::CWT; coiTolerance = exp(-2), fsample = 2000)
    nOctaves, totalWavelets, sRange, sWidth = getNWavelets(n1, c)
    # padding determines the actual number of elements
    n, nSpace = setn(n1, c)
    # indicates whether we should keep a spot for the father wavelet
    isAve = !(typeof(c.averagingType) <: NoAve)

    # Fourier equivalent frequencies
    Ŵ, ω = computeWavelets(n1, c)
    freqs = getMeanFreq(Ŵ, fsample)
    coi = directCoiComputation(n1, c; coiTolerance = coiTolerance)
    return sRange, freqs, coi
end

caveats(Y::AbstractArray{T}, w::ContWaveClass) where {T<:Number,S<:Real} = caveats(size(Y, 1), CWT(w), J1 = J1)

"""
    directCoiComputation(n1, c::CWT; coiTolerance = exp(-2)) -> coi

A straightforward computation of the cone of influence directly from the
constructed wavelets. `coi` is a binary matrix of dimensions `(signalLength)×(nscales+1)`
that indicates when the autocorrelation of the wavlet is greater than
`coiTolerance` of the maximum value.
"""
function directCoiComputation(n1, c::CWT; coiTolerance = exp(-2))
    ψ, ω = ContinuousWavelets.computeWavelets(n1, c, space = true)
    autoCorr = cat([sum(ψ .* conj.(circshift(ψ, (ii, 0))), dims = 1) for ii = 1:size(ψ, 1)]..., dims = 1) # auto correlation in time domain
    normedAutoCorr = autoCorr ./ maximum(abs.(autoCorr), dims = 1) # normalize by the max norm
    return abs.(normedAutoCorr) .>= coiTolerance # the cone is where the autocorrelation is larger than the tolerance
end

"""
    coi(n, s, wave::CWT)

for a signal of length `n` and wavelet transform of type wave, return a binary vector indicating the cone of influence from the edges so that the auto-correlation of the wavelet has decreased by a factor of ``\\textrm{e}^2``.
"""
function coi(n, s, ::Morlet)
    sqrt(2) * s
end

@doc raw"""
    crossSpectrum(X, Y, c::CWT)

Compute the cross spectrum of signals `X` and `Y`, which is defined as ``S(\\conj{C(X)})``
Note that unlike `cwt`, `crossSpectrum` only works for collections `X` and `Y` of vectors of shape (signalLength)×(nSignals), and outputs a collection of cross spectrums that has shape (signalLength)×(nscales+1)×(nSignalsX)×(nSignalsY).

# Examples
```jldoctest
julia> crossSpectrum(X, Y)
insert result of crossSpectrum(X, Y)
```
"""
function crossSpectrum(X, Y, c::CWT{B,S,W,N,isAn}) where {B,S,W,N,isAn}
    aveCXY, _ = sharedCrossSpectrum(X, Y, c)
    return aveCXY
end

@doc raw"""
    waveletCoherence(X, Y, c)

Compute the wavelet coherence between `X` and `Y`. This is given by the power of the cross spectrum, normalized by smoothed powers of both `X` and `Y`. Explicitly, that is ``\frac{|S(C^*(X)C(Y))|^2}{S(|C(X)|^2)S(|C(Y)|^2)}``

# Examples
```jldoctest
julia> waveletCoherence(X,Y,c)
insert result of waveletCoherence(X,Y,c)
```
"""
function waveletCoherence(X, Y, c)
    (aveCXY, aveWavelet, CX, CY) = sharedCrossSpectrum(X, Y, c)
    aveCX = abs.(dropdims(ContinuousWavelets.cwt(abs.(CX) .^ 2, c, aveWavelet), dims = 2))
    aveCY = abs.(dropdims(ContinuousWavelets.cwt(abs.(CY) .^ 2, c, aveWavelet), dims = 2))
    return abs.(aveCXY) .^ 2 ./ (aveCX .* aveCY)
end

"""
`waveletCoherence` depends on results precomputed by `crossSpectrum` in addition to the actual output of `crossSpectrum`.
"""
function sharedCrossSpectrum(X, Y, c)
    # ensure that the wavelet transform uses some sort of averaging
    isAve = !(typeof(c.averagingType) <: NoAve)
    if !isAve
        cAve = CWT(c.waveType, c.Q, ContinuousWavelets.boundaryType(c)(), ContinuousWavelets.Father(), c.averagingLength, c.frameBound, c.p, c.β)
    else
        cAve = c
    end

    # make sure both X and Y have the right shape
    if ndims(X) < 2
        X = reshape(X, (size(X)..., 1))
    elseif ndims(X) > 2
        X = reshape(X, (size(X, 1), :))
    end
    if ndims(Y) < 2
        Y = reshape(Y, (size(Y)..., 1))
    elseif ndims(Y) > 2
        Y = reshape(Y, (size(Y, 1), :))
    end
    daughters, _ = ContinuousWavelets.computeWavelets(size(Y, 1), cAve)
    aveWavelet = daughters[:, 1:1] ./ norm(daughters[:, 1:1], c.p) # we need an averaging function that returns the ave value, rather than the one that preserves the frame bound
    CX = ContinuousWavelets.cwt(X, cAve, daughters)
    CY = ContinuousWavelets.cwt(Y, cAve, daughters)
    # reshape so we can do an outer product
    CX = reshape(CX, (size(CX)[1:3]..., 1))
    CY = reshape(CY, (size(CY)[1:2]..., 1, size(CY, 3)))
    CXY = conj.(CX) .* CY # conjugate of the first times the second
    aveCXY = dropdims(ContinuousWavelets.cwt(CXY, cAve, aveWavelet), dims = 2) # get the average of the multiplication
    return (aveCXY, aveWavelet, CX, CY)
end


text:
# Initial ODR for DataType
const DataTypeODR = OnDiskRepresentation{(0, odr_sizeof(Vlen{String})),Tuple{String,Vector{Any}},Tuple{Vlen{String},Vlen{RelOffset}}}

const NULL_COMMITTED_DATATYPE = CommittedDatatype(RelOffset(0), 0)

function track_weakref!(f::JLDFile, header_offset::RelOffset, @nospecialize v)
    header_offset !== NULL_REFERENCE && (f.jloffset[header_offset] = WeakRef(v))
    nothing
end

## Generic machinery

# Carries the type and on-disk representation of data to be read from
# the disk
odr_sizeof(::ReadRepresentation{T,S}) where {T,S} = odr_sizeof(S)

# Determines whether a specific field type should be saved in the file
@noinline function hasfielddata(@nospecialize(T), encounteredtypes=DataType[])
    T === Union{} && return false
    !isconcretetype(T) && return true
    T = T::DataType
    T in encounteredtypes && return true
    push!(encounteredtypes, T)
    (ismutabletype(T) || T <: Type) && return true
    hasdata(T, encounteredtypes)
end

# Determines whether a specific type has fields that should be saved in the file
function hasdata(T::DataType, encounteredtypes=DataType[])
    isempty(T.types) && T.size != 0 && return true
    for ty in T.types
        hasfielddata(writeas(ty), copy(encounteredtypes)) && return true
    end
    false
end

# Gets the size of an on-disk representation
Base.@pure function odr_sizeof(::OnDiskRepresentation{Offsets,JLTypes,H5Types}) where {Offsets,JLTypes,H5Types}
    Offsets[end]+odr_sizeof(H5Types.parameters[end])
end

# Determines whether a type will have the same layout on disk as in memory
function samelayout(T::DataType)
    isempty(T.types) && return true
    offset = 0
    for i = 1:length(T.types)
        offset != fieldoffset(T, i) && return false
        ty = T.types[i]
        ty !== writeas(ty) && return false
        !samelayout(ty) && return false
        offset += ty.size
    end
    return offset == T.size
end
samelayout(::Type) = false

fieldnames(x::Type{T}) where {T<:Tuple} = [Symbol(i) for i = 1:length(x.types)]
fieldnames(@nospecialize x) = collect(Base.fieldnames(x))

# fieldodr gives the on-disk representation of a field of a given type,
# which is either always initialized (initialized=true) or potentially
# uninitialized (initialized=false)
@generated function fieldodr(::Type{T}, initialized::Bool) where T
    if isconcretetype(T)
        if !hasfielddata(T)
            # A ghost type, so no need to store at all
            return nothing
        elseif isa(T, DataType)
            if isbitstype(T)
                return :(odr(T))
            elseif !ismutabletype(T)
                return :(initialized ? odr(T) : RelOffset)
            end
        end
    end
    RelOffset
end

# h5fieldtype is fieldodr's HDF5 companion. It should give the HDF5
# datatype reflecting the on-disk representation.
@generated function h5fieldtype(f::JLDFile, writeas::Type{T}, readas::Type,
                                initialized::Initialized) where T
    if isconcretetype(T)
        if !hasfielddata(T)
            return nothing
        elseif isbitstype(T) || (isa(initialized, Type{Type{Val{true}}}) && !ismutabletype(T))
            return quote
                @lookup_committed f T
                $(if isempty(T.types)
                    # Opaque datatype
                    :(return commit(f, OpaqueDatatype(T.size), T, readas))
                else
                    # Compound type
                    :(return commit_compound(f, fieldnames(T), T, readas))
                end)
            end
        end
    end
    ReferenceDatatype()
end

# objodr gives the on-disk representation of a given object. This is
# almost always the on-disk representation of the type. The only
# exception is strings, where the length is encoded in the datatype in
# HDF5, but in the object in Julia.
@inline function objodr(x)
    writtenas = writeas(typeof(x))
    _odr(writtenas, typeof(x), odr(writtenas))
end
_odr(writtenas::DataType, readas::DataType, odr) =
    CustomSerialization(writtenas, readas, odr)

# h5type is objodr's HDF5 companion. It should give the HDF5 datatype
# reflecting the on-disk representation
#
# Performance note: this should be inferrable.
function h5type(f::JLDFile, writtenas, x)
    check_writtenas_type(writtenas)
    T = typeof(x)
    @lookup_committed f T
    if !hasdata(writtenas)
        commit(f, OpaqueDatatype(1), writtenas, T, WrittenAttribute(f, :empty, UInt8(1)))
    elseif isempty(writtenas.types) # bitstype
        commit(f, OpaqueDatatype(writtenas.size), writtenas, T)
    else
        commit_compound(f, fieldnames(writtenas), writtenas, T)
    end
end
check_writtenas_type(::DataType) = nothing
check_writtenas_type(::Any) = throw(ArgumentError("writeas(leaftype) must return a leaf type"))
h5type(f::JLDFile, x) = h5type(f, writeas(typeof(x)), x)

# Make a compound datatype from a set of names and types
function commit_compound(f::JLDFile, names::AbstractVector{Symbol},
                         writtenas::DataType, readas::Type)
    types = writtenas.types
    offsets = Int[]
    h5names = Symbol[]
    members = H5Datatype[]
    fieldtypes = RelOffset[]
    hasfieldtype = false

    offset = 0
    for i = 1:length(types)
        fieldty = types[i]
        fieldwrittenas = writeas(fieldty)
        !hasfielddata(fieldwrittenas) && continue
        dtype = h5fieldtype(f, fieldwrittenas, fieldty, Val{i <= ninitialized(writtenas)})
        dtype === nothing && continue
        push!(h5names, names[i])
        if isa(dtype, CommittedDatatype)
            # HDF5 cannot store relationships among committed
            # datatypes. We store these separately in an attribute.
            push!(fieldtypes, dtype.header_offset)
            dtype = f.datatypes[dtype.index]
            hasfieldtype = true
        else
            push!(fieldtypes, NULL_REFERENCE)
        end
        push!(members, dtype)
        push!(offsets, offset)
        offset += dtype.size::UInt32
    end

    @assert offset != 0
    compound = CompoundDatatype(offset, h5names, offsets, members)
    if hasfieldtype
        fieldtypeattr = WrittenAttribute(:field_datatypes,
                                         WriteDataspace(f, fieldtypes, DataType),
                                         ReferenceDatatype(),
                                         fieldtypes)
        commit(f, compound, writtenas, readas, fieldtypeattr)::CommittedDatatype
    else
        commit(f, compound, writtenas, readas)::CommittedDatatype
    end
end

# Write an HDF5 datatype to the file
function commit(f::JLDFile, dtype::H5Datatype, writeas::DataType, readas::DataType,
                attributes::WrittenAttribute...)
    io = f.io

    # This needs to be written this way or type inference gets unhappy...
    # Also needs to happen here so that we write the DataType type
    # before we try to find where this type will be written
    typeattr = WrittenAttribute(
        :julia_type, WriteDataspace(f, DataType, odr(DataType)), h5type(f, DataType, DataType), readas)

    offset = f.end_of_data

    seek(io, offset)
    id = length(f.datatypes)+1
    h5o = h5offset(f, offset)
    cdt = CommittedDatatype(h5o, id)
    f.datatype_locations[h5o] = cdt
    f.jlh5type[readas] = cdt
    push!(f.datatypes, dtype)
    f.types_group[@sprintf("%08d", id)] = h5o

    if writeas !== readas
        wrtypeattr = WrittenAttribute(:written_type,
                                      WriteDataspace(f, DataType, odr(DataType)),
                                      h5type(f, DataType, DataType), writeas)
        f.h5jltype[cdt] = ReadRepresentation{readas,CustomSerialization{writeas, odr(writeas)}}()
        commit(f, dtype, tuple(typeattr, wrtypeattr, attributes...))
    else
        f.h5jltype[cdt] = ReadRepresentation{writeas,odr(writeas)}()
        commit(f, dtype, tuple(typeattr, attributes...))
    end

    cdt::CommittedDatatype
end

# h5convert! stores the HDF5 representation of Julia data to a pointer. This
# method handles types with no padding or references where this is just a simple
# store
h5convert!(out::Pointers, ::Type{T}, ::JLDFile, x, ::JLDWriteSession) where {T} =
    (jlunsafe_store!(convert(Ptr{T}, out), x); nothing)

# We pack types that have padding using a staged h5convert! method
@generated function h5convert!(out::Pointers,
       ::OnDiskRepresentation{Offsets,Types,H5Types},
       file::JLDFile, x, wsession::JLDWriteSession) where {Offsets,Types,H5Types}
    T = x
    types = Types.parameters
    members = H5Types.parameters

    getindex_fn = isa(T, Type{T} where T<:Tuple) ? (:getindex) : (:getfield)
    ex = Expr(:block)
    args = ex.args
    for i = 1:length(Offsets)
        member = members[i]
        isa(member, Nothing) && continue

        offset = Offsets[i]
        conv = :(h5convert!(out+$offset, $(member), file, convert($(types[i]), $getindex_fn(x, $i)), wsession))
        if i > ninitialized(T) && (!isconcretetype(x.types[i]) || !isbitstype(x.types[i]))
            push!(args, quote
                if !isdefined(x, $i)
                    h5convert_uninitialized!(out+$offset, $(member))
                else
                    $conv
                end
            end)
        else
            push!(args, conv)
        end
    end
    push!(args, nothing)
    ex
end

# jl_canbeuninitialized specifies whether a given ODR could be uninitialized,
# which determines whether or not we'll try to call jlconvert_isinitialized to
# determine whether or not it is actually defined when reconstructing types.
jlconvert_canbeuninitialized(::Any) = false

# jlconvert converts data from a pointer into a Julia object. This method
# handles types where this is just a simple load
@inline jlconvert(::ReadRepresentation{T,T}, ::JLDFile, ptr::Ptr,
                  ::RelOffset) where {T} =
    jlunsafe_load(convert(Ptr{T}, ptr))

# When fields are undefined in the file but can't be in the workspace, we need
# to throw exceptions to prevent errors on null pointer loads
struct UndefinedFieldException
    ty::DataType
    fieldname::Symbol
end
Base.showerror(io::IO, x::UndefinedFieldException) =
    print(io, "field \"", x.fieldname, "\" of type ", x.ty,
          " must be defined in the current workspace, but was undefined in the file")

## References

h5type(::JLDFile, ::Type{RelOffset}, ::RelOffset) = ReferenceDatatype()
odr(::Type{RelOffset}) = RelOffset

@inline function h5convert!(out::Pointers, odr::Type{RelOffset}, f::JLDFile, x::Any,
                            wsession::JLDWriteSession)
    ref = write_ref(f, x, wsession)
    jlunsafe_store!(convert(Ptr{RelOffset}, out), ref)
    nothing
end
h5convert_uninitialized!(out::Pointers, odr::Type{RelOffset}) =
    (jlunsafe_store!(convert(Ptr{RelOffset}, out), NULL_REFERENCE); nothing)

# Reading references as references
jlconvert(::ReadRepresentation{RelOffset,RelOffset}, f::JLDFile, ptr::Ptr,
          ::RelOffset) =
    jlunsafe_load(convert(Ptr{RelOffset}, ptr))
jlconvert_canbeuninitialized(::ReadRepresentation{RelOffset,RelOffset}) = false

# Reading references as other types
@inline function jlconvert(::ReadRepresentation{T,RelOffset}, f::JLDFile, ptr::Ptr,
                           ::RelOffset) where T
    x = load_dataset(f, jlunsafe_load(convert(Ptr{RelOffset}, ptr)))
    (isa(x, T) ? x : convert(T, x))::T
end
jlconvert_canbeuninitialized(::ReadRepresentation{T,RelOffset}) where {T} = true
jlconvert_isinitialized(::ReadRepresentation{T,RelOffset}, ptr::Ptr) where {T} =
    jlunsafe_load(convert(Ptr{RelOffset}, ptr)) != NULL_REFERENCE

## Routines for variable-length datatypes

# Write variable-length data and store the offset and length to out pointer
@inline function store_vlen!(out::Pointers, odr, f::JLDFile, x::AbstractVector,
                             wsession::JLDWriteSession)
    jlunsafe_store!(convert(Ptr{UInt32}, out), length(x))
    obj = write_heap_object(f, odr, x, wsession)
    jlunsafe_store!(convert(Ptr{GlobalHeapID}, out)+4, obj)
    nothing
end

h5convert!(out::Pointers, ::Type{Vlen{T}}, f::JLDFile, x, wsession::JLDWriteSession) where {T} =
    store_vlen!(out, T, f, x, wsession)

@assert odr_sizeof(Vlen) == jlsizeof(UInt128)
h5convert_uninitialized!(out::Pointers, odr::Type{T}) where {T<:Vlen} =
    (jlunsafe_store!(convert(Ptr{Int128}, out), 0); nothing)

# Read variable-length data given offset and length in ptr
jlconvert(::ReadRepresentation{T,Vlen{S}}, f::JLDFile, ptr::Ptr, ::RelOffset) where {T,S} =
    read_heap_object(f, jlunsafe_load(convert(Ptr{GlobalHeapID}, ptr+4)), ReadRepresentation{T, S}())
jlconvert_canbeuninitialized(::ReadRepresentation{T,Vlen{S}}) where {T,S} = true
jlconvert_isinitialized(::ReadRepresentation{T,Vlen{S}}, ptr::Ptr) where {T,S} =
    jlunsafe_load(convert(Ptr{GlobalHeapID}, ptr+4)) != GlobalHeapID(RelOffset(0), 0)

## DataTypes

const H5TYPE_DATATYPE = CompoundDatatype(
    odr_sizeof(Vlen{String})+odr_sizeof(Vlen{RelOffset}),
    [:name, :parameters],
    [0, odr_sizeof(Vlen{String})],
    [H5TYPE_VLEN_UTF8, VariableLengthDatatype(ReferenceDatatype())]
)

function h5fieldtype(f::JLDFile, ::Type{T}, readas::Type, ::Initialized) where T<:DataType
    if !(readas <: DataType)
        @lookup_committed f readas
        return commit(f, H5TYPE_DATATYPE, DataType, readas)
    end

    @lookup_committed f DataType
    io = f.io
    offset = f.end_of_data

    seek(io, offset)
    id = length(f.datatypes)+1
    h5o = h5offset(f, offset)
    cdt = CommittedDatatype(h5o, id)
    f.datatype_locations[h5o] = cdt
    f.jlh5type[DataType] = cdt
    f.h5jltype[cdt] = ReadRepresentation{DataType,DataTypeODR()}()
    push!(f.datatypes, H5TYPE_DATATYPE)
    f.types_group[@sprintf("%08d", id)] = h5o

    commit(f, H5TYPE_DATATYPE, (WrittenAttribute(:julia_type, WriteDataspace(f, DataType, odr(DataType)), cdt, DataType),))

    cdt
end
fieldodr(::Type{T}, ::Bool) where {T<:DataType} = DataTypeODR()

h5type(f::JLDFile, ::Type{T}, x) where {T<:DataType} =
    h5fieldtype(f, DataType, typeof(x), Val{true})
odr(::Type{T}) where {T<:DataType} = DataTypeODR()

function typename(T::DataType)
    s = IOBuffer()
    join(s, fullname(T.name.module), '.')
    print(s, '.')
    print(s, T.name.name)
    return String(resize!(s.data, s.size))
end

function refs_from_types(f::JLDFile, types, wsession::JLDWriteSession)
    refs = RelOffset[
        if isa(x, DataType)
            # The heuristic here is that, if the field type is a committed data type,
            # then we commit the datatype and write it as a reference to the committed
            # datatype. Otherwise we write it as a name. This ensures that type
            # parameters that affect the structure of a type are written to the file,
            # so that we can reconstruct the type when the layout depends on the
            # parameters.
            dt = h5fieldtype(f, writeas(x), x, Val{true})
            if isa(dt, CommittedDatatype)
                (dt::CommittedDatatype).header_offset
            else
                write_ref(f, x, wsession)
            end
        else
            write_ref(f, x, wsession)
        end
    for x in types]
end

function h5convert!(out::Pointers, ::DataTypeODR, f::JLDFile, T::DataType, wsession::JLDWriteSession)
    t = typename(T)
    store_vlen!(out, UInt8, f, unsafe_wrap(Vector{UInt8}, t), f.datatype_wsession)
    if isempty(T.parameters)
        h5convert_uninitialized!(out+odr_sizeof(Vlen{UInt8}), Vlen{UInt8})
    else
        refs = refs_from_types(f, T.parameters, wsession)
        store_vlen!(out+odr_sizeof(Vlen{UInt8}), RelOffset, f, refs, f.datatype_wsession)
    end
    nothing
end

# This is a trick to compactly write long NTuple
# This uses that NTuple{N,T} == Tuple{T,T,T,T,...,T}
function h5convert!(out::Pointers, ::DataTypeODR, f::JLDFile, T::Type{NTuple{N,ET}}, wsession::JLDWriteSession) where {N,ET}
    if isempty(T.parameters)
        store_vlen!(out, UInt8, f, unsafe_wrap(Vector{UInt8}, "Tuple"), f.datatype_wsession)
        h5convert_uninitialized!(out+odr_sizeof(Vlen{UInt8}), Vlen{UInt8})
    else
        store_vlen!(out, UInt8, f, unsafe_wrap(Vector{UInt8}, "NTuple"), f.datatype_wsession)
        refs = refs_from_types(f, Any[N,ET], wsession)
        store_vlen!(out+odr_sizeof(Vlen{UInt8}), RelOffset, f, refs, f.datatype_wsession)
    end
    nothing
end

## Union Types

const H5TYPE_UNION = CompoundDatatype(
      odr_sizeof(Vlen{String})+2*odr_sizeof(Vlen{RelOffset}),
      [:description, :datatype, :unionall],
      [0, odr_sizeof(Vlen{String}), odr_sizeof(Vlen{String})+odr_sizeof(Vlen{RelOffset})],
      [H5TYPE_VLEN_UTF8, VariableLengthDatatype(ReferenceDatatype()), VariableLengthDatatype(ReferenceDatatype())]
      )

# ODR for UnionDataType
const UnionTypeODR = OnDiskRepresentation{
      (0, odr_sizeof(Vlen{String}), odr_sizeof(Vlen{String})+odr_sizeof(Vlen{RelOffset})),
      Tuple{String, Vector{Any}, Vector{Any}},
      Tuple{Vlen{String}, Vlen{RelOffset}, Vlen{RelOffset}}}

function h5fieldtype(f::JLDFile, ::Type{T}, readas::Type{S}, ::Initialized) where {T<:Union,S<:Union}
    @lookup_committed f Union
    commit(f, H5TYPE_UNION, Union, Union)
end
function h5fieldtype(f::JLDFile, ::Type{T}, readas::Type, ::Initialized) where T<:Union
    @lookup_committed f readas
    commit(f, H5TYPE_UNION, Union, readas)
end

fieldodr(::Type{T}, ::Bool) where {T<:Union} = UnionTypeODR()
h5fieldtype(f::JLDFile, ::Type{Union{}}, ::Initialized) = nothing
fieldodr(::Type{Union{}}, initialized::Bool) = nothing

h5type(f::JLDFile, ::Type{T}, x::Any) where {T<:Union} =
    h5fieldtype(f, Union, typeof(x), Val{true})
odr(::Type{Union}) = fieldodr(Union, true)

function h5convert!(out::Pointers, ::UnionTypeODR, f::JLDFile, x::Union, wsession::JLDWriteSession)
    # Write a description of the union type
    # This is not needed for loading the file but makes h5dump output clearer
    t = string(x)
    store_vlen!(out, UInt8, f, unsafe_wrap(Vector{UInt8}, t), f.datatype_wsession)
    out += odr_sizeof(Vlen{String})
    dts = filter(t -> t isa DataType, Base.uniontypes(x))
    uls = filter(t -> t isa UnionAll, Base.uniontypes(x))
    if !isempty(dts)
        refs = refs_from_types(f, dts, wsession)
        store_vlen!(out, RelOffset, f, refs, f.datatype_wsession)
    else
        h5convert_uninitialized!(out, Vlen{RelOffset})
    end
    if !isempty(uls)
        refs = RelOffset[write_ref(f, x, wsession) for x in uls]
        store_vlen!(out+odr_sizeof(Vlen{RelOffset}), RelOffset, f, refs, f.datatype_wsession)
    else
        h5convert_uninitialized!(out+odr_sizeof(Vlen{RelOffset}), Vlen{RelOffset})
    end
end

function jlconvert(::ReadRepresentation{Union, UnionTypeODR()}, f::JLDFile,
                   ptr::Ptr, header_offset::RelOffset)
    # Skip union type description in the beginning
    ptr += odr_sizeof(Vlen{String})
    # Reconstruct a Union by reading a list of DataTypes and UnionAlls
    # Lookup of RelOffsets is taken from jlconvert of DataTypes
    datatypes, = types_from_refs(f, ptr)
    unionalls, = types_from_refs(f, ptr+odr_sizeof(Vlen{RelOffset}))

    v = Union{datatypes..., unionalls...}
    track_weakref!(f, header_offset, v)
    v
end

function constructrr(::JLDFile, ::Type{T}, dt::CompoundDatatype, ::Vector{ReadAttribute}) where {T<:Union}
    dt == H5TYPE_UNION ? (ReadRepresentation{Union,UnionTypeODR()}(), true) :
                         throw(UnsupportedFeatureException())
end

## UnionAll

const UnionAllODR = OnDiskRepresentation{(0, 8),Tuple{TypeVar,Any},Tuple{JLD2.RelOffset,JLD2.RelOffset}}

# This needs its own h5convert! method, since otherwise we will attempt to specialize the
# generic h5convert! method for the specific UnionAll type rather than for UnionAll
# more generally.
function h5convert!(out::Pointers,
                    odr::UnionAllODR,
                    f::JLDFile, x::UnionAll, wsession::JLDWriteSession)
    h5convert!(out, RelOffset, f, x.var, f.datatype_wsession)
    h5convert!(out+odr_sizeof(RelOffset), RelOffset, f, x.body, f.datatype_wsession)
end

# The following two definitions are borrowed from BSON.
# They are used to generate instances of arbitrary types
# given their fields regarless of potential constructors.
# It is unclear to the author whether this approach is
# optimal.
newstruct(T) = ccall(:jl_new_struct_uninit, Any, (Any,), T)

function newstruct(T, fields)
    if !ismutabletype(T)
        return ccall(:jl_new_structv, Any, (Any,Ptr{Cvoid},UInt32), T, fields, length(fields))
    else
        # Manual inline of newstruct! to work around bug
        # https://github.com/MikeInnes/BSON.jl/issues/2#issuecomment-452204339
        x = newstruct(T)

        for (i, f) = enumerate(fields)
            ccall(:jl_set_nth_field, Nothing, (Any, Csize_t, Any), x, i-1, f)
        end
        x
    end
end

# jlconvert for empty objects
function jlconvert(::ReadRepresentation{T,nothing}, f::JLDFile, ptr::Ptr,
                              header_offset::RelOffset) where T
    T.size == 0 && return newstruct(T)::T

    # In this case, T is a non-empty object, but the written data was empty
    # because the custom serializers for the fields all resulted in empty
    # objects
    fields = map(T.types) do ty
        writtenas = writeas(ty)
        @assert writtenas.size == 0
        if writtenas === ty
            # This will usually equal `ty()` unless ty does not have a
            # constructor without arguments
            jlconvert(ReadRepresentation{ty,nothing}(), f, ptr, header_offset)
        else
            rconvert(ty,
                jlconvert(ReadRepresentation{writtenas,nothing}(), f, ptr, header_offset)
            )
        end
    end
    if T <: Tuple
        # Tuples are weird in that you can't instantiate them with Tuple{T,S}(t,s)
        return (fields...,)::T
    end
    return newstruct(T, fields)::T
end

# odr gives the on-disk representation of a given type, similar to
# fieldodr, but actually encoding the data for things that odr stores
# as references
function odr(::Type{T}) where T
    if !hasdata(T)
        # A pointer singleton or ghost. We need to write something, but we'll
        # just write a single byte.
        return nothing
    elseif isbitstype(T) && samelayout(T)
        # Has a specialized convert method or is an unpadded type
        return T
    end

    offsets = zeros(Int, length(T.types))
    odrs = Vector{Any}(undef, length(T.types))
    offset = 0
    for i = 1:length(T.types)
        ty = T.types[i]
        writtenas = writeas(ty)
        fodr = fieldodr(writtenas, i <= ninitialized(T))
        if writtenas !== ty && fodr !== nothing
            odrs[i] = CustomSerialization{writtenas,fodr}
        else
            odrs[i] = fodr
        end
        offsets[i] = offset
        offset += odr_sizeof(fodr)
    end

    OnDiskRepresentation{(offsets...,), Tuple{T.types...}, Tuple{odrs...}}()
end

abstract type DataMode end
struct ReferenceFree <: DataMode end
struct HasReferences <: DataMode end

datamode(::Type{CustomSerialization{WrittenAs,ODR}}) where {WrittenAs,ODR} = datamode(ODR)
datamode(::Union{Type{<:Vlen},Type{RelOffset}}) = HasReferences()
datamode(::DataType) = ReferenceFree()
datamode(::FixedLengthString) = ReferenceFree()
datamode(::Nothing) = ReferenceFree()
@generated function datamode(odr::OnDiskRepresentation{Offsets,JLTypes,H5Types} where {Offsets,JLTypes}) where H5Types
    for ty in H5Types.parameters
        datamode(ty) == HasReferences() && return HasReferences()
    end
    return ReferenceFree()
end

ext: jl


text:
export getHelmholtzOperator

"""
	H = getHelmholtzOperator(m,gamma,omega,Mesh)

	build Helmholtz operator
		H = Laplace - omega^2 * MassMatrix(m,gamma)

	Input:
		m                  - model
		gamma              - attenuation
		omega::Number      - frequency
		Mesh::RegularMesh  - mesh

	Output
		H::SparseMatrixCSC
"""
function getHelmholtzOperator(m,gamma,omega::Number,Mesh::RegularMesh)

Lap   = getNodalLaplacianMatrix(Mesh)
M     = getMassMatrix(vec(m),vec(gamma),Mesh)
# Get the Helmholtz operator (note the sign)
H = Lap .- omega^2 * M
return H
end


text:
"""
dataManager.jl s'occupe de l'accès aux données.
"""

const DATAPATH = "../data";

# retourne un Array{Station,1} d'objets Station
function getStations(fileName::String)
    f = openFile(fileName) # ouverture du fichier
    nbStations = parse(Int,readline(f)) # nombre de stations
    stations = Vector{Station}(undef,nbStations)
    for indexStation in 1:nbStations
        line = readline(f)
        splitLine = split(line," ")
        stations[indexStation] = Station(indexStation, splitLine[1],parse(Float64,splitLine[3]),parse(Float64,splitLine[2])) # construction de la station
    end
    return stations
end

# retourne un Array{Townsite,1} d'objets Townsite
function getTownsites(fileName::String,stations::Array{Station,1})
    f = openFile(fileName) # ouverture du fichier
    nbTownsites = parse(Int,readline(f)) # nombre de townsites
    townsites = Vector{Townsite}(undef,nbTownsites)
    for indexTownsite in 1:nbTownsites
        line = readline(f)
        splitLine = split(line," ")
        nbNearStations = length(splitLine) - 4 # nombre de stations proches
        nearStations = Vector{Station}(undef,nbNearStations)
        for indexStation in 1:nbNearStations
            nearStations[indexStation] = getStation(string(splitLine[4+indexStation]),stations) # on les stock dans un vecteur
        end
        townsites[indexTownsite] = Townsite(splitLine[1],parse(Float64,splitLine[3]),parse(Float64,splitLine[2]),parse(Int,splitLine[4]),nearStations)
    end
    return townsites
end

# retourne la station du vecteur stations dont le nom est stationName
function getStation(stationName::String,stations::Array{Station,1})
    found = false
    indexStation = 1
    while !found && indexStation <= length(stations)
        if stations[indexStation].name == stationName # on a trouvé la station
            found = true
        end
        indexStation += 1
    end
    return stations[indexStation-1]
end

# retourne le vecteur des densités des lotissements
function getDensities(townsites::Array{Townsite,1})
    n = length(townsites)
    densities = Vector{Int}(undef,n)
    for indexTownsite in 1:n
        densities[indexTownsite] = townsites[indexTownsite].density
    end
    return densities
end

# retourne le fichier ouvert correspondant au nom donné
function openFile(fileName::String)
    actual_path = pwd()
    cd(joinpath(DATAPATH)) # déplacement dans le repertoire où sont stockés les fichiers source
    f = open(fileName)
    cd(actual_path)
    return f
end

max_stars_repo_name: LucasBaussay/ProjetODC

max_issues_repo_name: LucasBaussay/ProjetODC

max_forks_repo_name: LucasBaussay/ProjetODC


text:
"""
`RCropSize <: ImageTransformation`

Description
============

Crops out an area of the specified pixel dimensions
at a randomized position of the given image

Usage
======

    RCropSize(64, 64)

    RCropSize(width = 64, height = 64)

Arguments
==========

- **`width`** : The desired width or the cropped out subimage
in pixels

- **`height`** : The desired height or the cropped out subimage
in pixels

Methods
========

- **`transform`** : Applies the transformation to the given Image
or set of images

Author(s)
==========

- Christof Stocker (Github: https://github.com/Evizero)

Examples
========

    using Augmentor
    using TestImages

    # load an example image
    img = testimage("lena")

    # Crop out a 64x32 area in the center of the image
    img_new = transform(RCropSize(64, 32), img)

see also
=========

`ImageTransformation`, `transform`
"""
immutable RCropSize <: ImageTransformation
    width::Int
    height::Int

    function RCropSize(width::Int, height::Int)
        width > 0 || throw(ArgumentError("width has to be greater than 0"))
        height > 0 || throw(ArgumentError("height has to be greater than 0"))
        new(width, height)
    end
end

RCropSize(; width::Int = 64, height::Int = 64) = RCropSize(width, height)

Base.showcompact(io::IO, op::RCropSize) = print(io, "Crop a random $(op.width)x$(op.height) window.")
multiplier(::RCropSize) = 1

function transform{T<:AbstractImage}(op::RCropSize, img::T)
    w = width(img)
    h = height(img)
    op.width <= w || throw(ArgumentError("image width is smaller than desired subimage"))
    op.height <= h || throw(ArgumentError("image height is smaller than desired subimage"))

    x_max = w - op.width + 1
    @assert x_max > 0
    y_max = h - op.height + 1
    @assert y_max > 0

    x = rand(1:x_max)
    y = rand(1:y_max)

    result = crop(img, x:(x+op.width-1), y:(y+op.height-1))
    _log_operation!(result, op)
end


text:
# Authors: Ali Siahkoohi, alisk@gatech.edu
# Date: March 2022

using DrWatson
@quickactivate :ConditionalNFs4Imaging

using InvertibleNetworks
using HDF5
using Random
using Statistics
using ProgressMeter
using Flux

# Random seed
Random.seed!(19)

args = read_config("train_amortized_imaging.json")
args = parse_input_args(args)

max_epoch = args["max_epoch"]
lr = args["lr"]
lr_step = args["lr_step"]
batchsize = args["batchsize"]
n_hidden = args["n_hidden"]
depth = args["depth"]
sim_name = args["sim_name"]
resume = args["resume"]

# Loading the existing weights, if any.
loaded_keys = resume_from_checkpoint(args, ["Params", "fval", "fval_eval", "opt", "epoch"])
Params = loaded_keys["Params"]
fval = loaded_keys["fval"]
fval_eval = loaded_keys["fval_eval"]
opt = loaded_keys["opt"]
init_epoch = loaded_keys["epoch"]

# Define raw data directory
mkpath(datadir("training-data"))
data_path = datadir("training-data", "training-pairs.h5")

# Download the dataset into the data directory if it does not exist
if isfile(data_path) == false
    run(`wget https://www.dropbox.com/s/53u8ckb9aje8xv4/'
        'training-pairs.h5 -q -O $data_path`)
end

# Load seismic images and create training and testing data
file = h5open(data_path, "r")
X_train = file["dm"][:, :, :, :]
Y_train = file["rtm"][:, :, :, :]
close(file)

Y_train[1:10, :, :, :] .= 0.0f0

nx, ny, nc, nsamples = size(X_train)

AN_x = ActNorm(nsamples)
AN_y = ActNorm(nsamples)
X_train = AN_x.forward(X_train)
Y_train = AN_y.forward(Y_train)

AN_params_x = get_params(AN_x)
AN_params_y = get_params(AN_y)

# Split in training/validation
ntrain = Int(floor((nsamples * 0.9)))
train_idx = randperm(nsamples)[1:ntrain]
val_idx = shuffle(setdiff(1:nsamples, train_idx))[1:64]

# Dimensions after wavelet squeeze to increase no. of channels
nx = Int(nx / 2)
ny = Int(ny / 2)
n_in = Int(nc * 4)

# Create network
CH = NetworkConditionalHINT(n_in, n_hidden, depth)
Params != nothing && put_params!(CH, convert(Array{Any,1}, Params))
CH = CH |> gpu

X_val = wavelet_squeeze(X_train[:, :, :, val_idx]) |> gpu
Y_val = wavelet_squeeze(Y_train[:, :, :, val_idx]) |> gpu

# Training
# Batch extractor
train_loader = Flux.DataLoader(train_idx, batchsize = batchsize, shuffle = true)
num_batches = length(train_loader)

# Optimizer
opt == nothing && (
    opt = Flux.Optimiser(
        Flux.ExpDecay(lr, 9.0f-1, num_batches * lr_step, 1.0f-6),
        Flux.ADAM(lr),
    )
)

# Training log keeper
fval == nothing && (fval = zeros(Float32, num_batches * max_epoch))
fval_eval == nothing && (fval_eval = zeros(Float32, max_epoch))

p = Progress(num_batches * (max_epoch - init_epoch + 1))

for epoch = init_epoch:max_epoch

    fval_eval[epoch] = loss_supervised(CH, X_val, Y_val; grad = false)

    for (itr, idx) in enumerate(train_loader)
        Base.flush(Base.stdout)

        # Augmentation
        if rand() > 5.0f-1
            X = X_train[:, end:-1:1, :, idx]
            Y = Y_train[:, end:-1:1, :, idx]
        else
            X = X_train[:, :, :, idx]
            Y = Y_train[:, :, :, idx]
        end

        # Apply wavelet squeeze.
        X = wavelet_squeeze(X)
        Y = wavelet_squeeze(Y)

        X = X |> gpu
        Y = Y |> gpu

        fval[(epoch-1)*num_batches+itr] = loss_supervised(CH, X, Y)[1]

        ProgressMeter.next!(
            p;
            showvalues = [
                (:Epoch, epoch),
                (:Itreration, itr),
                (:NLL, fval[(epoch-1)*num_batches+itr]),
                (:NLL_eval, fval_eval[epoch]),
            ],
        )

        # Update params
        for p in get_params(CH)
            Flux.update!(opt, p.data, p.grad)
        end
        clear_grad!(CH)
    end

    # Saving parameters and logs
    Params = get_params(CH) |> cpu
    save_dict =
        @strdict epoch max_epoch lr lr_step batchsize n_hidden depth sim_name Params fval fval_eval train_idx AN_params_x AN_params_y opt
    @tagsave(
        datadir(sim_name, savename(save_dict, "jld2"; digits = 6)),
        save_dict;
        safe = true
    )
end


text:
using PtFEM

data = Dict(
  # Frame(nels, nn, ndim, nst, nip, finite_element(nod, nodof))
  :struc_el => Frame(4, 5, 3, 1, 1, Line(2, 2)),
  :properties => [5.0e5;],
  :x_coords => [0.0, 1.25, 3.5, 4.0, 2.0],
  :y_coords => [0.0, 3.0, 2.0, 1.0, 1.5],
  :z_coords => [0.0, 0.0, 0.0, 0.0, 3.0],
  :g_num => [1 2 3 4;5 5 5 5],
  :support => [
    (1, [0 0 0]),
    (2, [0 0 0]),
    (3, [0 0 0]),
    (4, [0 0 0])
    ],
  :loaded_nodes => [
    (5, [20.0 -20.0 30.0])],
  :fixed_freedoms => [
    (5, 2, -0.0005)],
  :penalty => 1e19
)

data |> display
println()

@time fem, dis_df, fm_df = p42(data)
println()

display(dis_df)
println()
display(fm_df)


text:
function Makie.plot(result::FastAI.LRFinderResult)
    ticks = [round((10.)^i, digits=abs(i)) for i in -10:2]
    fig = Figure()
    ax = Axis(
        title = "Learning rate finder",
        titlesize = 20,
        fig[1, 1],
        xscale = log,
        xticks = (ticks, string.(ticks)),
        xminorticks = IntervalsBetween(5),
        xminorgridvisible=true,
        ygridcolor = :white,
        xlabelsize = 14,
        ylabelsize = 14,
        ylabelcolor = :gray,
        xtickcolor= :gray,
        xticklabelcolor= :black,
        xticklabelsize= 12,
        ytickcolor= :gray,
        yticklabelcolor= :gray,
        yticklabelsize= 12,
        ylabel = "Loss",
        xlabel = "Learning rate (log)")

    lines!(
        result.lrs,
        smoothvalues(result.losses, 0.98),
        color = :black,
    )

    hidespines!(ax)

    # plot suggestions
    ls = []
    for (estim, val) in zip(result.estimators, result.estimates)
        push!(ls, vlines!(ax, [val]))

    end

    leg = Legend(
        fig[2, 1], ls,
        ["$(estim): $(round(val, sigdigits=3))" for (estim, val) in zip(result.estimators, result.estimates)],
        framevisible=false,
        labelsize=14,
        orientation = :horizontal,)
    leg.tellheight[] = true
    leg.tellwidth[] = false

    return fig
end


text:
function load_desc(tree::ResDirTree)
    read(descfile(tree), String);
end

function save_desc(tree::ResDirTree, desc::String)
    write(descfile(tree), desc);
end

max_stars_repo_name: josePereiro/MyResultsManager

max_issues_repo_name: josePereiro/MyResultsManager

max_forks_repo_name: josePereiro/MyResultsManager


text:
"""
    TBReader(logdir)

Creates a TBReader object that can be used to deserialize data but cannot be used to
write.
"""
struct TBReader
    logdir::String
end

"""
An Union of types that can be deserialized with TensorBoardLogger.Deserialization
"""
TBReadable = Union{TBLogger, TBReader}

logdir(reader::TBReader) = reader.logdir


text:
using RecipesBase

@recipe function plot(h::Surface; xlims::NTuple{2, <:Number}, ylims::NTuple{2, <:Number})
    if !(xlims isa Tuple && ylims isa Tuple)
        error("Please provide 'xlims' and 'ylims' to plot an infinite surface.")
    end
    margin = 0.1
    h ∩ rectangle(xlims[1]-margin, ylims[1]-margin,
                  xlims[2]+margin, ylims[2]+margin)
end

@recipe function plot(c::ConvexPolygon)
	seriestype --> :shape
	linealpha --> 0.9
	fillalpha --> 0.3
	legend --> false
	showaxis --> false
	aspect_ratio --> :equal

	if isempty(c)
		[], []
	else
		v = vertices(c)
		push!(v, v[1])
		[p[1] for p in v], [p[2] for p in v]
	end
end

@recipe function plot(cs::Reunion{ConvexPolygon{T}}) where T
	seriestype --> :shape
	linealpha --> 0.9
	fillalpha --> 0.3
	legend --> false
	showaxis --> false
	aspect_ratio --> :equal

	v = []
	for c in cs.content
		if !isempty(c)
			append!(v, vertices(c))
			push!(v, vertices(c)[1])
			push!(v, Point(NaN, NaN))
		end
	end
	if length(v) > 0
		push!(v, v[1])
		[p[1] for p in v], [p[2] for p in v]
	else
		[], []
	end
end


text:
"""
Locally open locations of function calls in the source code
"""
module Where

import Base: typesof, insert!

include("main.jl")
export
    where,
    @where

end  # module


text:
using PowerModels
using Ipopt
import JSON

function run_powermodels(json_path)
    pm = PP2PM.load_pm_from_json(json_path)
    result = PowerModels.run_pf(pm, ACPPowerModel, Ipopt.IpoptSolver(),
                                    setting = Dict("output" => Dict("branch_flows" => true)))
    return result
end

max_stars_repo_name: Zamwell/pandapower

max_issues_repo_name: Zamwell/pandapower

max_forks_repo_name: gdgarcia/pandapower


text:
# This file is a part of DynareJulia. License is GNU GPL v.3: https://www.gnu.org/licenses/.

"""
    ExtendedDates

The `ExtendedDates` module extends Julia `Dates` module by providing dates at different frequencies:
YearDate`, `SemesterDate`, `QuarterDate`, `MonthDate`, `WeekDate`, `DayDate` and `UndatedDate` types.
All of them use their own proleptic Gregorian calendar at corresponding frequency. This is primarily aimed
at providing the time index of timeseries at various frequencies.

`Dates.Period` types `Semester` and `Undated` are also provided.

```jldoctest
julia> y = YearDate(2017)
2017

julia> s = SemesterDate(2017, 2)
2017-S2

julia> q = QuarterDate(2017, 3)
2017-Q3

julia> m = MonthDate(2017, 10)
2017-10

julia> w = WeekDate(2017, 52)
2017-W52

julia> d = DayDate(2017, 9, 15)
2017-09-15

julia> u = UndatedDate(3)
3

```

Please see the manual for more information.
"""
module ExtendedDates

import Base: ==, div, fld, mod, rem, gcd, lcm, +, -, *, /, %, broadcast
using Printf: @sprintf

using Base.Iterators
using Dates

include("types.jl")
include("periods.jl")
include("accessors.jl")
include("query.jl")
include("arithmetic.jl")
include("conversions.jl")
include("ranges.jl")
include("adjusters.jl")
include("rounding.jl")
include("io.jl")
include("parse.jl")

# re-export all Dates names
export Period,
    DatePeriod,
    TimePeriod,
    Year,
    Quarter,
    Month,
    Week,
    Day,
    Hour,
    Minute,
    Second,
    Millisecond,
    Microsecond,
    Nanosecond,
    TimeZone,
    UTC,
    TimeType,
    DateTime,
    Date,
    Time,
    # periods.jl
    canonicalize,
    # accessors.jl
    yearmonthday,
    yearmonth,
    monthday,
    year,
    month,
    week,
    day,
    hour,
    minute,
    second,
    millisecond,
    dayofmonth,
    microsecond,
    nanosecond,
    # query.jl
    dayofweek,
    isleapyear,
    daysinmonth,
    daysinyear,
    dayofyear,
    dayname,
    dayabbr,
    dayofweekofmonth,
    daysofweekinmonth,
    monthname,
    monthabbr,
    quarterofyear,
    dayofquarter,
    Monday,
    Tuesday,
    Wednesday,
    Thursday,
    Friday,
    Saturday,
    Sunday,
    Mon,
    Tue,
    Wed,
    Thu,
    Fri,
    Sat,
    Sun,
    January,
    February,
    March,
    April,
    May,
    June,
    July,
    August,
    September,
    October,
    November,
    December,
    Jan,
    Feb,
    Mar,
    Apr,
    May,
    Jun,
    Jul,
    Aug,
    Sep,
    Oct,
    Nov,
    Dec,
    # conversions.jl
    unix2datetime,
    datetime2unix,
    now,
    today,
    rata2datetime,
    datetime2rata,
    julian2datetime,
    datetime2julian,
    # adjusters.jl
    firstdayofweek,
    lastdayofweek,
    firstdayofmonth,
    lastdayofmonth,
    firstdayofyear,
    lastdayofyear,
    firstdayofquarter,
    lastdayofquarter,
    adjust,
    tonext,
    toprev,
    tofirst,
    tolast,
    # io.jl
    ISODateTimeFormat,
    ISODateFormat,
    ISOTimeFormat,
    DateFormat,
    RFC1123Format,
    @dateformat_str

# ExtendDates exported names
export Semester,
    Undated,
    YearDate,
    SemesterDate,
    QuarterDate,
    MonthDate,
    WeekDate,
    DayDate,
    UndatedDate,
    semester,
    quarter,
    undated,
    simpleperiod,
    FYear,
    FSemester,
    FQuarter,
    FMonth,
    FWeek,
    FDay,
    FUndated

end # module


text:
using InvertibleNetworks, LinearAlgebra, Test, Random

# Data
nx = 16
ny = 16
n_in = 4
batchsize = 4
X_prev = randn(Float32, nx, ny, n_in, batchsize)
X_curr = randn(Float32, nx, ny, n_in, batchsize)

# Network
architecture = ((0, 6), (1, 12), (0, 24), (-1, 12), (0, 6))
α0 = .2f0
H = NetworkHyperbolic(n_in, architecture; α=α0)

###################################################################################################
# Test invertibility

X_prev_, X_curr_ = H.inverse(H.forward(X_prev, X_curr)[1:2]...)
@show norm(X_curr - X_curr_), norm(X_curr_)
@test isapprox(norm(X_curr - X_curr_)/norm(X_curr), 0f0; atol=1e-2)

X_prev_, X_curr_ = H.forward(H.inverse(X_prev, X_curr)...)[1:2]
@show norm(X_curr - X_curr_), norm(X_curr_)
@test isapprox(norm(X_curr - X_curr_)/norm(X_curr), 0f0; atol=1e-2)

####################################################################################################
# Training

# Loss
function loss(H, X_prev, X_curr, ind)
    Y_prev, Y_curr, lgdet = H.forward(X_prev, X_curr)
    f = -log_likelihood(tensor_cat(Y_prev, Y_curr)) - lgdet
    ΔY = -∇log_likelihood(tensor_cat(Y_prev, Y_curr))
    ΔY_prev, ΔY_curr = tensor_split(ΔY)
    ΔX_prev, ΔX_curr, X_prev, X_curr = H.backward(ΔY_prev, ΔY_curr, Y_prev, Y_curr)
    return f, ΔX_prev, ΔX_curr, H.HL[ind].W.grad
end

# Data
X_prev = randn(Float32, nx, ny, n_in, batchsize)
X_curr = randn(Float32, nx, ny, n_in, batchsize)
X0_prev = randn(Float32, nx, ny, n_in, batchsize)
X0_curr = randn(Float32, nx, ny, n_in, batchsize)
dX_prev = X_prev - X0_prev
dX_curr = X_curr - X0_curr

# Gradient test w.r.t. input X0
Y_prev, Y_curr = H.forward(X_prev, X_curr)
f0, ΔX_prev, ΔX_curr = loss(H, X0_prev, X0_curr, 1)[1:3]
h = 0.1f0
maxiter = 6
err1 = zeros(Float32, maxiter)
err2 = zeros(Float32, maxiter)

print("\nGradient test hyperbolic network\n")
for j=1:maxiter
    f = loss(H, X0_prev + h*dX_prev,X0_curr + h*dX_curr, 1)[1]
    err1[j] = abs(f - f0)
    err2[j] = abs(f - f0 - h*dot(dX_prev, ΔX_prev) - h*dot(dX_curr, ΔX_curr))
    print(err1[j], "; ", err2[j], "\n")
    global h = h/2f0
end

@test isapprox(err1[end] / (err1[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err2[end] / (err2[1]/4^(maxiter-1)), 1f0; atol=1f1)

@show rate_1 = sum(err1[1:end-1]./err1[2:end])/(maxiter - 1)
@show rate_2 = sum(err2[1:end-1]./err2[2:end])/(maxiter - 1)

# Gradient test w.r.t. weights of hyperbolic network
H = NetworkHyperbolic(n_in, architecture; α=α0)
H0 = NetworkHyperbolic(n_in, architecture; α=α0)
H.forward(X_prev, X_curr)
H0.forward(X_prev, X_curr)   # evaluate to initialize actnorm layer

maxiter = 6
err3 = zeros(Float32, maxiter)
err4 = zeros(Float32, maxiter)

for i=1:length(H0.HL)
    h1 = 0.1f0
    print("\nGradient test invertible layer for layer $(i)\n")
    W0s = H0.HL[i].W.data
    dWg = H.HL[i].W.data - W0s; dWg *= norm(W0s)/norm(dWg)
    f01, ΔX1_prev, ΔX1_curr, ΔWg = loss(H0, X_prev, X_curr, i)
    for j=1:maxiter
        H0.HL[i].W.data = W0s + h1*dWg
        f = loss(H0, X_prev, X_curr, i)[1]
        err3[j] = abs(f - f01)
        err4[j] = abs(f - f01 - h1*dot(dWg, ΔWg))
        print(err3[j], "; ", err4[j], "\n")
        h1 = h1/2f0
    end

    @show local rate_1 = sum(err3[1:end-1]./err3[2:end])/(maxiter - 1)
    @show local rate_2 = sum(err4[1:end-1]./err4[2:end])/(maxiter - 1)
    H0.HL[i].W.data = W0s
    @test isapprox(err3[end] / (err3[1]/2^(maxiter-1)), 1f0; atol=2f1)
    @test isapprox(err4[end] / (err4[1]/4^(maxiter-1)), 1f0; atol=2f1)
end

###################################################################################################
# Jacobian-related tests

# Gradient test

# Initialization
H = NetworkHyperbolic(n_in, architecture; α=α0);
H.forward(randn(Float32, nx, ny, n_in, batchsize), randn(Float32, nx, ny, n_in, batchsize))
θ = deepcopy(get_params(H))
H0 = NetworkHyperbolic(n_in, architecture; α=α0);
H0.forward(randn(Float32, nx, ny, n_in, batchsize), randn(Float32, nx, ny, n_in, batchsize))
θ0 = deepcopy(get_params(H0))
X_prev = randn(Float32, nx, ny, n_in, batchsize)
X_curr = randn(Float32, nx, ny, n_in, batchsize)

# Perturbation (normalized)
dθ = θ-θ0; dθ .*= norm.(θ0)./(norm.(dθ).+1f-10)
dX_prev = randn(Float32, nx, ny, n_in, batchsize); dX_prev *= norm(X_prev)/norm(dX_prev)
dX_curr = randn(Float32, nx, ny, n_in, batchsize); dX_curr *= norm(X_curr)/norm(dX_curr)

# Jacobian eval
dY_prev, dY_curr, Y_prev, Y_curr = H.jacobian(dX_prev, dX_curr, dθ, X_prev, X_curr)

# Test
print("\nJacobian test\n")
h = 0.1f0
maxiter = 5
err5 = zeros(Float32, maxiter)
err6 = zeros(Float32, maxiter)
for j=1:maxiter
    set_params!(H, θ+h*dθ)
    Y_prev_, Y_curr_, _ = H.forward(X_prev + h*dX_prev, X_curr + h*dX_curr)
    err5[j] = norm(tensor_cat(Y_prev_, Y_curr_) - tensor_cat(Y_prev, Y_curr))
    err6[j] = norm(tensor_cat(Y_prev_, Y_curr_) - tensor_cat(Y_prev, Y_curr) - tensor_cat(h*dY_prev, h*dY_curr))
    print(err5[j], "; ", err6[j], "\n")
    global h = h/2f0
end

@test isapprox(err5[end] / (err5[1]/2^(maxiter-1)), 1f0; atol=1f1)
@test isapprox(err6[end] / (err6[1]/4^(maxiter-1)), 1f0; atol=1f1)

# Adjoint test
set_params!(H, θ)
dY_prev, dY_curr, Y_prev, Y_curr = H.jacobian(dX_prev, dX_curr, dθ, X_prev, X_curr)
dY_prev_ = randn(Float32, size(dY_prev)); dY_curr_ = randn(Float32, size(dY_curr))

dX_prev_, dX_curr_, dθ_, _, _ = H.adjointJacobian(dY_prev_, dY_curr_, Y_prev, Y_curr)

a = dot(tensor_cat(dY_prev, dY_curr), tensor_cat(dY_prev_, dY_curr_))
b = dot(tensor_cat(dX_prev, dX_curr), tensor_cat(dX_prev_, dX_curr_)) + dot(dθ, dθ_)
@test isapprox(a, b; rtol=1f-3)


text:
module DiffableDistributions

    ###########
    # imports #
    ###########
    using ForwardDiff
    typealias FDN ForwardDiff.ForwardDiffNumber

    import Distributions:
        Normal, AbstractNormal

    ############
    # includes #
    ############
    include("types.jl")

    ###########
    # exports #
    ###########
    export DNormal

end


text:
"""
# Description

Check if an expression contains grammatical rules that Volcanito can't handle.
If so, throw an error. Otherwise, return `nothing` to indicate that the code
was valid.

# Arguments

1. `e::Any`: An expression.

# Return Values

1. `_::Nothing`: If `e` is valid, return `nothing`. Otherwise an error is
    thrown.

# Examples

```
julia> validate(:x)

```
"""
function validate(@nospecialize(e::Any))::Nothing
    heads = (
        :comprehension,
        :block,
        :import,
        :export,
        :while,
        :for,
        :break,
        :continue,
        :let,
        :struct,
    )

    if isa(e, Expr)
        for head in heads
            if e.head == head
                error("Expr's with head == $head are not supported")
            end
        end
    end

    if isa(e, Expr)
        for i in 1:length(e.args)
            validate(e.args[i])
        end
        nothing
    elseif isa(e, Symbol)
        nothing
    else
        nothing
    end
end


text:
using Test
using Timeout

@testset "Timeout" begin
    t = @elapsed @test_throws InterruptException @timeout 0.5 sleep(10)
    @test t < 10
    @test_throws ErrorException @timeout 0.5 throw(ErrorException("bar"))
end