#####
##### Multitapered periodogram
#####

struct MTConfig{T,R1,F,P,T1,T2,W,R2}
    n_samples::Int
    fs::R1
    nfft::Int
    ntapers::Int
    freq::F
    plan::P
    fft_input_tmp::T1
    fft_output_tmp::T2
    window::W
    onesided::Bool
    r::R2 # inverse normalization; e.g. equal to `fs*N` for an unwindowed/untapered periodogram of a signal of length `N`
          # e.g. equal to `fs*ntapers*ones(ntapers)` when tapered by properly normalized window functions (i.e. norm-2 equal to 1 for each taper)
          # can be adjusted to weight the tapers, e.g. by the eigenvalues of the DPSS windows
    function MTConfig{T}(n_samples, nfft, ntapers, freq, fs, plan, fft_input_tmp,
                         fft_output_tmp, window, onesided, r) where {T}
        n_samples > 0 || throw(ArgumentError("`n_samples` must be positive"))
        nfft >= n_samples || throw(ArgumentError("Must have `nfft >= n_samples`"))
        ntapers > 0 || throw(ArgumentError("`ntapers` must be positive"))
        fs > 0 || throw(ArgumentError("`fs` must be positive"))
        if size(plan) != size(fft_input_tmp)
            throw(DimensionMismatch("""Must have `size(plan) == size(fft_input_tmp)`;
                got `size(plan)` = $(size(plan)) and `size(fft_input_tmp)` = $(size(fft_input_tmp))"""))
        end
        if size(fft_input_tmp) != (nfft,)
            throw(DimensionMismatch("""Must have `size(fft_input_tmp) == (nfft,)`;
                got `size(fft_input_tmp)` = $(size(fft_input_tmp)) and `nfft` = $(nfft)"""))
        end
        if size(fft_output_tmp) != (length(freq),)
            throw(DimensionMismatch("""Must have `size(fft_output_tmp) == (length(freq),)`;
                got `size(fft_output_tmp)` = $(size(fft_output_tmp)) and `length(freq)` = $(length(freq))"""))
        end
        if size(window) != (n_samples, ntapers)
            throw(DimensionMismatch("""Must have `size(window) == (n_samples, ntapers)`;
                got `size(window)` = $(size(window)) and `(ntapers, n_samples)` = $((n_samples, ntapers))"""))
        end
        if size(r) != (ntapers,)
            throw(DimensionMismatch("""Must have `size(r) == (ntapers,)`;
            got `size(r)` = $(size(r)) and `(ntapers,)` = $((ntapers,))"""))
        end
        if onesided && T <: Complex
            throw(ArgumentError("cannot compute one-sided FFT of a complex signal"))
        end
        return new{T,typeof(fs),typeof(freq),typeof(plan),typeof(fft_input_tmp),
                   typeof(fft_output_tmp),typeof(window),typeof(r)}(n_samples, fs, nfft,
                                                                    ntapers, freq, plan,
                                                                    fft_input_tmp,
                                                                    fft_output_tmp, window,
                                                                    onesided, r)
    end
end

# provides an `MTConfig` with `keep_only_large_evals` and `weight_by_evals` options used in some of the reference tests.
function dpss_config(::Type{T}, n_samples; nw=4, ntapers = 2nw-1, fs=1, keep_only_large_evals=false,
                     weight_by_evals=false, kwargs...) where {T}

    window = dpss(n_samples, nw, ntapers)

    if keep_only_large_evals
        evals = dpsseig(window, nw)
        keep_evals_mask = evals .> 0.9
        window = window[:, keep_evals_mask]
        evals = evals[keep_evals_mask]
        ntapers = size(window, 2)
    else
        evals = nothing
    end

    if weight_by_evals
        if evals === nothing
            evals = dpsseig(window, nw)
        end
        taper_weights = evals ./ sum(evals)
    else
        taper_weights = fill(1/ntapers, ntapers)
    end

    return MTConfig{T}(n_samples; window=window, nw=nw, ntapers=ntapers, taper_weights=taper_weights, fs=fs, kwargs...)
end

"""
    MTConfig{T}(n_samples; fs=1,
            nfft = nextpow(2, n_samples),
            window = nothing,
            nw = 4,
            ntapers = 2 * nw - 1,
            taper_weights = fill(1/ntapers, ntapers),
            onesided::Bool=T<:Real,
            fft_flags = FFTW.MEASURE)

Creates a config object which holds the configuration state
and temporary variables used in multitaper computations,
e.g. [`mt_pgram!`](@ref), [`mt_spectrogram`](@ref), [`MTSpectrogramConfig`](@ref),
[`MTCrossSpectraConfig`](@ref), and [`MTCoherenceConfig`](@ref).

An `MTConfig` can be re-used between computations as long
as none of the input arguments change.

* `n_samples`: the number of samples to be used as input when computing multitaper periodograms
  with this configuration. Used for pre-allocating temporary buffers.
* `fs`: the number of samples per second of the input signal
* `nfft`: length of input vector to the FFT; if `nfft > n_samples`, then the
  input signal will be zero-padded until it is of length `nfft`.
* `window`: window function to use for tapering. If left at the default of `nothing`,
  `window` will be set to `dpss(n_samples, nw, ntapers)`.
* `ntapers`: the number of tapers to use.
* `taper_weights = fill(1/ntapers, ntapers)`: how to weight the contribution of each taper.
  The default setting is to simply average them.
* `onesided`: whether or not to compute a "one-sided" FFT by using that real signal data
  yields conjugate-symmetry in Fourier space.
* `fft_flags`: flags to control how the FFT plan is generated.
"""
function MTConfig{T}(n_samples; fs=1, nfft=nextpow(2, n_samples), window=nothing, nw=4,
                     ntapers=2 * nw - 1, taper_weights = fill(1/ntapers, ntapers),
                    onesided::Bool=T <: Real,
                     fft_flags=FFTW.MEASURE) where {T}
    if onesided && T <: Complex
        throw(ArgumentError("cannot compute one-sided FFT of a complex signal"))
    end
    n_samples > 0 || throw(ArgumentError("`n_samples` must be positive"))
    nfft >= n_samples || throw(ArgumentError("Must have `nfft >= n_samples`"))
    freq = onesided ? rfftfreq(nfft, fs) : fftfreq(nfft, fs)
    fft_input_tmp = Vector{T}(undef, nfft)
    fft_output_tmp = Vector{fftouttype(T)}(undef, length(freq))
    plan = onesided ? plan_rfft(fft_input_tmp; flags=fft_flags) :
           plan_fft(fft_input_tmp; flags=fft_flags)
    if window === nothing
        r = fs ./ taper_weights
        window = dpss(n_samples, nw, ntapers)
    else
        r = fs .* vec(sum(abs2, window; dims=1)) ./ taper_weights
    end

    return MTConfig{T}(n_samples, nfft, ntapers, freq, fs, plan, fft_input_tmp,
                       fft_output_tmp, window, onesided, r)
end

function allocate_output(config::MTConfig{T}) where {T}
    return Vector{fftabs2type(T)}(undef, length(config.freq))
end

# Internal function used in [`mt_pgram!`](@ref) and [`mt_cross_power_spectra!`](@ref).
function mt_fft_tapered!(fft_output, signal, taper_index, config)
    # Create the input: tapered + zero-padded version of the signal
    fft_input = config.fft_input_tmp
    @inbounds for i in 1:length(signal)
        fft_input[i] = config.window[i, taper_index] * signal[i]
    end
    fft_input[(length(signal) + 1):(config.nfft)] .= 0

    # do the FFT
    mul!(fft_output, config.plan, fft_input)
    return nothing
end

"""
    mt_pgram(s; onesided=eltype(s)<:Real, nfft=nextfastfft(n), fs=1, nw=4, ntapers=iceil(2nw)-1, window=dpss(length(s), nw, ntapers))
    mt_pgram(signal::AbstractVector, config::MTConfig)

Computes the multitaper periodogram of a signal `s`.

If `window` is not specified, the signal is tapered with
`ntapers` discrete prolate spheroidal sequences with
time-bandwidth product `nw`. Each sequence is equally weighted;
adaptive multitaper is not (yet) supported.

If `window` is specified, each column is applied as a taper. The
sum of periodograms is normalized by the total sum of squares of
`window`.

Returns a `Periodogram`.

See also [`mt_pgram!`](@ref) and [`MTConfig`](@ref).
"""
mt_pgram

function mt_pgram(s::AbstractVector{T}; onesided::Bool=eltype(s)<:Real,
                  nfft::Int=nextfastfft(length(s)), fs::Real=1,
                  nw::Real=4, ntapers::Int=ceil(Int, 2nw)-1,
                  window::Union{AbstractMatrix,Nothing}=nothing) where T<:Number
    config = MTConfig{T}(length(s); fs = fs,
        nfft = nfft, window = window, nw = nw, ntapers = ntapers, onesided = onesided,
        fft_flags = FFTW.ESTIMATE)
    out = allocate_output(config)
    return mt_pgram!(out, s, config)
end

function mt_pgram!(output, s::AbstractVector{T}; onesided::Bool=eltype(s)<:Real,
    nfft::Int=nextfastfft(length(s)), fs::Real=1,
    nw::Real=4, ntapers::Int=ceil(Int, 2nw)-1,
    window::Union{AbstractMatrix,Nothing}=nothing) where T<:Number
    config = MTConfig{T}(length(s); fs = fs,
    nfft = nfft, window = window, nw = nw, ntapers = ntapers, onesided = onesided,
    fft_flags = FFTW.ESTIMATE)
    return mt_pgram!(output, s, config)
end

function mt_pgram(signal::AbstractVector, config::MTConfig)
    out = allocate_output(config)
    return mt_pgram!(out, signal, config)
end

"""
    mt_pgram!(output, s::AbstractVector{T}; onesided::Bool=eltype(s)<:Real,
        nfft::Int=nextfastfft(length(s)), fs::Real=1,
        nw::Real=4, ntapers::Int=ceil(Int, 2nw)-1,
        window::Union{AbstractMatrix,Nothing}=nothing) where T<:Number
    mt_pgram!(output::AbstractVector, signal::AbstractVector, config::MTConfig) -> Periodogram

Computes a multitapered periodogram, storing the output in `output`. Arguments:

* `signal::AbstractVector`: should be of length `config.n_samples`
* `output::AbstractVector`: should be of length `length(config.freq)`

Optionally pass an [`MTConfig`](@ref) object to preallocate temporary variables
and choose configuration settings; otherwise, keyword arguments may be passed
to choose those settings.

Returns a `Periodogram`.

See also [`mt_pgram`](@ref) and [`MTConfig`](@ref).
"""
mt_pgram!

function mt_pgram!(output::AbstractVector, signal::AbstractVector, config::MTConfig)
    if length(output) != length(config.freq)
        throw(DimensionMismatch("""Expected `output` to be of length `length(config.freq)`;
                got `length(output) = $(length(output)) and `length(config.freq)` = $(length(config.freq))"""))
    end
    if length(signal) != config.n_samples
        throw(DimensionMismatch("""Expected `signal` to be of length `config.n_samples`;
                got `length(signal) == $(length(signal)) and `config.n_samples` = $(config.n_samples)"""))
    end
    fft_output = config.fft_output_tmp

    output .= 0
    for taper_index in 1:(config.ntapers)
        mt_fft_tapered!(fft_output, signal, taper_index, config)
        fft2pow!(output, fft_output, config.nfft, config.r[taper_index], config.onesided)
    end
    return Periodogram(output, config.freq)
end

#####
##### Multitapered spectrogram
#####

struct MTSpectrogramConfig{T,C<:MTConfig{T}}
    n_samples::Int
    n_overlap_samples::Int
    time::FloatRange{Float64}
    mt_config::C
end

"""
    MTSpectrogramConfig(n_samples, mt_config::MTConfig{T}, n_overlap_samples) where {T}
    MTSpectrogramConfig{T}(n_samples, samples_per_window, n_overlap_samples; fs=1, kwargs...) where {T}

Creates a `MTSpectrogramConfig` which holds configuration and temporary variables for [`mt_spectrogram`](@ref) and [`mt_spectrogram!`](@ref).
Any keyword arguments accepted by [`MTConfig`](@ref) may be passed here, or an `MTConfig` object itself.
"""
function MTSpectrogramConfig(n_samples::Int, mt_config::MTConfig{T}, n_overlap_samples::Int) where {T}
    samples_per_window = mt_config.n_samples
    if samples_per_window <= n_overlap_samples
        throw(ArgumentError("Need `samples_per_window > n_overlap_samples`; got `samples_per_window` = $(samples_per_window) and `n_overlap_samples` = $(n_overlap_samples)."))
    end

    f = samples_per_window / 2
    hop = samples_per_window - n_overlap_samples

    len = n_samples < samples_per_window ? 0 : div(n_samples - samples_per_window, samples_per_window - n_overlap_samples) + 1
    time = range(f, step = hop, length = len) / mt_config.fs
    return MTSpectrogramConfig{T,typeof(mt_config)}(n_samples, n_overlap_samples, time,
                                                    mt_config)
end

# add a method for if the user specifies the type, i.e. `MTSpectrogramConfig{T}`
function MTSpectrogramConfig{T}(n_samples::Int, mt_config::MTConfig{T}, n_overlap_samples::Int) where {T}
    return MTSpectrogramConfig(n_samples, mt_config, n_overlap_samples)
end

# Create the `MTConfig` if it's not passed
function MTSpectrogramConfig{T}(n_samples::Int, samples_per_window::Int,
                                n_overlap_samples::Int; fs=1, kwargs...) where {T}
    return MTSpectrogramConfig(n_samples, MTConfig{T}(samples_per_window; fs=fs, kwargs...), n_overlap_samples)
end

"""
    mt_spectrogram!(output, signal::AbstractVector{T}, n::Int=length(signal) >> 3,
        n_overlap::Int=n >> 1; fs::Int=1, onesided::Bool=T <: Real, kwargs...) where {T}
    mt_spectrogram!(destination::AbstractMatrix, signal::AbstractVector, config::MTSpectrogramConfig)

Computes a multitaper spectrogram using the parameters specified in `config`. Arguments:

* `destination`: `length(config.mt_config.freq)` x `length(config.time)` matrix. This can be created by `DSP.allocate_output(config)`.
* `signal`: vector of length `config.n_samples`
* `config`: optionally, pass an [`MTSpectrogramConfig`](@ref) object to hold temporary variables and configuration settings. Otherwise, settings arguments may be passed directly.

Returns a `Spectrogram`.

See also [`mt_spectrogram`](@ref).
"""
mt_spectrogram!

function mt_spectrogram!(output, signal::AbstractVector{T}, n::Int=length(signal) >> 3,
    n_overlap::Int=n >> 1; fs::Int=1, onesided::Bool=T <: Real,
    kwargs...) where {T}
    config = MTSpectrogramConfig{T}(length(signal), n, n_overlap; fs=fs, onesided=onesided,
        fft_flags=FFTW.ESTIMATE, kwargs...)
    return mt_spectrogram!(output, signal, config)
end

@views function mt_spectrogram!(destination::AbstractMatrix, signal::AbstractVector,
                                config::MTSpectrogramConfig)
    if size(destination) != (length(config.mt_config.freq), length(config.time))
        throw(DimensionMismatch("""Expected `destination` to be of size `(length(config.mt_config.freq), length(config.time))`;
                got `size(destination) == $(size(destination)) and
                `(length(config.mt_config.freq), length(config.time))` = $((length(config.mt_config.freq), length(config.time)))"""))
    end
    if length(signal) != config.n_samples
        throw(DimensionMismatch("""Expected `signal` to be of length `config.n_samples`;
                got `length(signal) = $(length(signal)) and `config.n_samples` = $(config.n_samples)"""))
    end
    samples_per_window = config.mt_config.n_samples
    subepochs = arraysplit(signal, samples_per_window, config.n_overlap_samples)
    @assert length(subepochs) == length(config.time)
    for (time_index, subepoch) in enumerate(subepochs)
        mt_pgram!(destination[:, time_index], subepoch, config.mt_config)
    end
    return Spectrogram(destination, config.mt_config.freq, config.time)
end

function allocate_output(config::MTSpectrogramConfig{T}) where {T}
    return Matrix{fftabs2type(T)}(undef, length(config.mt_config.freq), length(config.time))
end

function mt_spectrogram(signal::AbstractVector, config::MTSpectrogramConfig)
    output = allocate_output(config)
    return mt_spectrogram!(output, signal, config)
end

"""
    mt_spectrogram(signal::AbstractVector{T}, n::Int=length(s) >> 3,
                                  n_overlap::Int=n >> 1; fs::Int=1,
                                  onesided::Bool=T <: Real, kwargs...) where {T}
    mt_spectrogram(signal::AbstractVector, config::MTSpectrogramConfig)

Compute a multitaper spectrogram, returning a `Spectrogram` object.
Optionally pass a [`MTSpectrogramConfig`](@ref) object; otherwise, any additional keyword
arguments accepted by [`MTConfig`](@ref) may be passed to configure the tapering.

Returns a `Spectrogram`.

See also [`mt_spectrogram!`](@ref).
"""
mt_spectrogram

function mt_spectrogram(signal::AbstractVector{T}, n::Int=length(signal) >> 3,
                        n_overlap::Int=n >> 1; fs::Int=1, onesided::Bool=T <: Real,
                        kwargs...) where {T}
    config = MTSpectrogramConfig{T}(length(signal), n, n_overlap; fs=fs, onesided=onesided,
                                    fft_flags=FFTW.ESTIMATE, kwargs...)
    X = allocate_output(config)
    return mt_spectrogram!(X, signal, config)
end

"""
    mt_spectrogram(signal::AbstractVector, mt_config::MTConfig,
                   n_overlap::Int=mt_config.n_samples >> 1) -> Spectrogram

Compute a multitaper spectrogram using an an [`MTConfig`](@ref) object to choose
the window size. Note that all the workspace variables are contained in the [`MTConfig`](@ref) object, and this object can be re-used between spectrogram computations, so that only the output needs to be allocated.

Returns a `Spectrogram`.

See also [`mt_spectrogram!`](@ref).

## Example

```julia
signal1 = sin.(10_000) .+ randn(10_000)
mt_config = MTConfig{Float64}(1000) # 1000 samples per window
spec1 = mt_spectrogram(signal1, mt_config, 500)
signal2 = cos.(2_000) .+ randn(2_000)
spec2 = mt_spectrogram(signal2, mt_config, 250) # same `mt_config`, different signal, different overlap
```
"""
function mt_spectrogram(signal::AbstractVector, mt_config::MTConfig{T},
                        n_overlap::Int=mt_config.n_samples >> 1) where {T}
    config = MTSpectrogramConfig{T}(length(signal), mt_config, n_overlap)
    X = allocate_output(config)
    return mt_spectrogram!(X, signal, config)
end

#####
##### Multitapered cross spectral matrix
#####

# the following code was initially ported from MNE-python before being partially rewritten.
# MNE-python's 3-clause BSD license can be found here: <https://github.com/mne-tools/mne-python/blob/c6b22e82b64da0ed0f52a5f28bd4711e47090f9e/LICENSE.txt>.

"""
    CrossPowerSpectra{T,F,A<:AbstractArray{T, 3}}

Access the power (an `n_channels` x `n_channels` x `length(freq)` array) via the function [`power`](@ref), and the frequencies by the function [`freq`](@ref).

See also [`mt_cross_power_spectra`](@ref) and [`mt_cross_power_spectra!`](@ref).
"""
struct CrossPowerSpectra{T,F,A<:AbstractArray{T, 3}}
    power::A
    freq::F
end

freq(c::CrossPowerSpectra) = c.freq
power(c::CrossPowerSpectra) = c.power

function check_onesided_real(mt_config::MTConfig{T}) where T
    if !((T <: Real) && mt_config.onesided)
        throw(ArgumentError("Only real data is supported (with the default choice of `onesided=true`) for this operation."))
    end
    return nothing
end

struct MTCrossSpectraConfig{T,T1,T2,T3,T4,F,T5,T6,C<:MTConfig{T}}
    n_channels::Int
    normalization_weights::T1
    x_mt::T2
    demean::Bool
    mean_per_channel::T3
    demeaned_signal::T4
    freq::F
    freq_range::T5
    freq_inds::T6
    ensure_aligned::Bool
    mt_config::C
    function MTCrossSpectraConfig{T,T1,T2,T3,T4,F,T5,T6,C}(n_channels, normalization_weights,
            x_mt, demean, mean_per_channel, demeaned_signal, freq, freq_range,
            freq_inds, ensure_aligned, mt_config) where {T,T1,T2,T3,T4,F,T5,T6,C}
        check_onesided_real(mt_config) # this restriction is artifical; the code needs to be generalized
        return new{T,T1,T2,T3,T4,F,T5,T6,C}(n_channels, normalization_weights, x_mt,
            demean, mean_per_channel, demeaned_signal, freq, freq_range,
            freq_inds, ensure_aligned, mt_config)
    end
end

"""
    MTCrossSpectraConfig{T}(n_channels, n_samples; fs=1, demean=false, freq_range=nothing,
                            ensure_aligned = T == Float32 || T == Complex{Float32}, kwargs...) where {T}
    MTCrossSpectraConfig(n_channels, mt_config::MTConfig{T}; demean=false, freq_range=nothing,
                         ensure_aligned = T == Float32 || T == Complex{Float32})

Creates a configuration object used for [`mt_cross_power_spectra`](@ref) and [`mt_cross_power_spectra!`](@ref).

* `n_channels`: the number of channels of the input.
* `n_samples`: the number of samples for each channel of the input.
* `demean`: if `true`, the channelwise mean will be subtracted from the input signals before the cross spectral powers are computed.
* `freq_range`: if `nothing`, all frequencies are retained. Otherwise, only frequencies between `first(freq_range)` and `last(freq_range)` are retained.
* `ensure_aligned = T == Float32 || T == Complex{Float32}`: perform an extra copy to ensure that the FFT output is memory-aligned
* Additionally, either pass an [`MTConfig`](@ref) object, or keyword arguments such as `fs` accepted by [`MTConfig`](@ref).

Returns a `CrossPowerSpectra` object.
"""
function MTCrossSpectraConfig{T}(n_channels, n_samples; fs=1, demean=false,
                                 freq_range=nothing,
                                 ensure_aligned = T == Float32 || T == Complex{Float32},
                                 kwargs...) where {T}
    mt_config = MTConfig{T}(n_samples; fs=fs, kwargs...)
    return MTCrossSpectraConfig{T}(n_channels, mt_config; demean=demean, freq_range=freq_range, ensure_aligned=ensure_aligned)
end

# extra method to ensure it's ok to pass the redundant type parameter {T}
MTCrossSpectraConfig{T}(n_channels, mt_config::MTConfig{T}; demean=false,
        freq_range=nothing, ensure_aligned = T == Float32 || T == Complex{Float32}) where {T} = MTCrossSpectraConfig(n_channels, mt_config; demean=demean,
        freq_range=freq_range, ensure_aligned = ensure_aligned)

function MTCrossSpectraConfig(n_channels, mt_config::MTConfig{T}; demean=false,
        freq_range=nothing, ensure_aligned = T == Float32 || T == Complex{Float32}) where {T}

    n_samples = mt_config.n_samples
    if demean
        mean_per_channel = Vector{T}(undef, n_channels)
        demeaned_signal = Matrix{T}(undef, n_channels, n_samples)
    else
        mean_per_channel = nothing
        demeaned_signal = nothing
    end

    normalization_weights = 2 ./ mt_config.r

    x_mt = Array{fftouttype(T),3}(undef, length(mt_config.freq), mt_config.ntapers,
                                  n_channels)
    if freq_range !== nothing
        freq_mask = first(freq_range) .< mt_config.freq .< last(freq_range)
        freq_inds = findall(freq_mask)
        freq = mt_config.freq[freq_mask]
    else
        freq_inds = eachindex(mt_config.freq)
        freq = mt_config.freq
    end
    return MTCrossSpectraConfig{T,typeof(normalization_weights),typeof(x_mt),
                                typeof(mean_per_channel),typeof(demeaned_signal),
                                typeof(freq),typeof(freq_range),typeof(freq_inds),
                                typeof(mt_config)}(n_channels, normalization_weights, x_mt, demean,
                                                   mean_per_channel, demeaned_signal, freq,
                                                   freq_range, freq_inds, ensure_aligned, mt_config)
end

function allocate_output(config::MTCrossSpectraConfig{T}) where {T}
    return Array{fftouttype(T),3}(undef, config.n_channels, config.n_channels,
                                  length(config.freq))
end

"""
    mt_cross_power_spectra!(output, signal::AbstractMatrix; fs=1, kwargs...)
    mt_cross_power_spectra!(output, signal::AbstractMatrix, config::MTCrossSpectraConfig)

Computes multitapered cross power spectra between channels of a signal. Arguments:

* `output`: `n_channels` x `n_channels` x `length(config.freq)`. Can be created by `DSP.allocate_output(config)`.
* `signal`: `n_channels` x `n_samples`
* `config`: `MTCrossSpectraConfig{T}`: optionally pass a [`MTCrossSpectraConfig`](@ref) to
  preallocate temporary and choose configuration settings.
  Otherwise, one may pass any keyword arguments accepted by this object.

Produces a `CrossPowerSpectra` object holding the `n_channels` x `n_channels` x `n_frequencies`
output array and the corresponding frequencies (accessed by [`freq`](@ref)).

See also [`mt_cross_power_spectra`](@ref) and [`MTCrossSpectraConfig`](@ref).
"""
mt_cross_power_spectra!

function mt_cross_power_spectra!(output, signal::AbstractMatrix{T}; fs=1, kwargs...) where {T}
    n_channels, n_samples = size(signal)
    config = MTCrossSpectraConfig{T}(n_channels, n_samples; fs=fs, fft_flags=FFTW.ESTIMATE,
                                     kwargs...)
    return mt_cross_power_spectra!(output, signal, config)
end

@views function mt_cross_power_spectra!(output, signal::AbstractMatrix,
                                   config::MTCrossSpectraConfig)
    if size(signal) != (config.n_channels, config.mt_config.n_samples)
        throw(DimensionMismatch("Size of `signal` does not match `(config.n_channels, config.mt_config.n_samples)`;
        got `size(signal)`=$(size(signal)) but `(config.n_channels, config.mt_config.n_samples)`=$((config.n_channels, config.mt_config.n_samples))"))
    end
    if size(output) != (config.n_channels, config.n_channels, length(config.freq_inds))
        throw(DimensionMismatch("Size of `output` does not match `(config.n_channels, config.n_channels, length(config.freq_inds))`;
        got `size(output)`=$(size(output)) but `(config.n_channels, config.n_channels, length(config.freq_inds))`=$((config.n_channels, config.n_channels, length(config.freq_inds)))"))
    end

    if config.demean
        mean!(config.mean_per_channel, signal)
        config.demeaned_signal .= signal .- config.mean_per_channel
        signal = config.demeaned_signal
    end

    x_mt = config.x_mt

    if config.ensure_aligned
        mt_fft_tapered_multichannel_ensure_aligned!(x_mt, signal, config)
    else
        mt_fft_tapered_multichannel!(x_mt, signal, config)
    end
    x_mt[1, :, :] ./= sqrt(2)
    if iseven(config.mt_config.nfft)
        x_mt[end, :, :] ./= sqrt(2)
    end
    cs_inner!(output, config.normalization_weights, x_mt, config)
    return CrossPowerSpectra(output, config.freq)
end

function mt_fft_tapered_multichannel_ensure_aligned!(x_mt, signal, config)
    fft_output = config.mt_config.fft_output_tmp
    for k in 1:(config.n_channels), taper in 1:config.mt_config.ntapers
        # we do this in two steps so that we are sure `fft_output` has the memory alignment FFTW expects (without needing the `FFTW.UNALIGNED` flag)
        mt_fft_tapered!(fft_output, signal[k, :], taper, config.mt_config)
        x_mt[:, taper, k] .= fft_output
    end
end

@views function mt_fft_tapered_multichannel!(x_mt, signal, config)
    for k in 1:(config.n_channels), taper in 1:config.mt_config.ntapers
        mt_fft_tapered!(x_mt[:, taper, k], signal[k, :], taper, config.mt_config)
    end
end

function cs_inner!(output, normalization_weights, x_mt, config)
    freq_inds = config.freq_inds
    n_channels = config.n_channels
    @boundscheck checkbounds(output, 1:n_channels, 1:n_channels, 1:length(freq_inds))
    @boundscheck checkbounds(normalization_weights, 1:length(normalization_weights))
    @boundscheck checkbounds(x_mt, freq_inds, 1:length(normalization_weights), 1:n_channels)
    output .= zero(eltype(output))
    # Up to the `normalization_weights` scaling, we have
    # J_k^l(f) = x_mt[k, f, l]
    # Ŝ^lm(f) = output[l, m, f]
    # using the notation from https://en.wikipedia.org/wiki/Multitaper#The_method
    # so the formula `Ŝ^lm(f) = \sum_k conj(J_k^l(f)) * (J_k^m(f))`` becomes the following loop:
    @inbounds for (fi, f) in enumerate(freq_inds),
                  m in 1:n_channels,
                  l in 1:n_channels,
                  k in 1:length(normalization_weights)

        output[l, m, fi] += normalization_weights[k] * x_mt[f, k, l] * conj(x_mt[f, k, m])
    end
    return nothing
end

"""
    mt_cross_power_spectra(signal::AbstractMatrix{T}; fs=1, kwargs...) where {T}
    mt_cross_power_spectra(signal::AbstractMatrix, config::MTCrossSpectraConfig)

Computes multitapered cross power spectra between channels of a signal. Arguments:

* `signal`: `n_channels` x `n_samples`
* Optionally pass an [`MTCrossSpectraConfig`](@ref) object to preallocate temporary variables
and choose configuration settings. Otherwise, any keyword arguments accepted by [`MTCrossSpectraConfig`](@ref) may be passed here.

Produces a `CrossPowerSpectra` object holding the `n_channels` x `n_channels` x `n_frequencies`
output array (accessed by [`power`](@ref)) and the corresponding frequencies (accessed by [`freq`](@ref)).

See also [`mt_cross_power_spectra!`](@ref) and [`MTCrossSpectraConfig`](@ref).
"""
mt_cross_power_spectra

function mt_cross_power_spectra(signal::AbstractMatrix{T}; fs=1, kwargs...) where {T}
    n_channels, n_samples = size(signal)
    config = MTCrossSpectraConfig{T}(n_channels, n_samples; fs=fs, fft_flags=FFTW.ESTIMATE,
                                     kwargs...)
    return mt_cross_power_spectra(signal, config)
end

function mt_cross_power_spectra(signal::AbstractMatrix, config::MTCrossSpectraConfig)
    output = allocate_output(config)
    return mt_cross_power_spectra!(output, signal, config)
end

#####
##### Multitapered coherences
#####

struct MTCoherenceConfig{T,T1,C<:MTCrossSpectraConfig{T}}
    cs_matrix::T1
    cs_config::C
end

"""
    MTCoherenceConfig{T}(n_channels, n_samples; fs=1, demean=false, freq_range=nothing, kwargs...) where T
    MTCoherenceConfig(cs_config::MTCrossSpectraConfig{T}) where {T}
    MTCoherenceConfig(n_channels, mt_config::MTConfig{T}; demean=false, freq_range=nothing,
        ensure_aligned = T == Float32 || T == Complex{Float32}) where {T}

Creates a configuration object for coherences from a [`MTCrossSpectraConfig`](@ref). Provides a helper method
with the same arugments as `MTCrossSpectraConfig` to construct the `MTCrossSpectraConfig` object.

See also [`mt_coherence`](@ref) and [`mt_coherence!`](@ref).
"""
function MTCoherenceConfig(cs_config::MTCrossSpectraConfig{T}) where {T}
    cs_matrix = allocate_output(cs_config)
    return MTCoherenceConfig{T,typeof(cs_matrix),typeof(cs_config)}(cs_matrix, cs_config)
end

# add a method to cover the case in which the user specifies the `{T}` here
MTCoherenceConfig{T}(cs_config::MTCrossSpectraConfig{T}) where {T} = MTCoherenceConfig(cs_config)

function MTCoherenceConfig{T}(n_channels, n_samples; fs=1, demean=false,
                              freq_range=nothing, kwargs...) where {T}
    cs_config = MTCrossSpectraConfig{T}(n_channels, n_samples; fs=fs, demean=demean,
                                        freq_range=freq_range, kwargs...)
    cs_matrix = allocate_output(cs_config)
    return MTCoherenceConfig{T,typeof(cs_matrix),typeof(cs_config)}(cs_matrix, cs_config)
end

# ensure it's OK to pass the extra {T} type parameter
function MTCoherenceConfig{T}(n_channels, mt_config::MTConfig{T}; demean=false,
    freq_range=nothing, ensure_aligned = T == Float32 || T == Complex{Float32}) where T
    return MTCoherenceConfig(n_channels, mt_config; demean=demean,
    freq_range=freq_range, ensure_aligned = ensure_aligned)
end

function MTCoherenceConfig(n_channels, mt_config::MTConfig{T}; demean=false,
    freq_range=nothing, ensure_aligned = T == Float32 || T == Complex{Float32}) where {T}
    cs_config = MTCrossSpectraConfig(n_channels, mt_config; demean=demean,
              freq_range=freq_range, ensure_aligned=ensure_aligned)
    cs_matrix = allocate_output(cs_config)
    return MTCoherenceConfig{T,typeof(cs_matrix),typeof(cs_config)}(cs_matrix, cs_config)
end

function allocate_output(config::MTCoherenceConfig{T}) where {T}
    return Array{real(T)}(undef, config.cs_config.n_channels, config.cs_config.n_channels, length(config.cs_config.freq))
end

"""
    coherence_from_cs!(output::AbstractArray, cs_matrix)

Compute the pairwise channel coherences from a cross spectral matrix, storing the result in `output`.
"""
function coherence_from_cs!(output::AbstractArray{T}, cs_matrix) where T
    n_channels = size(cs_matrix, 2)
    n_freqs = size(cs_matrix, 3)
    @boundscheck checkbounds(output, 1:n_channels, 1:n_channels, 1:n_freqs)
    @boundscheck checkbounds(cs_matrix, 1:n_channels, 1:n_channels, 1:n_freqs)
    output .= zero(T)
    @inbounds for ch2 in 1:n_channels
        for ch1 in (ch2 + 1):n_channels # lower triangular matrix
            for f in 1:n_freqs
                output[ch1, ch2, f] += abs(cs_matrix[ch1, ch2, f]) /
                                         sqrt(real(cs_matrix[ch1, ch1, f] *
                                               cs_matrix[ch2, ch2, f]))
            end
        end
    end
    output .+= PermutedDimsArray(output, (2, 1, 3)) # symmetrize
    # diagonal elements should be `1` for any frequency
    @inbounds for i in 1:n_channels
        output[i, i, :] .= one(T)
    end
    return nothing
end

"""
    Coherence

Holds an `n_channels` x `n_channels` x `length(freq)` array consisting of the pairwise coherences between each channel for each frequency which is accessed by [`coherence`](@ref), as well
as the frequencies accessed by [`freq`](@ref).

See also [`mt_coherence`](@ref) and [`mt_coherence!`](@ref).
"""
struct Coherence{T,F,A<:AbstractArray{T, 3}}
    coherence::A
    freq::F
end

freq(c::Coherence) = c.freq

"""
    coherence(c::Coherence)

Given an `Coherence` object, returns an `n_channels` x `n_channels` x `length(freq(c))`
array consisting of the pairwise coherences between each channel for each frequency.
"""
coherence(c::Coherence) = c.coherence

"""
    mt_coherence!(output, signal::AbstractMatrix; fs=1, freq_range=nothing, demean=false, kwargs...)
    mt_coherence!(output, signal::AbstractMatrix, config::MTCoherenceConfig)

Computes the pairwise coherences between channels.

* `output`: `n_channels` x `n_channels` matrix
* `signal`: `n_samples` x `n_channels` matrix
* `config`: optional configuration object that pre-allocates temporary variables and choose settings.

Returns a `Coherence` object.

See also [`mt_coherence`](@ref) and [`MTCoherenceConfig`](@ref).
"""
mt_coherence!

function mt_coherence!(output, signal::AbstractMatrix,
                       config::MTCoherenceConfig)
    if size(signal) != (config.cs_config.n_channels, config.cs_config.mt_config.n_samples)
        throw(DimensionMismatch("Size of `signal` does not match `(config.cs_config.n_channels, config.cs_config.mt_config.n_samples)`;
            got `size(signal)`=$(size(signal)) but `(config.cs_config.n_channels, config.cs_config.mt_config.n_samples)`=$((config.cs_config.n_channels, config.cs_config.mt_config.n_samples))"))
    end
    if size(output) != (config.cs_config.n_channels, config.cs_config.n_channels, length(config.cs_config.freq))
        throw(DimensionMismatch("Size of `output` does not match `(config.cs_config.n_channels, config.cs_config.n_channels, length(config.cs_config.freq))`;
        got `size(output)`=$(size(output)) but `(config.cs_config.n_channels, config.cs_config.n_channels, length(config.cs_config.freq))`=$((config.cs_config.n_channels, config.cs_config.n_channels, length(config.cs_config.freq)))"))
    end
    cs = mt_cross_power_spectra!(config.cs_matrix, signal, config.cs_config)
    coherence_from_cs!(output, power(cs))

    return Coherence(output, config.cs_config.freq)
end

function mt_coherence!(output, signal::AbstractMatrix{T}; fs=1, freq_range=nothing, demean=false,
    kwargs...) where {T}
    n_channels, n_samples = size(signal)
    config = MTCoherenceConfig{T}(n_channels, n_samples; fs=fs, demean=demean,
                freq_range=freq_range,
                fft_flags=FFTW.ESTIMATE, kwargs...)
    return mt_coherence!(output, signal, config)
end

"""
    mt_coherence(signal::AbstractMatrix{T}; fs=1, freq_range = nothing, demean=false, kwargs...) where T
    mt_coherence(signal::AbstractMatrix, config::MTCoherenceConfig)

Arguments:

* `signal`: `n_channels` x `n_samples` matrix
* Optionally pass an `MTCoherenceConfig` to pre-allocate temporary variables and choose configuration settings, otherwise, see [`MTCrossSpectraConfig`](@ref) for the meaning of the keyword arguments.

Returns a `Coherence` object.

See also [`mt_coherence`](@ref) and [`MTCoherenceConfig`](@ref).
"""
mt_coherence

function mt_coherence(signal::AbstractMatrix{T}; fs=1, freq_range=nothing, demean=false,
                      kwargs...) where {T}
    n_channels, n_samples = size(signal)
    config = MTCoherenceConfig{T}(n_channels, n_samples; fs=fs, demean=demean,
                                  freq_range=freq_range,
                                  fft_flags=FFTW.ESTIMATE, kwargs...)
    return mt_coherence(signal, config)
end

function mt_coherence(signal::AbstractMatrix, config::MTCoherenceConfig)
    cohs = allocate_output(config)
    return mt_coherence!(cohs, signal, config)
end


#=
Single-point derivatives of scalar->scalar maps.
=#
function finite_difference_derivative(f, x::T, fdtype::Type{T1}=Val{:central},
    returntype::Type{T2}=eltype(x), f_x::Union{Nothing,T}=nothing) where {T<:Number,T1,T2}

    epsilon = compute_epsilon(fdtype, x)
    if fdtype==Val{:forward}
        return (f(x+epsilon) - f(x)) / epsilon
    elseif fdtype==Val{:central}
        return (f(x+epsilon) - f(x-epsilon)) / (2*epsilon)
    elseif fdtype==Val{:complex} && returntype<:Real
        return imag(f(x+im*epsilon)) / epsilon
    end
    fdtype_error(returntype)
end
# Single point derivative implementations end here.

#=
Multi-point implementations of scalar derivatives for efficiency.
=#
struct DerivativeCache{CacheType1, CacheType2, fdtype, returntype}
    fx      :: CacheType1
    epsilon :: CacheType2
end

function DerivativeCache(
    x          :: AbstractArray{<:Number},
    fx         :: Union{Nothing,AbstractArray{<:Number}} = nothing,
    epsilon    :: Union{Nothing,AbstractArray{<:Real}} = nothing,
    fdtype     :: Type{T1} = Val{:central},
    returntype :: Type{T2} = eltype(x)) where {T1,T2}

    if fdtype==Val{:complex} && !(eltype(returntype)<:Real)
        fdtype_error(returntype)
    end

    if fdtype!=Val{:forward} && typeof(fx)!=Nothing
        @warn("Pre-computed function values are only useful for fdtype==Val{:forward}.")
        _fx = nothing
    else
        # more runtime sanity checks?
        _fx = fx
    end

    if typeof(epsilon)!=Nothing && typeof(x)<:StridedArray && typeof(fx)<:Union{Nothing,StridedArray} && 1==2
        @warn("StridedArrays don't benefit from pre-allocating epsilon.")
        _epsilon = nothing
    elseif typeof(epsilon)!=Nothing && fdtype==Val{:complex}
        @warn("Val{:complex} makes the epsilon array redundant.")
        _epsilon = nothing
    else
        if typeof(epsilon)==Nothing || eltype(epsilon)!=real(eltype(x))
          epsilon = fill(zero(real(eltype(x))), size(x))
        end
        _epsilon = epsilon
    end
    DerivativeCache{typeof(_fx),typeof(_epsilon),fdtype,returntype}(_fx,_epsilon)
end

#=
Compute the derivative df of a scalar-valued map f at a collection of points x.
=#
function finite_difference_derivative(
    f,
    x          :: AbstractArray{<:Number},
    fdtype     :: Type{T1} = Val{:central},
    returntype :: Type{T2} = eltype(x),      # return type of f
    fx         :: Union{Nothing,AbstractArray{<:Number}} = nothing,
    epsilon    :: Union{Nothing,AbstractArray{<:Real}} = nothing) where {T1,T2}

    df = fill(zero(returntype), size(x))
    finite_difference_derivative!(df, f, x, fdtype, returntype, fx, epsilon)
end

function finite_difference_derivative!(
    df         :: AbstractArray{<:Number},
    f,
    x          :: AbstractArray{<:Number},
    fdtype     :: Type{T1} = Val{:central},
    returntype :: Type{T2} = eltype(x),
    fx         :: Union{Nothing,AbstractArray{<:Number}} = nothing,
    epsilon    :: Union{Nothing,AbstractArray{<:Real}}   = nothing) where {T1,T2}

    cache = DerivativeCache(x, fx, epsilon, fdtype, returntype)
    finite_difference_derivative!(df, f, x, cache)
end

function finite_difference_derivative!(df::AbstractArray{<:Number}, f, x::AbstractArray{<:Number},
    cache::DerivativeCache{T1,T2,fdtype,returntype}) where {T1,T2,fdtype,returntype}

    fx, epsilon = cache.fx, cache.epsilon
    if typeof(epsilon) != Nothing
        epsilon_factor = compute_epsilon_factor(fdtype, eltype(x))
        @. epsilon = compute_epsilon(fdtype, x, epsilon_factor)
    end
    if fdtype == Val{:forward}
        if typeof(fx) == Nothing
            @. df = (f(x+epsilon) - f(x)) / epsilon
        else
            @. df = (f(x+epsilon) - fx) / epsilon
        end
    elseif fdtype == Val{:central}
        @. df = (f(x+epsilon) - f(x-epsilon)) / (2 * epsilon)
    elseif fdtype == Val{:complex} && returntype<:Real
        epsilon_complex = eps(eltype(x))
        @. df = imag(f(x+im*epsilon_complex)) / epsilon_complex
    else
        fdtype_error(returntype)
    end
    df
end

#=
Optimized implementations for StridedArrays.
Essentially, the only difference between these and the AbstractArray case
is that here we can compute the epsilon one by one in local variables and avoid caching it.
=#
function finite_difference_derivative!(df::StridedArray, f, x::StridedArray,
    cache::DerivativeCache{T1,T2,fdtype,returntype}) where {T1,T2,fdtype,returntype}

    epsilon_factor = compute_epsilon_factor(fdtype, eltype(x))
    if fdtype == Val{:forward}
        fx = cache.fx
        @inbounds for i ∈ eachindex(x)
            epsilon = compute_epsilon(Val{:forward}, x[i], epsilon_factor)
            x_plus = x[i] + epsilon
            if typeof(fx) == Nothing
                df[i] = (f(x_plus) - f(x[i])) / epsilon
            else
                df[i] = (f(x_plus) - fx[i]) / epsilon
            end
        end
    elseif fdtype == Val{:central}
        @inbounds for i ∈ eachindex(x)
            epsilon = compute_epsilon(Val{:central}, x[i], epsilon_factor)
            epsilon_double_inv = one(typeof(epsilon)) / (2*epsilon)
            x_plus, x_minus = x[i]+epsilon, x[i]-epsilon
            df[i] = (f(x_plus) - f(x_minus)) * epsilon_double_inv
        end
    elseif fdtype == Val{:complex}
        epsilon_complex = eps(eltype(x))
        @inbounds for i ∈ eachindex(x)
            df[i] = imag(f(x[i]+im*epsilon_complex)) / epsilon_complex
        end
    else
        fdtype_error(returntype)
    end
    df
end


using Gadfly, DataArrays, RDatasets

plot(dataset("ggplot2", "diamonds"), x="Price", Geom.histogram(bincount=30))


# ---
# title: 98. Validate Binary Search Tree
# id: problem98
# author: Indigo
# date: 2021-06-02
# difficulty: Medium
# categories: Tree, Depth-first Search, Recursion
# link: <https://leetcode.com/problems/validate-binary-search-tree/description/>
# hidden: true
# ---
#
# Given the `root` of a binary tree, _determine if it is a valid binary search
# tree (BST)_.
#
# A **valid BST** is defined as follows:
#
#   * The left subtree of a node contains only nodes with keys **less than** the node's key.
#   * The right subtree of a node contains only nodes with keys **greater than** the node's key.
#   * Both the left and right subtrees must also be binary search trees.
#
#
#
# **Example 1:**
#
# ![](https://assets.leetcode.com/uploads/2020/12/01/tree1.jpg)
#
#
#
#     Input: root = [2,1,3]
#     Output: true
#
#
# **Example 2:**
#
# ![](https://assets.leetcode.com/uploads/2020/12/01/tree2.jpg)
#
#
#
#     Input: root = [5,1,4,null,null,3,6]
#     Output: false
#     Explanation: The root node's value is 5 but its right child's value is 4.
#
#
#
#
# **Constraints:**
#
#   * The number of nodes in the tree is in the range `[1, 104]`.
#   * `-231 <= Node.val <= 231 - 1`
#
#
## @lc code=start
using LeetCode

function is_valid_BST(root)
    isnothing(root) && return true
    arr = Int[]
    function pre_tra(root::TreeNode{Int})
        pre_tra(root.left)
        push!(arr, root.val)
        pre_tra(root.right)
    end
    pre_tra(::Nothing) = nothing
    pre_tra(root)
    return issorted(arr)
end
## @lc code=end


#=
@doc """
    cmp(op1, op2)

Compare `op1` and `op2`. Return a positive value if `op1>op2`, zero if
`op1=op2`, and a negative value if `op1<op2`. Both `op1` and `op2` are
considered to their full own precision, which may differ. If one of
the operands is NaN, set the `erange` flag and return zero.

Note: These functions may be useful to distinguish the three possible
cases. If you need to distinguish two cases only, it is recommended to
use the predicate functions (e.g., `equal_p` for the equality); they
behave like the IEEE 754 comparisons, in particular when one or both
arguments are NaN. But only floating-point numbers can be compared
(you may need to do a conversion first).
"""
=#

#=
e has type mpft_exp_t which must be at least 32 bits and it is defined
in mpfr.h according to the value of _MPFR_EXP_FORMAT. I set to it
Cint.
=#
export cmp_2exp
for (specC, tJ, tC) in ((:ui, CulongMax,  Culong),
                        (:si, ClongMax,   Clong))
    @eval begin
        function cmp_2exp(op1::BigFloat, op2::$tJ, e)
            isnan(op1) && throw(DomainError("op1 is NaN"))
            ccall(($(string("mpfr_cmp_", specC, "_2exp")), :libmpfr), Int32,
                  (Ref{BigFloat}, $tC, Cint), op1, op2, Int32(e))
        end
    end
end

@doc """
    cmp_2exp(op1, op2, e)

Compare `op1` and `op2 × 2^e`: returns a positive value if `op1>op2 ×
2^e`, zero if `op1=op2 × 2^e` or a negative number if `op1<op2 ×
2^e`. Throws an exception if `op1` is NaN.
""" cmp_2exp

for fJ = (:number_p, :regular_p)
    @eval begin
        export $fJ
        ($fJ)(op::BigFloat) =
            ccall(($(string(:mpfr_, fJ)), :libmpfr), Int32, (Ref{BigFloat},), op) != 0
    end
end

@doc """
    number_p(op)

Return true if `op` is an ordinary number (i.e., neither NaN nor an
infinity). Return false otherwise.
""" number_p

@doc """
    regular_p(op)

Return true if `op` is a regular number (i.e., neither NaN, nor an
infinity nor zero). Return false otherwise.
""" regular_p

export cmpabs
"""
    cmpabs(op1, op2)

Compare `|op1|` and `|op2|`. Return a positive value if `|op1|>|op2|`,
zero if `|op1|=|op2|`, and a negative value if `|op1|<|op2|`. Raise a
DomainError exception if one of the operands is NaN.
"""
function cmpabs(op1::BigFloat, op2::BigFloat)
    isnan(op1) || isnan(op2) && throw(DomainError("Operands must not be NaNs"))
    ccall((:mpfr_cmpabs, "libmpfr"), Int32, (Ref{BigFloat}, Ref{BigFloat}),
          op1, op2)
end

export unordered_p
"""
    unordered_p(op1, op2)

Return true if `op1` or `op2` is a NaN (i.e., they cannot be
compared), false otherwise.
"""
unordered_p(op1::BigFloat, op2::BigFloat) =
    ccall((:mpfr_unordered_p, "libmpfr"), Int32, (Ref{BigFloat}, Ref{BigFloat}),
          op1, op2) != 0


" The following initializes the master problem and runs a loop iterating between adding scenarios
and computing a solution for the current subset of scenarios. It is described in Algorithms 1 and 3 of the paper.
  * It is assumed throughout that the number of uncertain constraints is equal to data.ncons.
  * It is assumed throught that f(ξ,x) = ξᵀCx h(x).
  * Uncertain constraints are assumed to be linear functions g_ℓ(ξ,x) = ξᵀ(A_ℓ)x."

function scenario_generation(heuristic,warmup)
  # cc and aa record the uncertainty vectors ξ generated so far and multiplies them by the matrices C and A_ℓ.
  # Specifically:
  # cc[j,:] = Cᵀξ^j for each j in [m]
  # aa[ℓ][j,:] = Aᵀ_ℓ ξ^j for eacsh j in [m] and ℓ in [L]
  aa, cc = populate_scenarios()
  tΞ = [] #set of scenarios generated so far
  incumbent = []
  if !heuristic
    # If the algorithm is called in exact mode then a heuristic will be run first to obtain an upper bound
    # This upper bound serves to filter the set of solutions to be generated X(UB)
    if heuristicmode == "DUAL"
      @timeit to "heuristic" heurvalue,incumbent = heuristic_dualization()
    elseif heuristicmode == "HEURCG"
      @timeit to "heuristic" heurvalue,incumbent = scenario_generation(true,warmup)
    end
    # Current upper bound is set to the best solution value returned by the heuristic
    ub = heurvalue
    data.heuristic_ub = ub
    !warmup && @info "exact mode"
    # All the solutions in the set X(UB) are generated by using a problem-specific enumeration algorithm
    @timeit to "gen_sol" solutions = generate_solutions(ub)
  else
    # If the algoritm is called in heuristic mode then it is initalized simply with a set of solutions
    # Below this is done simply by starting with the optimal solution using the data of the first scenario and its neighbors
    !warmup && @info "heuristic mode"

    if data.constraint_uncertainty
      a_vec = aa[:][1,:]
    else
      a_vec = []
    end
    # A deterministic optimal solution is obtained using the data realization of first scenario
    detsol = find_new_solution(cc[1,:],a_vec,Inf,[],[],false,false,TIME_LIM)
    # The incumbent is initialized by creating k copies of this deterministic solution
    incumbent = Array{Int64}(undef, k, n)
    for k in K incumbent[k,:] = round_solution(detsol) end
    # The set of solutions is enriched by adding some neighbors of the incumbent
    @timeit to "gen_sol" solutions = generate_neighbors(incumbent,aa)
    # An initial (poor) upper bound is obtained by computing the robust cost of detsol
    ub = compute_basic_ub(detsol)
  end
  # In some pathological cases where the number of enumerated solutions is < k we add copies of the first solution
  for kk in (size(solutions,1)+1):k push!(solutions,solutions[1]) end
  # S is the number of solutions generated
  S = 1:size(solutions,1)
  # M is the number of scenarios added so far
  M = 1:size(cc,1)

  # Now we proceed to the iterations of the algorithm
  # solutions containts all solutions from X(UB) in the exact mode and a subset of these solutions in the heuristic mode

  # indicator new_scenario_added is used as a stopping criteria of convergence
  new_scenario_added = true

  # First initialize the structures that will be used throughout the algorithm
  # x_sol_array contains all solution vectors
  # sol_items contains the indices of non-zero variables in a solution vector
  x_sol_array = Vector{Vector{Float64}}()
  sol_items = Vector{Vector{Int}}()
  # Populate x_sol_array and sol_items using the enumerated solutions
  fill_solutions_bit(solutions, x_sol_array, sol_items)
  # The following vector contains h(x) for each x in X
  hh = [h(sol_items[s]) for s in S]
  # Starting lb: min_x∈X ξ^Tx for some ξ ∈ Ξ
  # In the following the initial scenario is used as ξ
  lb = minimum(sum(x_sol_array[s][l]*cc[1,l]*hh[s] for l in N) for s in S)

  # d records cost of each solution in each scenario, specifically
  # d[s,j] = ξ_jᵀCx_s h(x_s) for s in [r] and j in [m]
  d = Matrix{Float64}(undef,size(solutions,1),size(cc,1))
  # a_cov records constraint satisfaction of each solution in each scenario, specifically
  # (a_cov[ℓ][s,j] = 1) ⇔ (ξ_jᵀ(A_ℓ)x_s h(x_s) ≤ b_ℓ), for s in [r] and j in [m]
  a_cov = Vector{Matrix{Bool}}()
  if data.constraint_uncertainty
    for cons in 1:data.ncons
      push!(a_cov,Matrix{Bool}(undef,size(solutions,1),size(cc,1)))
    end
  end
  # initialize the number of iterations
  NUMITER = 0

  # calculate the initial gap
  gap = 100*round((ub - lb)/abs(ub), digits = 4)

  if time()-STARTTIME < TIME_LIM
    println("IT LB  UB  gap time")
    while new_scenario_added
      ξ = []
      NUMITER += 1
      if heuristic
        # In the heuristic mode d and a_cov are recalculated from scratch in each iteration
        # as this is not very costly with a small set of solutions
        @timeit to "re-calculate d" d = calculate_d(x_sol_array,cc,S,M,hh)
        if data.constraint_uncertainty
          @timeit to "re-calculate a_cov" a_cov = calculate_a(x_sol_array,aa,S,M)
        end
      else
        # In the exact mode d and a_cov are updated each time a new scenario is introduced
        # This permits reusing all d and a_cov values for the previously added scenarios
        @timeit to "update d" d = update_d(x_sol_array,cc,d,S,M,NUMITER,hh)
        if data.constraint_uncertainty
          @timeit to "update a_cov" a_cov = update_a(x_sol_array,aa,a_cov,S,M,NUMITER)
        end
      end
      # The master problem is solved using a binary search algorithm
      # It returns the current optimal solution and its value
      if !(data.constraint_uncertainty)
        @timeit to "binary search" x̃, ṽ_down, ṽ_up = binary_search(x_sol_array,sol_items,lb,ub,min(gap,1),d,a_cov,cc,δ)
      else
        @timeit to "binary search" x̃, ṽ_down, ṽ_up = binary_search(x_sol_array,sol_items,lb,ub,ϵ,d,a_cov,cc,δ)
      end

      # The optimal value of the master problem provides a lower bound
      lb = ṽ_down
      # If the master problem is solved succesfully then we proceed to the separation
      # If not then the binary search algorithm was stopped due to time limit restrictions
      # In this case it returns an empty solution permitting the algorithm to skip the separation step and break out of the while loop
      if length(x̃) > 0
        # Separation problem returns an identified scenario along with the separation objective value
        @timeit to "separation" ξ, sep_value = build_and_solve_separation(x̃, ṽ_up)
        # Because of the tolerance garantee when building the covcering matrix, it may happen that a given scenario be generated more than once
        # If that happens, we consider instead that the solution did not return any scenario.
        if data.constraint_uncertainty
          @timeit to "scenario generated?" begin
            if length(ξ) > 0 && length(tΞ) > 0
              for ξ′ in tΞ
                if norm(ξ-ξ′) < ϵ*min(norm(ξ),norm(ξ′))
                  @warn "scenario already generated"
                  ξ = []
                  sep_value = 0
                  break
                end
              end
            end
          end
        end
        length(ξ) > 0 && push!(tΞ,ξ)

        if !data.constraint_uncertainty
          # If there is no constraint uncertainty then upper bound updates are done whenever the separation value
          # improves upon the current bound. The upper bound is obtained simply as the separation value.
          # As the adversarial function is continuous this is a valid upper bound.
          if sep_value < ub  && sep_value > -Inf
            ub = sep_value
            incumbent = x̃
          end
        end
      end
      # Update the optimality gap
      gap = 100*round((ub - lb)/abs(ub), digits = 4)

      # If a scenario is returned from the separation problem and the convergence is not yet achieved
      # then add the realization of this scenario to vectors cc and aa.
      # Otherwise the indicator new_scenario_added is turned false to break out of the while loop.
      # The convergence is tested using the relative gap between ub and lb in the case where there is only objective uncertainty
      # and using the value returned by the function S(x̃,ṽ) in the case of contraint uncertainty
      if !(data.constraint_uncertainty)
        if length(ξ) == 0 || gap ≤ ϵ
          new_scenario_added = false
        end
      else
        if length(ξ) == 0 || sep_value ≤ ϵ
          new_scenario_added = false
        end
      end

      if new_scenario_added
        # We compute Cᵀξ and Aᵀ_ℓ ξ for the scenario ξ just obtained
        # println("Scenario added: $ξ")
        new_c, new_a = add_scenario(ξ)
        # M is increased by 1
        M = 1:(length(M)+1)
        # cc and aa are augmented with the new realization
        cc = [cc; new_c']
        if data.constraint_uncertainty
          for cons in 1:data.ncons
            aa[cons] = [aa[cons]; new_a[cons]']
          end
        end
      else
        # The current solution leads to S(x̃,ṽ) so the ub on its optimal value is an ub for the true problem
        ub = ṽ_up
        gap = 100*round((ub - lb)/abs(ub), digits = 4)
      end

      if heuristic && new_scenario_added
        # HEURISTIC MODE: the new scenario offers an opportunity to enrich the set of solutions
        if data.constraint_uncertainty
          a_vec = aa[:][end,:]
        else
          a_vec = []
        end
        # A good solution for the newly added scenario and its neighbors are added to the set of enumerated solutions
        @timeit to "relax and fix" new_solution = relax_and_fix(cc[end,:],a_vec,ub)

        if length(new_solution)>0
          @timeit to "generate neighbors" solutions_to_add = generate_neighbors(reshape(round_solution(new_solution),length(new_solution),1)',aa)
          for sol in solutions_to_add
            if !(sol in solutions)
              push!(solutions,sol)
            end
          end
        end
        # The following updates data-structures used throughout with the newly discovered solutions
        x_sol_array = Vector{Vector{Float64}}()
        sol_items = Vector{Vector{Int}}()
        fill_solutions_bit(solutions, x_sol_array, sol_items)
        S = 1:size(solutions,1)
        hh = [ h(sol_items[s]) for s in S ]
      end
      TotalTime = round(time()-STARTTIME, digits = 2)
      LowerBound = round(lb, digits = 4)
      println("$NUMITER $LowerBound $ub $gap $TotalTime")
      TotalTime > TIME_LIM && break
    end
  end
  if heuristic
    return ub,incumbent
  else
    return lb, ub, incumbent, gap, length(M), NUMITER, length(S)
  end
end

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

"The following computes an optimal solution for the current set of scenarios using
a binary search algorithm based on the optimal solution cost, in line with the radius
formulation used for the p-center problem. It is described in Algorithm 2 of the paper."

function binary_search(x_sol_array,sol_items,lb,ub,ϵ_BS,d,a_cov,cc,δ)

  " For a given coverage matrix identify the non-dominated solutions."
  function calculate_nondom_solutions(cov,last,M,S)
    map = [BitSet([j for j in M if cov[s, j]]) for s in S]
    discarded = falses(length(S))
    for s in 1:(length(S)+1-k)
      if length(S)-length(discarded)==k break end
      if discarded[s] continue end
      for ss in 1:(length(S)+1-k)
        if s == ss continue
        elseif discarded[ss] continue
        elseif issubset(map[ss], map[s])
            discarded[ss] = true
            break
        elseif issubset(map[s], map[ss])
            discarded[s] = true
            break
        end
      end
    end
    SS = [s for s in S if !discarded[s]]
    return SS
  end

  " For a given coverage matrix identify the non-dominated scenarios."
  function calculate_nondom_scenarios(cov,M,S)
    map = [BitSet([s for s in S if cov[s, j]]) for j in M]
    discarded = falses(length(M))
    for m in 1:(length(M)+1-k)
        if discarded[m] continue end
        for mm in 1:(length(M)+1-k)
          if m == mm continue
          elseif discarded[mm] continue
          elseif issubset(map[mm],map[m])
            discarded[m] = true
            break
          elseif issubset(map[m],map[mm])
            discarded[mm] = true
            break
          end
        end
    end
    MM = [m for m in M if !discarded[m]]
    return MM
  end

  " For a given radius, test whether the associated covering problem is feasible using
  a MILP formulation with slack variables."
  function build_and_solve_covering_MIP(SS,MM,cov)
    time_remaining = TIME_LIM - (time()-STARTTIME)
    cover = create_model(time_remaining)
    @variable(cover, y[SS], Bin)
    @variable(cover, slack[MM] ≥ 0)
    @constraint(cover, covering[m in MM], sum(y[s] for s in SS if cov[s,m]) >= 1-slack[m]) # covering constraints
    @constraint(cover, lessthank, sum(y[s] for s in SS) == k) # cardinality constraint
    @objective(cover, Min, sum(slack[m] for m in MM))
    optimize!(cover)
    sol_used = []
    if termination_status(cover) == MOI.OPTIMAL
      for s in SS
        if value(y[s]) > 0.9
          push!(sol_used,s)
        end
      end
      return objective_value(cover),sol_used
    else
      return Inf,sol_used
    end
  end

  " Translate the binary search solution to a solution vector x."
  function compute_x_values(S, sol_items, solutions_used)::Array{Int64,2}
      x = zeros(k, data.n)
      k_ = 0
      for s in solutions_used
          k_ += 1
          for l in sol_items[S[s]]
              x[k_ ,l] = 1
          end
      end
      return x
  end

  S = 1:length(x_sol_array)
  M = 1:size(d,2)
  # Indicator for verifying if this is the last iteration of the algorithm
  # When last=true a feasible solution corresponding to the ub is calculated
  last = false
  # Count the number of iterations of the binary search algorithm
  ITER_BS = 0
  # Initialize the optimal solution vector
  x = []
  # The while loop is executed at least for the first and the last iteration
  # and until the convergence criteria on the relative gap is met
  while (ub - lb)/abs(ub) > ϵ_BS/100 || last || ITER_BS == 0
    ITER_BS += 1
    time()-STARTTIME > TIME_LIM && break

    if last
      # In the last iteraion r is set large enough to find a feasible solution
      r = ub
    else
      # Otherwise binary search classical update of r
      r = (lb+ub)/2
    end
    # A coverage matrix is built based on whether a solution covers a scenario with a cost <= r
    # For Smuggler, increase the precision to ϵ
    @timeit to "build cov" cov = [d[s,m] ≤ r+ϵ*abs(r) for s in S, m in M]
    if data.constraint_uncertainty
      for cons in 1:data.ncons
        # The objective coverage matrix is multiplied with the constraint coverage matrix to filter out infeasible solutions
        # (this might be sped-up by avoiding repeating 0 elements and using sparse representations)
        @timeit to "update cov" cov = cov .& a_cov[cons]
      end
    end
    infeasible = false
    # Non-dominated solutions and scenarios are filtered based on the coverage matrix
    @timeit to "dominance solutions" SS = calculate_nondom_solutions(cov,last,M,S)
    @timeit to "dominance scenarios" MM = calculate_nondom_scenarios(cov,M,S)
    # A MIP determining whether it is possible to cover all scenarios by choosing K solutions is solved
    # (notice only non-dominated scenarios and solutions are sent to this MIP)
    @timeit to "solve covering problem" objval,solutions_selected = build_and_solve_covering_MIP(SS,MM,cov)
    # If its optimal value is >= 1 then the MIP is infeasible
    infeasible = objval > 0.9
    if !infeasible
      # If the covering MIP is feasible then a feasible solution with value r can be found and therefore
      # the upper bound is updated
      ub = r
      if last
        # If this is the last iteration of the algorithm a feasible solution is returned
        x = compute_x_values(S, sol_items, solutions_selected)
        break
      end
    else
      # If the covering MIP is infeasible then the lower bound can be increased to r
      if last && objval < Inf
        @error "Shouldn't be here"
        exit()
      elseif last && objval >= Inf
        return [], lb, ub
      else
        lb = r
      end
    end
    # When convergence is achieved last iteration indicator is turned true and one last iteration of BS is done
    if (ub - lb)/abs(ub) ≤ ϵ_BS
      last = true
    end
  end
  # The algorithm returns with an optimal solution and a lower bound on the optimal master value
  # Note that if the algorithm is broken due to time limit restrictions then we return x=[]
  return x, lb, ub
end

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

function update_d(x_sol_array,cc,d,S,M,NUMITER,hh)
  if NUMITER == 1
    for s in S, m in M
      d[s,m] = hh[s]*sum(x_sol_array[s][l]*cc[m,l] for l in N)
    end
  else
    new_costs = [hh[s]*sum(x_sol_array[s][l]*cc[length(M),l] for l in N) for s in S]
    d = [d new_costs]
  end
  return d
end

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

function update_a(x_sol_array,aa,a_cov,S,M,NUMITER)
  if NUMITER == 1
    for cons in data.ncons, s in S, m in M
      a_cov[cons][s,m] = sum(x_sol_array[s][l]*aa[cons][m,l] for l in N) ≤ data.B[cons]
    end
  else
    for cons in data.ncons
      new_cov = [sum(x_sol_array[s][l]*aa[cons][length(M),l] for l in N) ≤ data.B[cons] for s in S]
      a_cov[cons] = [a_cov[cons] new_cov]
    end
  end
  return a_cov
end

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

function calculate_d(x_sol_array,cc,S,M,hh)
  d = Matrix{Float64}(undef,length(S),length(M))
  for s in S, m in M
    d[s,m] = hh[s]*sum(x_sol_array[s][l]*cc[m,l] for l in N)
  end
  return d
end

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

function calculate_a(x_sol_array,aa,S,M)
  a_cov = Vector{Matrix{Bool}}()
  for cons in 1:data.ncons
    push!(a_cov,Matrix{Bool}(undef,length(S),length(M)))
    for s in S, m in M
      a_cov[cons][s,m] = sum(x_sol_array[s][l]*aa[cons][m,l] for l in N) ≤ data.B[cons]
    end
  end
  return a_cov
end

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

function find_new_solution(c_vec,a_vec,UB,fixedzero,fixedone,relax,bound,TIME)
  sol_model = create_model(TIME)
  if relax
    @variable(sol_model, 0<=x[i in N]<=1, container=Array)
  else
    @variable(sol_model, x[i in N], Bin, container=Array)
  end
  add_constraints_X(sol_model,x)
  if data.constraint_uncertainty
    @constraint(sol_model, [cons in 1:data.ncons], sum(a_vec[cons][i]*x[i] for i in N) ≤ data.B[cons] )
  end
  @constraint(sol_model,[i in fixedzero],x[i]==0)
  @constraint(sol_model,[i in fixedone],x[i]==1)
  if bound @constraint(sol_model,sum(x[i]*c_vec[i] for i in N) ≤ UB) end
  @objective(sol_model,Min,sum(x[i]*c_vec[i] for i in N))
  optimize!(sol_model)
  if termination_status(sol_model) == MOI.OPTIMAL
    return value.(x)
  else
    return []
  end
end

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

function relax_and_fix(cc,aa,UB)
  fixedzero = []
  fixedone = []
  time_remaining = TIME_LIM-(time()-STARTTIME)
  @timeit to "relaxation" solution = find_new_solution(cc,aa,UB,fixedzero,fixedone,true,true,time_remaining)
  if size(solution,1) > 0
    for i in N
      if solution[i]>=1-ϵ
        push!(fixedone,i)
      elseif solution[i]<=ϵ
        push!(fixedzero,i)
      end
    end
    @timeit to "integer" solution = find_new_solution(cc,aa,UB,fixedzero,fixedone,false,false,time_remaining)
    if size(solution,1) > 0
      return solution
    else
      return []
    end
  else
    return []
  end
end

#-----------------------------------------------------------------------------------|
" Traslate the enumerated solutions (vector of bool) to solution vectors in 0/1 (x_sol_array). For each solution we also
store the vector of indices of the non-zero variables (sol_items)."
function fill_solutions_bit(solutions, x_sol_array, sol_items)
    S = 1:size(solutions,1)
    for s in S
        sol = zeros(data.n)
        for i in N
            if solutions[s][i]
              sol[i] = 1
            end
        end
        push!(x_sol_array, sol)
        items = Vector{Int}()
        for l in N
            if sol[l] == 1
                push!(items, l)
            end
        end
        push!(sol_items, items)
    end
end

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

" Round the variable values that pass the solver tolerance checks."
function round_solution(sol)
  for i=1:length(sol)
    if sol[i] <= ϵ
      global sol[i] = 0
    elseif sol[i] >= 1-ϵ
      global sol[i]= 1
    end
  end
  return sol
end


#    Copyright (c) 2021 J.A. Duffek
#    Copyright (c) 2000 D.M. Spink
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program; if not, see <http://www.gnu.org/licenses/>.

"""
    nrbtestsrf()::NURBS2D

Construct a simple test surface.

# Output:
- `srf`: NURBS data structure of a surface

# Examples:
```julia
julia> srf = nrbtestsrf()
```
"""
function nrbtestsrf()::NURBS2D
# allocate multi-dimensional array of control points
pnts = zeros(typeof(1.0),(3,5,5));

# define a grid of control points
# in this case a regular grid of u,v points
# pnts[1:3,u,v]

pnts[:,:,1] = [ 0.0  3.0  5.0  8.0 10.0;  # w*x
                0.0  0.0  0.0  0.0  0.0;  # w*y
                2.0  2.0  7.0  7.0  8.0]; # w*z

pnts[:,:,2] = [ 0.0  3.0  5.0  8.0 10.0;
                3.0  3.0  3.0  3.0  3.0;
                0.0  0.0  5.0  5.0  7.0];

pnts[:,:,3] = [ 0.0  3.0  5.0  8.0 10.0;
                5.0  5.0  5.0  5.0  5.0;
                0.0  0.0  5.0  5.0  7.0];

pnts[:,:,4] = [ 0.0  3.0  5.0  8.0 10.0;
                8.0  8.0  8.0  8.0  8.0;
                5.0  5.0  8.0  8.0 10.0];

pnts[:,:,5] = [ 0.0  3.0  5.0  8.0 10.0;
               10.0 10.0 10.0 10.0 10.0;
                5.0  5.0  8.0  8.0 10.0];

# knots
knots = Vector{Vector{typeof(1.0)}}(undef,2);
knots[1] = vec([0 0 0 1/3 2/3 1 1 1]); # knots along u
knots[2] = vec([0 0 0 1/3 2/3 1 1 1]); # knots along v

# make nurbs surface
return nrbmak(pnts,knots);
end # nrbtestsrf


using FreqTables
using Test
using NamedArrays

x = repeat(["a", "b", "c", "d"], outer=[100]);
# Values not in order to test discrepancy between index and levels with CategoricalArray
y = repeat(["D", "C", "A", "B"], inner=[10], outer=[10]);

tab = @inferred freqtable(x)
@test tab == [100, 100, 100, 100]
@test names(tab) == [["a", "b", "c", "d"]]
@test @inferred prop(tab) == [0.25, 0.25, 0.25, 0.25]
tab = @inferred freqtable(y)
@test tab == [100, 100, 100, 100]
@test names(tab) == [["A", "B", "C", "D"]]
tab = @inferred freqtable(x, y)
@test tab == [30 20 20 30;
              30 20 20 30;
              20 30 30 20;
              20 30 30 20]
@test names(tab) == [["a", "b", "c", "d"], ["A", "B", "C", "D"]]

pt = @inferred prop(tab)
@test pt == [0.075  0.05  0.05 0.075;
             0.075  0.05  0.05 0.075;
              0.05 0.075 0.075  0.05;
              0.05 0.075 0.075  0.05]
pt = @inferred prop(tab, margins=2)
@test pt == [0.3 0.2 0.2 0.3;
             0.3 0.2 0.2 0.3;
             0.2 0.3 0.3 0.2;
             0.2 0.3 0.3 0.2]
pt = @inferred prop(tab, margins=1)
@test pt == [0.3 0.2 0.2 0.3;
             0.3 0.2 0.2 0.3;
             0.2 0.3 0.3 0.2;
             0.2 0.3 0.3 0.2]
pt = @inferred prop(tab, margins=(1, 2))
@test pt == [1.0 1.0 1.0 1.0;
             1.0 1.0 1.0 1.0;
             1.0 1.0 1.0 1.0;
             1.0 1.0 1.0 1.0]

tbl = @inferred prop(rand(5, 5, 5, 5), margins=(1, 2))
sumtbl = sum(tbl, dims=(3,4))
@test all(x -> x ≈ 1.0, sumtbl)

@test_throws MethodError prop()
@test_throws MethodError prop(("a","b"))
@test_throws MethodError prop((1, 2))
@test_throws MethodError prop([1,2,3], "a")
@test_throws MethodError prop([1,2,3], 1, 2)
@test_throws ArgumentError prop([1,2,3], margins=2)
@test_throws ArgumentError prop([1,2,3], margins=0)

tab = @inferred freqtable(x, y,
                          subset=1:20,
                          weights=repeat([1, .5], outer=[10]))
@test tab == [2.0 3.0
              1.0 1.5
              3.0 2.0
              1.5 1.0]
@test names(tab) == [["a", "b", "c", "d"], ["C", "D"]]
pt = @inferred prop(tab)
@test pt == [4 6; 2 3; 6 4; 3 2] / 30.0
pt = @inferred prop(tab, margins=2)
@test pt == [8  12; 4   6; 12  8; 6   4] / 30.0
pt = @inferred prop(tab, margins=1)
@test pt == [6 9; 6 9; 9 6; 9 6] / 15.0
pt = @inferred prop(tab, margins=(1, 2))
@test pt == [1.0 1.0; 1.0 1.0; 1.0 1.0; 1.0 1.0]

using CategoricalArrays
cx = CategoricalArray(x)
cy = CategoricalArray(y)

tab = @inferred freqtable(cx)
@test tab == [100, 100, 100, 100]
@test names(tab) == [["a", "b", "c", "d"]]
tab = @inferred freqtable(cy)
@test tab == [100, 100, 100, 100]
@test names(tab) == [["A", "B", "C", "D"]]
tab = @inferred freqtable(cx, cy)
@test tab == [30 20 20 30;
              30 20 20 30;
              20 30 30 20;
              20 30 30 20]
@test names(tab) == [["a", "b", "c", "d"], ["A", "B", "C", "D"]]

tab = @inferred freqtable(cx, cy,
                          subset=1:20,
                          weights=repeat([1, .5], outer=[10]))
@test tab == [0.0 0.0 2.0 3.0
              0.0 0.0 1.0 1.5
              0.0 0.0 3.0 2.0
              0.0 0.0 1.5 1.0]
@test names(tab) == [["a", "b", "c", "d"], ["A", "B", "C", "D"]]

const ≅ = isequal
mx = Array{Union{String, Missing}}(x)
my = Array{Union{String, Missing}}(y)
mx[1] = missing
my[[1, 10, 20, 400]] .= missing

mcx = categorical(mx)
mcy = categorical(my)

tab = @inferred freqtable(mx)
tabc = @inferred freqtable(mcx)
@test tab == tabc == [99, 100, 100, 100, 1]
@test names(tab) ≅ names(tabc) ≅ [["a", "b", "c", "d", missing]]
tab = @inferred freqtable(my)
tabc = @inferred freqtable(mcy)
@test tab == tabc == [100, 99, 99, 98, 4]
@test names(tab) ≅ names(tabc) ≅ [["A", "B", "C", "D", missing]]
tab = @inferred freqtable(mx, my)
tabc = @inferred freqtable(mcx, mcy)
@test tab == tabc == [30 20 20 29 0;
                      30 20 20 29 1;
                      20 30 30 20 0;
                      20 29 29 20 2;
                      0   0  0  0 1]
@test names(tab) ≅ names(tabc) ≅ [["a", "b", "c", "d", missing],
                                  ["A", "B", "C", "D", missing]]

tab = @inferred freqtable(mx, skipmissing=true)
tabc = @inferred freqtable(mcx, skipmissing=true)
@test tab == tabc == [99, 100, 100, 100]
@test names(tab) ≅ names(tabc) ≅ [["a", "b", "c", "d"]]
tab = @inferred freqtable(my, skipmissing=true)
tabc = @inferred freqtable(mcy, skipmissing=true)
@test names(tab) ≅ names(tabc) ≅ [["A", "B", "C", "D"]]
@test tab == tabc == [100, 99, 99, 98]
tab = @inferred freqtable(mx, my, skipmissing=true)
tabc = @inferred freqtable(mcx, mcy, skipmissing=true)
@test tab == tabc == [30 20 20 29;
                      30 20 20 29;
                      20 30 30 20;
                      20 29 29 20]

using DataFrames

for docat in [false, true]
    iris = DataFrame(SepalLength=[4.8, 4.3, 5.8, 5.7, 5.4, 5.7, 5.7, 6.2,
                                  5.1, 5.7, 6.3, 5.8, 7.1, 6.3, 6.5, 7.6, 4.9],
                     SepalWidth=[3, 3, 4, 4.4, 3.9, 3, 2.9, 2.9, 2.5, 2.8,
                                 3.3, 2.7, 3, 2.9, 3, 3, 2.5],
                     Species=["Iris-setosa", "Iris-setosa", "Iris-setosa",
                              "Iris-setosa", "Iris-setosa", "Iris-versicolor",
                              "Iris-versicolor", "Iris-versicolor", "Iris-versicolor",
                              "Iris-versicolor", "Iris-virginica", "Iris-virginica",
                              "Iris-virginica", "Iris-virginica", "Iris-virginica",
                              "Iris-virginica", "Iris-virginica"])
    if docat
        iris.LongSepal = categorical(iris.SepalLength .> 5.0)
    else
        iris.LongSepal = iris.SepalLength .> 5.0
    end
    for cols in ((:Species, :LongSepal), ("Species", "LongSepal"),
                 ("Species", :LongSepal), (:Species, "LongSepal"))
        tab = freqtable(iris, cols...)
        @test tab == [2 3
                    0 5
                    1 6]
        if all(x -> x isa AbstractString, cols)
            @test dimnames(tab) == ["Species", "LongSepal"]
        else
            @test dimnames(tab) == [:Species, :LongSepal]
        end
        @test names(tab) == [["Iris-setosa", "Iris-versicolor", "Iris-virginica"], [false, true]]
        @test (names(tab, 2) isa CategoricalArray) == docat
    end

    tab = freqtable(iris, :Species, :LongSepal, subset=iris.SepalWidth .< 3.8)
    @test tab == [2 0
                  0 5
                  1 6]
    @test names(tab[1:2, :]) == [["Iris-setosa", "Iris-versicolor"], [false, true]]
    @test (names(tab, 2) isa CategoricalArray) == docat
    iris_nt = (Species = iris.Species, LongSepal = iris.LongSepal)
    @test freqtable(iris, :Species, :LongSepal) == freqtable(iris_nt, :Species, :LongSepal)

    @test_throws ArgumentError freqtable(iris)
    @test_throws ArgumentError freqtable(nothing, :Species, :LongSepal)
end

# Issue #5
@test @inferred freqtable([Set(1), Set(2)]) == [1, 1]
@test @inferred freqtable([Set(1), Set(2)], [Set(1), Set(2)]) == [1 0; 0 1]

@test_throws ArgumentError freqtable()
@test_throws ArgumentError freqtable(DataFrame())

# Integer dimension
using NamedArrays

df = DataFrame(A = 101:103, B = ["x","y","y"]);
intft = freqtable(df, :A, :B)
@test names(intft) == [[101,102,103],["x","y"]]
@test intft == [1 0;
                0 1;
                0 1]

@test_throws BoundsError intft[101,"x"]
@test intft[Name(101),"x"] == 1

# proptable
df = DataFrame(x = [1, 2, 1, 2], y = [1, 1, 2, 2], z = ["a", "a", "c", "d"])

for cols in ((:x, :z), ("x", "z"), ("x", :z), (:x, "z"))
    tab = proptable(df, cols...)
    if all(x -> x isa AbstractString, cols)
        @test dimnames(tab) == ["x", "z"]
    else
        @test dimnames(tab) == [:x, :z]
    end
    @test tab == [0.25 0.25 0.0
                  0.25 0.0  0.25]
    @test names(tab) == [[1, 2], ["a", "c", "d"]]
end

tab = proptable(df, :x, :z, margins=1)
@test tab == [0.5 0.5 0.0
              0.5 0.0 0.5]

tab = proptable(df, :x, :y, margins=(1,2))
@test tab == [1.0 1.0
              1.0 1.0]
@test names(tab) == [[1, 2], [1, 2]]

@test_throws ArgumentError proptable(df)
@test_throws ArgumentError proptable(nothing, :x, :y)
@test_throws MethodError proptable(df, :x, :y, 1, 2)


using Test
using MarriageMarkets

include("setup_tests.jl") # multiprocess model computations

@testset "MarriageMarkets" begin

	@testset "Static model" begin

		@testset "Unidimensional model" begin
			include("test_static_unidim.jl")
		end

		@testset "Multidimensional model" begin
			include("test_static_multidim.jl")
		end

	end # static model

	@testset "Search model" begin

		@testset "Basic Shimer-Smith model" begin
			include("test_search_base.jl")
		end

		@testset "Stochastic model" begin
			include("test_search_stoch.jl")
		end

		@testset "Multidimensional types" begin
			include("test_search_multidim.jl")
		end

	end # search model

end # MarriageMarkets


#include("MCMCplot.jl"); traceplot("MCMC_samples_residual_variance.txt","plotly",4); savefig("plot.png");
using DelimitedFiles,Plots,Plots.PlotMeasures,StatsPlots

function traceplot(file,backend="plotly",nplots=4)
#catch errors when no backends are installed
    if backend == "pyplot"
        pyplot(size=(300*nplots,200*nplots))
    elseif backend == "plotly"
        plotly(size=(300*nplots,200*nplots))
    end
    mychain,mylabel=readdlm(file,',',header=true)
    if nplots > length(mylabel)
        nplots=length(mylabel)
    end
    mychain = mychain[:,1:nplots]
    mylabel = mylabel[1:nplots]

    steps = 1:size(mychain,1)
    h1=plot(mychain, layout=(nplots,1),title= reshape(mylabel,1,length(mylabel)),
         label="",title_location=:left,titlefont=12,xlabel="iterations",ylabel = "residual variance",ytickfont=font(8))
    # add horizontal line
    hline!(h1,cumsum(mychain,dims=1)./steps,layout=(nplots,1),label="",color=:red,linewidth=0.1)
    #density plot
    h2=density(mychain,layout= (nplots,1),label="",orientation = :horizontal,xlim =[0,1.5],ylim = [-10,25])
    plot(h1, h2, link = :y, layout = grid(1,2,widths = [0.7,0.3]))
end


using Test
using DomainIntegrals, DomainSets, BasisFunctions, GridArrays, FillArrays

using BasisFunctions.Test: generic_test_discrete_measure

@testset "Discrete measure" begin
    g = ScatteredGrid(rand(2))
    μ = discretemeasure(g)
    generic_test_discrete_measure(μ)
    @test weights(μ) isa FillArrays.Ones
    @test !isnormalized(μ)

    g2 = g×g
    μ = discretemeasure(g2)
    generic_test_discrete_measure(μ)
    @test Matrix(weights(μ))≈ones(2,2)
    @test !isnormalized(μ)

    μ = DiracWeight(rand())
    generic_test_discrete_measure(μ)
    @test weights(μ) isa FillArrays.Ones
    @test isnormalized(μ)

    μ = discretemeasure(PeriodicEquispacedGrid(3,0,1))
    @test isuniform(μ)
    generic_test_discrete_measure(μ)
    μ = discretemeasure(MidpointEquispacedGrid(3,0,1))
    @test isuniform(μ)
    generic_test_discrete_measure(μ)
    μ = discretemeasure(EquispacedGrid(3,0,1))
    @test isuniform(μ)
    generic_test_discrete_measure(μ)
    μ = BasisFunctions.NormalizedDiracComb(EquispacedGrid(3,0,1))
    @test isnormalized(μ)
    generic_test_discrete_measure(μ)
    μ = discretemeasure(MidpointEquispacedGrid(3,0,1),rand(3))
    generic_test_discrete_measure(μ)

    g = FourierGrid(10)
    mg = rescale(FourierGrid(10),-1, 1)
    m = forward_map(mg)
    μ = discretemeasure(mg)
    generic_test_discrete_measure(μ)
    @test forward_map(μ) == m
    @test supermeasure(μ) == discretemeasure(g)
    @test BasisFunctions.apply_map(discretemeasure(g),m) ≈ μ
    @test BasisFunctions.apply_map(μ,inverse(m))≈discretemeasure(g)

    g = PeriodicEquispacedGrid(3,-1,1)×PeriodicEquispacedGrid(4,-1,1)
    μ = discretemeasure(g)
    generic_test_discrete_measure(μ)
    @test μ isa BasisFunctions.DiscreteProductWeight
    @test components(μ) == (component(μ,1),component(μ,2))
end


export CKSVAR_to_TAR, CKSVAR_to_companion, CKSVAR_to_companionFD

function split_by_var(C)
    # Splits 𝐂 into 𝐂₁ corresponding to 𝐘ₜ (all bar final row of 𝐂) and 𝐂₂ corresponding to 𝐘bar*ₜ (final row). I.e. 𝐂 = [𝐂₁; 𝐂₂].
    k = size(C,1)
    return C[1:k-1,:], C[k,:]'
end

function split_by_lag(C, nlags)
    # Given p lags, this splits 𝐂 as per split_by_var and then separates 𝐂ᵢ into components 𝐂ᵢ¹,…,𝐂ᵢᵖ where 𝐂ᵢʲ is the component of 𝐂ᵢ corresponding to the jth lag.
    if ndims(C) == 1
        k = size(C)
        n = 1
    else
        k, n = size(C)
    end
    @assert rem(n, nlags) == 0;
    m = Int(n / nlags)
    C1, C2 = split_by_var(C)
    C1s = []
    C2s = []
    for i in 1:nlags
        C1i = C1[:, (i-1)*m+1:i*m]
        C2i = C2[:, (i-1)*m+1:i*m]
        push!(C1s, C1i)
        push!(C2s, C2i)
    end
    return C1s, C2s
end

function atomize_C(C, nlags)
    # Subdivides components of 𝐂 separated by lag into those corresponding to 𝐘₁ₜ and 𝐘₂ₜ.
    C1s, C2s = split_by_lag(C, nlags)
    C11s = []; C12s = []
    for (i, Ci) in enumerate(C1s)
        push!(C11s, C1s[i][:,1:end-1])
        push!(C12s, C1s[i][:,end])
    end
    C1s_atom = hcat(C11s, C12s)
    C21s = []; C22s = []
    for (i, Ci) in enumerate(C2s)
        push!(C21s, C2s[i][:,1:end-1])
        push!(C22s, C2s[i][:,end])
    end
    C2s_atom = hcat(C21s, C22s)
    return C1s_atom, C2s_atom
end

"""
    CKSVAR_to_TAR(C, Cstar, βtilde, nlags)

Returns pair `(Σ, X)` where `Σ` is a set of matrices corresponding to the CKSVAR model (with censored variable in levels) in TVAR form, and `X` is a set encoding state-space constraints.
"""
function CKSVAR_to_TAR(C, Cstar, βtilde, nlags)
    k,= size(C)
    C1s, C2s = atomize_C(C, nlags)
    Cs = reshape(hcat(C1s, C2s), (nlags, 2, 2))  # index notation is Cs[lag, rhsvariable, lhsvariable]
    Cstar1s, Cstar2s = split_by_lag(Cstar, nlags)
    Cstars = [Cstar1s Cstar2s]  # index notation is Cstars[lag, lhsvariable]

    vals = perms(nlags)

    Σ = []
    Es = []; Ds = []; Xs = []
    for val in vals
        CTAR = zeros(k*nlags, k*nlags)
        CTAR[(k+1):end, 1:k*(nlags-1)] = I(k*(nlags-1))  # I
        CTAR[k+1:2k-1, k] = -βtilde * indicator(val[1], 0, indregion=:below)  # -β𝟙{Y₂ₜ₋₁ < 0}

        CTAR2s = []
        CTAR2stars = []
        # lag = 1
        push!(CTAR2s, Cs[1, 1, 2])  # C₂₁₁
        push!(CTAR2stars, Cs[1, 2, 2] * indicator(val[1], 0) + Cstars[1, 2] - [Cs[1, 1, 2] * βtilde * indicator(val[1], 0, indregion=:below)])  # C₂₁* + C₂₂₁𝟙{Y₂ₜ₋ᵢ ≥ 0} - C₂₁₁β𝟙{Y₂ₜ₋ᵢ < 0}
        for lag in 2:nlags
            push!(CTAR2s, Cs[lag, 1, 2])  # C₂₁ᵢ
            push!(CTAR2stars, Cstars[lag, 2] + Cs[lag, 2, 2] * indicator(val[lag], 0))  # C₂ᵢ* + C₂₂ᵢ𝟙{Y₂ₜ₋ᵢ ≥ 0}
        end

        CTAR1s = []
        CTAR1stars = []
        # lag = 1
        push!(CTAR1s, Cs[1, 1, 1])  # C₁₁₁
        push!(CTAR1stars, Cstars[1, 1] + Cs[1, 2, 1] * indicator(val[1], 0) - Cs[1, 1, 1] * βtilde * indicator(val[1], 0 , indregion=:below))  # C₁₁* + C₁₂₁𝟙{Y₂ₜ₋ᵢ* ≥ 0} - C₁₁₁β𝟙{Y₂ₜ₋ᵢ* < 0}
        for lag in 2:nlags
            push!(CTAR1s, Cs[lag,1,1])  # C₁₁ᵢ
            push!(CTAR1stars, Cstars[lag, 1] + Cs[lag, 2, 1] * indicator(val[lag], 0))  # C₁ᵢ* + C₁₂ᵢ𝟙{Y₂ₜ₋ᵢ* ≥ 0}
        end

        for i in 1:nlags
            CTAR[1:k, k*(i-1)+1:i*k] = hcat(vcat(CTAR1s[i], CTAR2s[i]), vcat(CTAR1stars[i], CTAR2stars[i]))
        end
        push!(Σ, CTAR)

        E = zeros(nlags, k*nlags)
        D = zeros(1, k*nlags)
        for i in 1:nlags
            E[i, k*i] = val[i]
        end
        push!(Es, E)
        push!(Ds, D)
    end
    for i in 1:length(vals)
        X = Vector([Es[i], Ds[i]])
        push!(Xs, X)
    end

    Vector{Array{Float64, 2}}(Σ), Vector(Xs)
    # NOTE the ordering of the matrices in Σ
end

"""
    CKSVAR_to_companion(C, Cstar, βtilde, nlags)

Converts CKSVAR model (with censored variable in levels) into companion form.
"""
function CKSVAR_to_companion(C, Cstar, βtilde, nlags)
    k = size(C,1)
    C1s, C2s = atomize_C(C, nlags)
    Cs = reshape(hcat(C1s, C2s), (nlags, 2, 2))  # index notation is Cs[lag, rhsvariable, lhsvariable]
    Cstar1s, Cstar2s = split_by_lag(Cstar, nlags)
    Cstars = [Cstar1s Cstar2s]  # index notation is Cstars[lag, lhsvariable]
    K = k + (nlags-1) * (k+1)
    Σ = []
    for val in [1, -1]
        A = zeros(K, K)
        A[1:k-1, 1:k-1] = Cs[1, 1, 1]  # C₁₁₁
        A[k, 1:k-1] = Cs[1, 1, 2]  # C₂₁₁
        A[1:k-1, k] = Cstars[1, 1] + Cs[1, 2, 1]*indicator(val, 0, indregion=:above) - Cs[1, 1, 1]*βtilde*indicator(val, 0, indregion=:below)  # C₁₁* + C₁₂₁𝟙{Y₁ₜ₋₁* ≥ 0} - C₁₁₁β𝟙{Y₁ₜ₋₁* < 0}
        A[k,k] = Cstars[1, 2][1] + Cs[1, 2, 2][1]*indicator(val, 0, indregion=:above) - Cs[1, 1, 2]*βtilde*indicator(val, 0, indregion=:below)  # C₂₁* + C₂₂₁𝟙{Y₂ₜ₋₁ ≥ 0} - C₂₁₁𝟙{Y₁ₜ₋₁* < 0}
        for lag in 2:nlags
            C_lag = vcat(hcat(Cs[lag, 1, 1], Cs[lag, 2, 1], Cstars[lag, 1]), hcat(Cs[lag, 1, 2], Cs[lag, 2, 2], Cstars[lag, 2]))  # [C₁₁ᵢ C₁₂ᵢ C₁ᵢ*; C₂₁ᵢ C₂₂ᵢ C₂ᵢ*] (ith lag)
            J = (lag - 1)*(k + 1)
            A[1:k, J:J+k] = C_lag
        end
        A[k+1:2k-1, 1:k-1] = I(k-1)  # I
        A[k+1:2k-1, k] = -βtilde * indicator(val, 0, indregion=:below)  # -β𝟙{Y₂ₜ₋₁ < 0}
        A[2k, k] = indicator(val, 0, indregion=:above)  # 𝟙{Y₂ₜ₋₁ ≥ 0}
        A[2k+1, k] = 1
        if nlags > 2
            A[2k+2:K, k+1:K-k-1] = I(K - 2k - 1)  # I
        end
        push!(Σ, A)
    end
    Vector{Array{Float64, 2}}(Σ)
end

"""
    CKSVAR_to_companionFD(F, Fstar, βtilde, nlags; diff=true)

Converts CKSVAR model estimated with the censored variable entering in first differences into companion form. The default setting `diff=true` returns the companion form with the censored variable entering in first differences. To retrieve the companion form with censored variable entering in levels, set `diff=false`.
"""
function CKSVAR_to_companionFD(F, Fstar, βtilde, nlags; diff = true)
    k = size(F, 1)
    F1s, F2s = atomize_C(F, nlags)
    Fs = reshape(hcat(F1s, F2s), (nlags, 2, 2))  # index notation is Fs[lag, rhsvariable, lhsvariable]
    Fstar1s, Fstar2s = split_by_lag(Fstar, nlags-1)
    Fstars = [Fstar1s Fstar2s]  # index notation is Fstars[lag, lhsvariable]
    if diff == false
        K = k + (nlags-1) * (k+1)
        Σ = []
        for val in [1, -1]
            A = zeros(K, K)
            A[1:k-1, 1:k-1] = Fs[1, 1, 1]  # C₁₁₁
            A[k, 1:k-1] = Fs[1, 1, 2]  # C₂₁₁
            A[1:k-1, k] = Fstars[1, 1] + βtilde + (Fs[1, 2, 1] - βtilde)*indicator(val, 0, indregion=:above) - Fs[1, 1, 1]*βtilde*indicator(val, 0, indregion=:below)  # F₁₁* + β + (F₁₂₁ - β)𝟙{Y₁ₜ₋₁* ≥ 0} - C₁₁₁β𝟙{Y₁ₜ₋₁* < 0}
            A[k,k] = 1. + Fstars[1, 2][1] + Fs[1, 2, 2][1]*indicator(val, 0, indregion=:above) - Fs[1, 1, 2]*βtilde*indicator(val, 0, indregion=:below)  # 1 + F₂₁* + F₂₂₁𝟙{Y₂ₜ₋₁ ≥ 0} - C₂₁₁β𝟙{Y₁ₜ₋₁* < 0}
            for lag in 2:nlags-1  # i = 2,…,p-1
                F_lag = vcat(hcat(Fs[lag,1,1], Fs[lag,2,1] - Fs[lag-1,2,1], Fstars[lag, 1] - Fstars[lag-1,1]), hcat(Fs[lag, 1, 2], Fs[lag,2,2] - Fs[lag-1,2,2], Fstars[lag, 2] - Fstars[lag-1,2]))
                J = (lag - 1)*(k + 1)
                A[1:k, J:J+k] = F_lag
            end
            F_lag = vcat(hcat(Fs[nlags,1,1], -Fs[nlags-1,2,1], -Fstars[nlags-1,1]), hcat(Fs[nlags, 1, 2], -Fs[nlags-1,2,2], -Fstars[nlags-1,2]))  # lag = p
            J = (nlags - 1)*(k + 1)
            A[1:k, J:J+k] = F_lag
            A[k+1:2k-1, 1:k-1] = I + zeros(k-1, k-1)  # I
            A[k+1:2k-1, k] = -βtilde * indicator(val, 0, indregion=:below)  # -β𝟙{Y₂ₜ₋₁ < 0}
            A[2k, k] = indicator(val, 0, indregion=:above)  # 𝟙{Y₂ₜ₋₁ ≥ 0}
            A[2k+1, k] = 1
            if nlags > 2
                A[2k+2:K, k+1:K-k-1] = I(K-2k-1)  # I
            end
            push!(Σ, A)
        end
    end
        if diff == true
            K = nlags * k
            Σ = []
            for δ in [1., 0.]
                A = zeros(K, K)
                A[1:k-1, 1:k-1] = Fs[1, 1, 1]  # C₁₁₁
                A[k, 1:k-1] = Fs[1, 1, 2]  # C₂₁₁
                A[1:k-1, k] = Fstars[1, 1] + Fs[1, 1, 1]*βtilde*(δ - 1) + Fs[1, 2, 1]*δ  # F₁₁* + C₁₁₁β(δ - 1) + F₁₂₁δ
                A[k,k] = 1. + Fstars[1, 2][1] + Fs[1, 1, 2]*βtilde*(δ - 1) + Fs[1, 2, 2][1]*δ  # F₂₁* + C₂₁₁β(δ - 1) + F₂₂₁δ
                for lag in 2:nlags-1  # i = 2,…,p-1
                    F_lag = vcat(hcat(Fs[lag,1,1], Fs[lag,2,1], Fstars[lag, 1]), hcat(Fs[lag, 1, 2], Fs[lag,2,2], Fstars[lag, 2]))
                    J = (lag - 1)*(k + 1)
                    A[1:k, J:J+k] = F_lag
                end
                F_lag = vcat(Fs[nlags,1,1], Fs[nlags, 1, 2])  # lag = p
                J = (nlags - 1)*(k + 1)
                A[1:k, J:J+k-2] = F_lag
                A[k+1:2k-1, 1:k-1] = I + zeros(k-1, k-1)  # I
                A[k+1:2k-1, k] = βtilde*(δ - 1)  # β(δ - 1)
                A[2k, k] = δ  # δ
                if nlags > 2
                    A[2k+1, k] = 1
                    A[2k+2:K, k+1:K-k-1] = I(K-2k-1)  # I
                end
                push!(Σ, A)
            end
        end
    Vector{Array{Float64, 2}}(Σ)
end


# basis.jl
#
# Types and methods for defining and working with basis vectors of vector spaces

# BasisVector and Basis
#-----------------------
# immutable type for specifying basis vectors: every vector space should have a canonical basis, where the basis
# vectors can be specified using an identifier of some custom type T (e.g. user defined)
immutable BasisVector{S<:VectorSpace,T}
    space::S
    identifier::T
end

space(b::BasisVector) = b.space
Base.in(b::BasisVector,V::VectorSpace) = space(b) == V

Base.show(io::IO, b::BasisVector) = print(io, "BasisVector($(b.space),$(b.identifier))")

# a basis is just an iterator over the canonical basis vectors of the space
immutable Basis{S<:VectorSpace}
    space::S
end

space(B::Basis) = B.space
basis(V::VectorSpace) = Basis(V) # lowercase convenience method

Base.show(io::IO, B::Basis) = print(io, "Basis($(B.space))")


using PauliStrings
using Test

@testset "PauliString - Structure & Basic Arithmetic" begin

	p_string = PauliString(4,[1,3],[2,3],3.3)

	#
	#	check construction
	#
	p_string.N == 4
	p_string.sites == [1,3]
	p_string.baseIdx == [2,3]
	@test coef(p_string) == 3.3

	#
	#	check set_coef!
	#
	set_coef!(p_string, 2.0)
	@test coef(p_string) == 2.0

	#
	#	check multiplication
	#
	p_start = deepcopy(p_string)
	p_string = 3.0*p_string
	@test coef(p_string) == 6.0

	p_string = p_string*2.0
	@test coef(p_string) == 12.0

	#
	#	check equal and parallel:
	#	parallel means same site same baseIdx different coef
	#
	p_string∥p_start == true
	@test iszero(p_string == p_start)

	# check addition
	p_start += p_start # 4
	p_start *= 3.0
	@test isone(p_string == p_start)

	#
	#	check conj
	#
	p_string *= 1.0im
	p_string_conj = conj(p_string)

	@test coef(p_string_conj) == -12.0im

	conj!(p_string)
	@test isone(p_string_conj == p_string)

end

@testset "PauliStrings - Convertion to SparseArray" begin

	@show nothing

end

@testset "PauliStrings - Commutation & Multiplication" begin

	# test multiplication
	p1 = PauliString(1,[1],[2],1.0)
	p2 = PauliString(1,[1],[3],1.0)
	p3 = PauliString(1,[1],[1],1.0im)
	@test p3 ==  p1*p2
	@test p3 == conj(p2*p1)

	p1 = PauliString(2,[1,2],[2,1],1.0)
	p2 = PauliString(2,[1],[3],1.0)
	p3 = PauliString(2,[1,2],[1,1],1.0im)

	@test p3 ==  p1*p2
	@test p3 == conj(p2*p1)

	#commutation
	@test commutePauli(p1,p2) == 2*p3

	p1 = PauliString(2,[1,2],[2,1],1.0)
	p2 = PauliString(2,[1,2],[3,2],1.0)
	p3 = PauliString(2,[1,2],[1,3],-1.0)

	@test p3 ==  p1*p2
	@test p3 == p2*p1

	# commutation
	@test commutePauli(p1,p2) == nothing

end


#=
  http://cplint.eu/example/inference/viral.swinb
  """
  Viral Marketing

  A firm is interested in marketing a new product to its customers. These are connected in a social network
  that is known to the firm: the network represents the trust relationships between customers. The firm has
  decided to adopt a marketing strategy that involves giving the product for free to a number of its
  customers, in the hope that these influence the other customers and entice them to buy the product.
  The firm wants to choose the customers to which marketing is applied so that its return is maximized.
  This involves computing the probability that the non-marketed customers will acquire the product given the
  action to the marketed customers.

  This viral marketing scenario is inspired by [1].

  We can model this domain with an LPAD where the predicate trust/2 encodes the links between customers in the
  social network and predicate has/1 is true of customers that possess the product, either received as a gift
  or bought. Predicate trust/2 is defined by a number of certain facts, while predicate has/1 is defined by
  two rules, one expressing the prior probability of a customer to buy the product and one expressing the
  fact that if a trusted customer has the product, then there is a certain probability that the trusting
  customer buys the product.
  """

  Reference G. V. den Broeck, I. Thon, M. van Otterlo, L. D. Raedt,
  Dtproblog: A decision-theoretic probabilistic prolog, in: M. Fox, D. Poole (Eds.),
  24th AAAI Conference on Artificial Intelligence, AAAI’10, Atlanta, Georgia, USA,
  July 11-15, 2010, AAAI Press, 2010, pp. 1217–1222.

  """
  We want to compute the probability that customer 2 buys the product if we perform the action of
  giving the product to customer 3. We need to use causal reasoning so we use the action do(has(3))
  as evidence:
  ?- prob(has(2),do(has(3)),P).

  P = 0.136

  If instead we compute the effect of the action using regular probabilistic inference, we get:
  ?- prob(has(2),has(3),P).
  P = 0.4065135474609725
  So not distinguishing seeing from doing leads to an overly optimistic estimate.

  ""

  Cf ~/blog/viral_marketing.blog
     ~/psi/viral_marketing.psi
     ~/webppl/viral_marketing.wppl

  Distribution of values for has(p*)
  The BLOG model (-n 100000): given has(p3)
                  has(p1): 0.25285605483625506
                  has(p2): 0.32750708813609203
                  has(p3): 1.0
                  has(p4): 0.39824124623193263

  The PSI model: given has(p3)
                 has(p1): 0.254777070063694
                 has(p2): 0.331210191082803,
                 has(p3): 1,
                 has(p4): 0.4

  This Turing model,using
     chains = sample(model, MH(), MCMCThreads(), 100_000, num_chains)

  Summary Statistics
  parameters      mean       std   naive_se      mcse           ess      rhat
      Symbol   Float64   Float64    Float64   Float64       Float64   Float64

    has_a[1]    0.2512    0.4337     0.0007    0.0033    11421.0186    1.0006
    has_a[2]    0.3407    0.4739     0.0007    0.0035    10540.0175    1.0002
    has_a[3]    1.0000    0.0045     0.0000    0.0000   122493.4137    1.0000
    has_a[4]    0.3973    0.4893     0.0008    0.0034    12741.2279    1.0003

Quantiles
  parameters      2.5%     25.0%     50.0%     75.0%     97.5%
      Symbol   Float64   Float64   Float64   Float64   Float64

    has_a[1]    0.0000    0.0000    0.0000    1.0000    1.0000
    has_a[2]    0.0000    0.0000    0.0000    1.0000    1.0000
    has_a[3]    1.0000    1.0000    1.0000    1.0000    1.0000
    has_a[4]    0.0000    0.0000    0.0000    1.0000    1.0000

Distributions of variable (num:0)
(0.0, 0.0, 1.0, 0.0)    =>  122148 (0.30537)
(0.0, 0.0, 1.0, 1.0)    =>  81756 (0.20439)
(0.0, 1.0, 1.0, 0.0)    =>  56870 (0.142175)
(0.0, 1.0, 1.0, 1.0)    =>  38753 (0.0968825)
(1.0, 0.0, 1.0, 0.0)    =>  36105 (0.0902625)
(1.0, 1.0, 1.0, 0.0)    =>  25975 (0.0649375)
(1.0, 0.0, 1.0, 1.0)    =>  23722 (0.059305)
(1.0, 1.0, 1.0, 1.0)    =>  14663 (0.0366575)
(0.0, 0.0, 0.0, 1.0)    =>  4 (1.0e-5)
(1.0, 0.0, 0.0, 1.0)    =>  3 (7.5e-6)
(1.0, 0.0, 0.0, 0.0)    =>  1 (2.5e-6)
 21.266611 seconds (229.77 M allocations: 11.048 GiB, 23.33% gc time)

=#

using Turing, StatsPlots, DataFrames
include("jl_utils.jl")

@model function viral_marketing()
    p1,p2,p3,p4 = 1:4
    people = [p1,p2,p3,p4]

    trustMatrix = [[0,0,0,0], # p1
                   [1,0,0,0], # p2
                   [1,1,0,0], # p3
                   [1,0,1,0]  # p4
                  ]

    function trusts(a, b)
        return trustMatrix[a][b]
    end

    has_a = tzeros(4)
    #
    # Note: It seems that one have to create the extra array has_a
    #       to get the proper probabilities for has(p1..p4)
    #       Using
    #         # ...
    #         if s > 0
    #           t ~ flip(0.4)
    #         else
    #           t ~ flip(0.1)
    #         end
    #         return t
    #       don't differentiate between the p's.
    #       And this means that we can convert this to a for loop...
    #
    #=
    function has(p)
        s = sum(Int64[trusts(p,q)==1 && has(q)==1 for q in people])
        if s > 0
            has_a[p] ~ flip(0.4)
        else
            has_a[p] ~ flip(0.1)
        end
        return has_a[p]
    end
    =#

    # This is a little faster than using has(p)
    for p in people
        # s = sum(Int64[trusts(p,q)==1 && has_a[q]==1 for q in people])
        s = count(q->trusts(p,q)==1 && has_a[q]==1, people)
        if s > 0
            has_a[p] ~ flip(0.4)
        else
            has_a[p] ~ flip(0.1)
        end
    end

    # true ~ Dirac(has(p3)==1)
    # return has(p1),has(p2),has(p3),has(p4)
    true ~ Dirac(has_a[p3]==1)
    return has_a[p1],has_a[p2],has_a[p3],has_a[p4]

end

model = viral_marketing()

num_chains = 4

# chains = sample(model, MH(), MCMCThreads(), 100_000, num_chains)
chains = sample(model, MH(), MCMCThreads(), 10_000, num_chains)
# chains = sample(model, MH(), 10_000)

# chains = sample(model, PG(15), MCMCThreads(), 1_000, num_chains)

# chains = sample(model, SMC(1000), MCMCThreads(), 10_000, num_chains)
# chains = sample(model, SMC(1000), 10_000)
# chains = sample(model, IS(), 10_000)
#
display(chains)

genq = generated_quantities(model, chains)
show_var_dist_pct(genq)


function twoSum(l::Vector{T}, target::T) where T
    len = length(l)
    for i in 1:(len - 1)
        for j in (i+1):len
            l[i] + l[j] == target && return [i,j]
        end
    end
    return []
end