text:
#####
##### 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


text:
using Gaugefields
using Wilsonloop

function stoutsmearing(NX,NY,NZ,NT,NC)
    Nwing = 1
    Dim = 4
    U = Initialize_Gaugefields(NC,Nwing,NX,NY,NZ,NT,condition = "hot")
    L = [NX,NY,NZ,NT]

    comb = 6
    factor = 1/(comb*U[1].NV*U[1].NC)
    temp1 = similar(U[1])
    temp2 = similar(U[1])
    plaq_t = calculate_Plaquette(U,temp1,temp2)*factor
    println(" plaq_t = $plaq_t")

    nn = CovNeuralnet()
    ρ = [0.1]
    layername = ["plaquette"]
    st = STOUT_Layer(layername,ρ,L)
    push!(nn,st)

    show(nn)

    @time Uout,Uout_multi,_ = calc_smearedU(U,nn)
    plaq_t = calculate_Plaquette(Uout,temp1,temp2)*factor
    println("plaq_t = $plaq_t")

    gauge_action = GaugeAction(U)
    plaqloop = make_loops_fromname("plaquette")
    append!(plaqloop,plaqloop')
    β = 5.7/2
    push!(gauge_action,β,plaqloop)

    μ = 1
    dSdUμ = similar(U)
    for μ=1:Dim
        dSdUμ[μ] = calc_dSdUμ(gauge_action,μ,U)
    end

    @time dSdUbareμ = back_prop(dSdUμ,nn,Uout_multi,U)

end

NX = 4
NY = 4
NZ = 4
NT = 4
NC = 3
stoutsmearing(NX,NY,NZ,NT,NC)

ext: jl


text:
function test_cat()
  ϵ = eps(Float64)
  rtol = sqrt(ϵ)
  @testset ExtendedTestSet "Concatenation" begin
    A = simple_sparse_matrix(Float64, 100, 100)
    B = simple_matrix(Float64, 100, 10)
    C = simple_matrix(Float64, 100, 90)
    D = [A B C]
    Ao = LinearOperator(A)
    Bo = LinearOperator(B)
    Co = LinearOperator(C)
    Do = [Ao Bo Co]
    Do2 = LinearOperator(D)
    rhs = simple_vector(Float64, 200)

    @test(norm(Do * rhs - D * rhs) <= rtol * norm(D * rhs))  # throws InexactError()
    @test(norm(Do2 * rhs - D * rhs) <= rtol * norm(D * rhs))

    @test_throws LinearOperatorException [LinearOperator(ones(5, 5)) opEye(3)]

    # test vcat
    A = simple_sparse_matrix(Float64, 100, 100)
    B = simple_matrix(Float64, 10, 100)
    C = simple_matrix(Float64, 90, 100)
    D = [A; B; C]
    Ao = LinearOperator(A)
    Bo = LinearOperator(B)
    Co = LinearOperator(C)
    Do = [Ao; Bo; Co]
    Do2 = LinearOperator(D)
    rhs = simple_vector(Float64, 100)

    @test(norm(Do * rhs - D * rhs) <= rtol * norm(D * rhs))
    @test(norm(Do2 * rhs - D * rhs) <= rtol * norm(D * rhs))

    @test_throws LinearOperatorException [LinearOperator(ones(5, 5)); opEye(3)]
    K = [opEye(2) opZeros(2, 3); opZeros(3, 2) opEye(3)]
    v = simple_vector(Float64, 5)
    @test all(v .== K * v)

    K = [opEye(2); sparse(1.0I, 2, 2)]
    v = simple_vector(Float64, 2)
    @test all(K * v .== [v; v])
  end
end

test_cat()


text:
using Test, AlgoritmiStruttureDati.Laboratorio3
a = [0, 1, 2, 2, 4, 5, 6, 6, 6, 9]

@testset "freqTwo.jl" begin
    @test freqTwo(a, 0, 9) == 0
    @test freqTwo(a, 2, 2) == 1
    @test freqTwo(a, 6, 6) == 2
end


text:
# Partition the first M rows of v into M/N sets of length N
function partition(v::Array{Float64,2}, M::Int64, N::Int64)
    npartitions = int(floor(M/N));
    vparts = zeros(N,5,npartitions);   # preallocate results
    steps = [1:npartitions];
    for i in steps
        binstart = N*(i-1) + 1;
        binend = N*i;
        vparts[:,:,i] = v[binstart:binend,:];
    end
    return vparts;
end


text:
# Custom plugin for Documentation workflow
struct CI <: PkgTemplates.FilePlugin
    ci_yml::String

    # Default constructor; don't use this
    CI(ci_yml::AbstractString) = new(ci_yml)
end

# Construtor with kwargs; use this. Don't use the default value
CI(; ci_yml::AbstractString = "") = CI(ci_yml)

# Specifying source template and destination
PkgTemplates.source(p::CI) = p.ci_yml
PkgTemplates.destination(::CI) = "./.github/workflows/ci.yml"

# view definition
PkgTemplates.view(::CI, ::Template, pkg::AbstractString) = Dict(
    "PKG" => pkg
)


text:
# coding: utf-8

using DataFrames

doc=
"""
強化学習 & 逆強化学習
2017/07/18
"""

# ランダムな状態遷移
function get_random_state(num_state)
  return rand(1:num_state)
end

# ランダムな行動選択
function get_random_action(num_action)
  return rand(1:num_action)
end

# 可能な状態遷移
function get_possible_actions(D, state, num_state)
  state < 1 && throw("invaid state: $state")
  D == nothing && return get_random_state(num_state)
  return collect(keys(D[state]))
end

# 可能な状態遷移の中で最も高いq値を選択
function get_maxq_from_possible_actions(Q, state, possible_actions)
  return maximum([Q[state][i] for i in possible_actions])
end

# 特徴ベクトルを用いて価値関数を更新
function update_q!(Q, f, w, num_state, num_action)
  for s in collect(1:num_state)
    for a in collect(1:num_action)
      Q[s][a] = dot(w[1]', f[s][a])
    end
  end
end

# 特徴ベクトルを用いて報酬関数を更新
function update_r!(R, f, w, num_state, num_action)
  for s in collect(1:num_state)
    for a in collect(1:num_action)
      R[s][a] = dot(w[1]', f[s][a])
    end
  end
end

# q学習
function q_learn!(Q, R, D, w, f, gamma, learn_count, goal_state, num_state, num_action)
  # 初期状態を定義
  state = get_random_state(num_state)
  # シミュレーター
  for i in collect(1:learn_count)
    # i%100 == 0 && @printf("ite = %d \n", i)
    # 可能な状態遷移
    possible_actions = get_possible_actions(D, state, num_state)
    # ランダムに行動選択
    action = possible_actions[rand(1:end)]
    if D != nothing
      next_state = D[state][action]
    else
      next_state = get_random_state(num_state)
    end

    #update_q!(Q, f, w, num_state, num_action)
    Q[state][action] = dot(w[1]', f[state][action])
    next_possible_actions = get_possible_actions(D, next_state, num_state)
    max_q_next_s = get_maxq_from_possible_actions(Q, next_state, next_possible_actions)
    delta = R[state][action] + (gamma * max_q_next_s - Q[state][action])

    w[1] = w[1]+ 0.1 * (delta * f[state][action])

    state = next_state
    # 終了判定
    if state == goal_state
      # 初期状態を変更して続行
      state = get_random_state(goal_state)
    end
  end
end

# 学習結果の確認
function run_greedy(Q, R, D, num_state, num_action, start_state, goal_state,　end_action, greedy)
  state_log = []
  action_log = []
  i, flg, state = 0, 0, start_state
  while flg != goal_state
    #@printf "current state: %d \n"  state
    # 現在の状態において可能な行動を取得
    possible_actions = get_possible_actions(D, state, num_action)
    # Q-value(s, a)を最大化する行動を選択
    max_Q = 0
    best_action_candidates = []
    for a in possible_actions
      if Q[state][a] > max_Q
        push!(best_action_candidates, a)
        max_Q = Q[state][a]
      elseif Q[state][a] == max_Q
        push!(best_action_candidates, a)
      end
    end
    # 0.1の確率でランダムに行動
    if ((0.9 < rand()) | (length(best_action_candidates) == 0)) & greedy == true
      best_action = possible_actions[rand(1:end)]
    else
    # 最も高いQ値を選択
      best_action = best_action_candidates[end]
    end
    #@printf "-> choose action: %d \n" best_action
    push!(state_log, state)
    push!(action_log, best_action)

    end_action == best_action && return [state_log, action_log]
    # とったactionの値が次の状態の値と等価
    if D != nothing
      state, flg, i = D[state][best_action], D[state][best_action], i + 1
    else
      ns = get_random_state(num_state)
      state, flg, i = ns, ns, i + 1
    end
    # state == goal_state && @printf("state is %d \= goal \n", state)
    # state != goal_state && @printf("state is %d \= not_goal \n", state)
  end
  return [state_log, action_log]
end

# onehot
function onehot(dim, index)
  temp = zeros(dim)
  temp[index] = 1
  return temp
end

# reinforcement learning
function eg_rl_grid()
  # ゴール状態
  goal_state = 6

  # 行動数 左：右：上：下：停止
  num_action = 5

  # 状態数
  num_state = 6

  # 報酬関数
  R = Dict{Any, Any}(1 => Dict(1 => -10, 2 => -10, 3 => -10, 4 => 0, 5 => -10),
                     2 => Dict(1 => -10, 2 => -10, 3 => 100, 4 => 0, 5 => -10),
                     3 => Dict(1 => 0, 2 => -10, 3 => -10, 4 => -10, 5 => -10),
                     4 => Dict(1 => 0, 2 => 0, 3 => 0, 4 => -1, 5 => -10),
                     5 => Dict(1 => -10, 2 => 0, 3 => 0, 4 => 100, 5 => -10),
                     6 => Dict(1 => -10, 2 => -10, 3 => -10, 4 => 0, 5 => 100))

  # 状態遷移関数
  D =  Dict(1 => Dict( 4 => 5, 5 => 1),
  	        2 => Dict( 3 =>　6, 4 => 3, 5 => 2),
  	        3 => Dict( 1 => 4, 5 => 3),
  	        4 => Dict( 1 => 5, 2 => 3, 3 => 2, 5 => 4),
  	        5 => Dict( 2 => 4, 3 => 1, 4 => 6, 5 => 5),
  	        6 => Dict( 4 => 2, 5 => 6))

  # 重みベクトル
  w = [zeros(num_state * num_action)'']
  #w = [zeros(num_state + num_action)'']
  @show w

  # 特徴ベクトル
  f = Dict(x => Dict() for x in collect(1:num_state))
  dim = 1
  for s in collect(1:num_state)
    for a in collect(1:num_action)
      f[s][a] = onehot(num_state * num_action, dim)
      #f[s][a] = vcat(onehot(num_state, s), onehot(num_action, a))
      dim += 1
    end
  end

  # 行動価値関数
  Q = Dict(x => Dict(y => 0.0 for y in collect(1:num_action)) for x in collect(1:num_state))

  # 割引率
  gamma = 0.8

  # 試行回数
  learn_count = 10000

  # q学習
  q_learn!(Q, R, D, w, f, gamma, learn_count, goal_state, num_state, num_action)
  println("R function = ")
  println(R)
  println("Q function = ")
  println(Q)

  # 学習結果
  expert_log = []
  push!(expert_log, run_greedy(Q, R, D, num_state, num_action, 1, 6, nothing, false))
  push!(expert_log, run_greedy(Q, R, D, num_state, num_action, 2, 6, nothing, false))
  push!(expert_log, run_greedy(Q, R, D, num_state, num_action, 3, 6, nothing, false))
  push!(expert_log, run_greedy(Q, R, D, num_state, num_action, 4, 6, nothing, false))
  push!(expert_log, run_greedy(Q, R, D, num_state, num_action, 5, 6, nothing, false))
  push!(expert_log, run_greedy(Q, R, D, num_state, num_action, 6, 6, nothing, false))

  # 最短経路
  println("greedy result = ")
  println(expert_log[1][1], ":", expert_log[1][2])
  println(expert_log[2][1], ":", expert_log[2][2])
  println(expert_log[3][1], ":", expert_log[3][2])
  println(expert_log[4][1], ":", expert_log[4][2])
  println(expert_log[5][1], ":", expert_log[5][2])
  println(expert_log[6][1], ":", expert_log[6][2])
  return expert_log
end

gamma = 0.8

# state
num_state = 6

# action
num_action = 5

# エキスパートの行動ログ
expert_log = eg_rl_grid()
state_log = [x[1] for x in expert_log]
action_log =　[x[2] for x in expert_log]

# 特徴ベクトル
f = Dict(x => Dict() for x in collect(1:num_state))
dim = 1
for s in collect(1:num_state)
  for a in collect(1:num_action)
      f[s][a] = onehot(num_state * num_action, dim)
    dim += 1
  end
end

# u_expert
function get_u_expert(gamma, f, state_logs, action_logs)
  ue = []
  for (sl, al) in zip(state_logs, action_logs)
    ue_ = []
    for (i, (s, a)) in enumerate(zip(sl, al))
      push!(ue_, gamma^i * f[s][a])
    end
    push!(ue, sum(ue_))
  end
  return sum(ue) / length(state_logs)
end

# u
function get_u(gamma, f, state_log, action_log)
  u = []
  for (i, (s, a)) in enumerate(zip(state_log, action_log))
    push!(u, gamma^i * f[s][a])
  end
  return sum(u) / length(u)
end

ue = get_u_expert(gamma, f, state_log, action_log)''
u = [get_u(gamma, f, state_log[3], action_log[3])'']
mu =  [u[end]]
w = [ue - mu[end]]
t = [norm(w[end])]

@show mu
@show w
@show t

# 状態遷移関数
# 状態遷移関数
D =  Dict(1 => Dict( 4 => 5, 5 => 1),
          2 => Dict( 3 =>　6, 4 => 3, 5 => 2),
         　3 => Dict( 1 => 4, 5 => 3),
          4 => Dict( 1 => 5, 2 => 3, 3 => 2, 5 => 4),
          5 => Dict( 2 => 4, 3 => 1, 4 => 6, 5 => 5),
          6 => Dict( 4 => 2, 5 => 6))

# 報酬関数の初期化
R = Dict(x => Dict() for x in collect(1:num_state))
@show R

# 行動価値関数の初期化
Q = Dict(x => Dict(y => 0.0 for y in collect(1:num_action)) for x in collect(1:num_state))

for demo_i in collect(1:100)
  # 閾値
　 @show demo_i t[end]

  t[end] < 0.000000001 && break

  # 報酬関数の初期化
  R = Dict(x => Dict() for x in collect(1:num_state))
  # 報酬関数の更新
  update_r!(R, f, [w[end]], num_state, num_action)

  # 行動価値関数の初期化
　 # Q = Dict(x => Dict(y => 0.0 for y in collect(1:num_action)) for x in collect(1:num_state))

  # Q学習
　　q_learn!(Q, R, D, [w[end]], f, 0.8, 1000, 6, num_state, num_action)

  # best policyrun_greedy
  policy = [run_greedy(Q, R, D, num_state, num_action, ns, 6, nothing, true) for ns in collect(1:num_state)]

  # update
  push!(u, get_u_expert(gamma, f, [x[1] for x in policy], [x[2] for x in policy])'')
  push!(mu, mu[end] .+ ( ((u[end] - mu[end])' * (ue - mu[end])) / ((u[end] - mu[end])' * (u[end] - mu[end])) ) .* (u[end] - mu[end]))
  push!(w, ue - mu[end])
  push!(t, norm(w[end]))

end

println(" ---------- irl resutl ----------")
policy = [run_greedy(Q, R, D, num_state, num_action, ns, 6, nothing, false) for ns in collect(1:num_state)]
println(policy[1])
println(policy[2])
println(policy[3])
println(policy[4])
println(policy[5])
println(policy[6])
println("R function = ")
println(R)
println("Q function = ")
println(Q)


text:
using KLIEPInference
using LinearAlgebra, Random
using Distributions, JLD

Random.seed!(543)

m = parse(Int,ARGS[1])
sgn = parse(Int,ARGS[2])

γx = chain(m; lenC=5, lb=0.2, ub=0.4, sgn=sgn)
γy = copy(γx)

@save "graphs/params_exp3_$(m)_$(sgn).jld" m sgn γx γy


text:
##
# cauchymoment implements the moments of the cauchy transform of a space
#  note that it starts from k==1
##

cauchymoment(S...)=stieltjesmoment(S...)/(-2π*im)

#=
# gives \int_-1^1 x^(k-1)/(z-x) dx
function stieltjeslegendremoment(k::Integer,z)
    if k==1
        return -(log(z-1)-log(z+1))
    elseif k==2
        return -(2-z*log(1+z)+z*log(z-1))
    end

    error("stieltjesmoment not implemented for "*string(S.a)*string(S.b))
end

function stieltjeslegendremoment(k::Integer,z,s::Bool)
    if k==1
        return -(log(1-z)+(s?1:-1)*π*im-log(z+1))
    elseif k==2
        return -(2-z*log(1+z)+z*log(1-z) + (s?1:-1)*π*im*z)
    end

    error("stieltjesmoment not implemented for "*string(S.a)*string(S.b))
end

stieltjesmoment(S::Chebyshev,k::Integer,z)=stieltjeslegendremoment(k,tocanonical(S,z))
function stieltjesmoment(J::Jacobi,k::Integer,z)
    if k==1
        return stieltjeslegendremoment(k,z)
    elseif k==2
        if J.a==J.b
            return (1+J.a)*stieltjeslegendremoment(k,z)
        else
            return (J.b-J.a)/2*stieltjeslegendremoment(1,z) + (2+J.a+J.b)/2*stieltjeslegendremoment(2,z)
        end
    end

    error("Not implemented")
end

function stieltjesmoment(S::PolynomialSpace,k::Integer,z,s::Bool)
    z=tocanonical(S,z)

    if k==1
        return -(log(1-z)+(s?1:-1)*π*im-log(z+1))
    elseif k==2
        return -(2-z*log(1+z)+z*log(1-z) + (s?1:-1)*π*im*z)
    end

    error("stieltjesmoment not implemented for "*string(S.a)*string(S.b))
end

# represents atan(sqrt(z))/sqrt(z)
function sqrtatansqrt(x)
    if  isreal(x) && x ≤ 0
        y=sqrt(-x)
        log((1+y)/(1-y))/(2y)
    else
        sqrt(1/x)*atan(sqrt(x))
    end
end

function sqrtatansqrt(x,s::Bool)
    y=sqrt(-x)
    (log((1+y)/(y-1))-(s?1:-1)*π*im)/(2y)
end
=#

stieltjesmoment(S::WeightedJacobi,n::Int,z) = stieltjesjacobimoment(S.space.a,S.space.b,n,z)
stieltjesmoment(S::WeightedJacobi,z) = stieltjesjacobimoment(S.space.a,S.space.b,z)

stieltjesmoment(S::WeightedJacobiQ,n::Int,z) = stieltjesjacobimoment(S.space.a,S.space.b,n,z)
stieltjesmoment(S::WeightedJacobiQ,z) = stieltjesjacobimoment(S.space.a,S.space.b,z)

stieltjesmoment(S::Jacobi,n::Int,z) = stieltjesjacobimoment(S.a,S.b,n,z)
stieltjesmoment(S::Jacobi,z) = stieltjesjacobimoment(S.a,S.b,z)
stieltjesmoment(S::JacobiQ,n::Int,z) = stieltjesjacobimoment(S.a,S.b,n,z)
stieltjesmoment(S::JacobiQ,z) = stieltjesjacobimoment(S.a,S.b,z)

normalization(n::Int,α::Real,β::Real) = 2^(α+β)*gamma(n+α+1)*gamma(n+β+1)/gamma(2n+α+β+2)
stieltjesjacobimoment(α::Real,β::Real,n::Int,z) = (x = 2./(1-z);normalization(n,α,β)*(-x).^(n+1).*_₂F₁(n+1,n+α+1,2n+α+β+2,x))
stieltjesjacobimoment(α::Real,β::Real,z) = stieltjesjacobimoment(α,β,0,z)

function logjacobimoment(α::Real,β::Real,n::Int,z)
    x = 2./(1-z)
    if n == 0
        2normalization(0,α,β)*(log(z-1)-dualpart(_₂F₁(dual(zero(α)+eps(α+β),one(β)),α+1,α+β+2,x)))
        # For testing purposes only, should be equivalent to above within radius of convergence
        #2normalization(0,α,β)*(log(z-1)-(α+1)/(α+β+2)*x.*_₃F₂(α+2,α+β+3,x))
    else
        -2normalization(n,α,β)/n*(-x).^n.*_₂F₁(n,n+α+1,2n+α+β+2,x)
    end
end
logjacobimoment(α::Real,β::Real,z) = logjacobimoment(α,β,0,z)

stieltjeslegendremoment(n::Int,z) = stieltjesjacobimoment(zero(real(eltype(z))),zero(real(eltype(z))),n,z)
stieltjeslegendremoment(z) = stieltjeslegendremoment(0,z)

logabslegendremoment(z) = real(z*log((z+1)/(z-1))+log(z^2-1)-2)
@vectorize_1arg Number logabslegendremoment

#=
#TODO: this is for x^k but we probably want P_k

#These formulae are from mathematica
function stieltjesjacobimoment(α,β,k::Integer,z)
    if α == β == 0
        return stieltjeslegendremoment(k,z)
    elseif isapprox(α,0) && isapprox(β,0.5)
        if k==1
            return -2*sqrt(2)*(sqrtatansqrt(2/(z-1))-1)
        elseif k==2
            return 2*sqrt(2)*(1/3*(3z-2)-z*sqrtatansqrt(2/(z-1)))
        end
    elseif isapprox(α,0) && isapprox(β,-0.5)
        if k==1
            return 2sqrt(1/(z-1))*atan(sqrt(2/(z-1)))
        elseif k==2
            return z/sqrt(1-z)*log((1+1/sqrt(1-z))*(1-sqrt(2)/sqrt(1-z))/((1-1/sqrt(1-z))*(1+sqrt(2)/sqrt(1-z))))-
                2*sqrt(2)-2z/sqrt(1-z)*asinh(sqrt(-1/z))
        end
    elseif !isapprox(α,0) && isapprox(β,0.)
        return (-1)^k*stieltjesjacobimoment(β,α,k,-z)
    end
    error("stieltjesmoment not implemented for JacobiWeight "*string(α)*string(β))
end

function stieltjesjacobimoment(α,β,k::Integer,z,s::Bool)
    if α == β == 0
        return stieltjeslegendremoment(k,z,s)
    elseif isapprox(α,0) && isapprox(β,0.5)
        if k==1
            return -2*sqrt(2)*(sqrtatansqrt(2/(z-1))-1,!s)
        elseif k==2
            return 2*sqrt(2)*(1/3*(3z-2)-z*sqrtatansqrt(2/(z-1),!s))
        end
    elseif isapprox(α,0) && isapprox(β,-0.5)
        if k==1
            return 2im*(s?-1:1)*sqrt(1/(1-z))*atan((s?-1:1)*im*sqrt(2/(1-z)))
        elseif k==2 && real(z)<0
            return z/sqrt(1-z)*(logabs((1+1/sqrt(1-z))*(1-sqrt(2)/sqrt(1-z))/((1-1/sqrt(1-z))*(1+sqrt(2)/sqrt(1-z))))-π*im*(s?1:-1))-
                            2*sqrt(2)-2z/sqrt(1-z)*asinh(sqrt(-1/z))
        elseif k==2
            return z/sqrt(1-z)*log((1+1/sqrt(1-z))*(1-sqrt(2)/sqrt(1-z))/((1-1/sqrt(1-z))*(1+sqrt(2)/sqrt(1-z))))-
                             2*sqrt(2)-2z/sqrt(1-z)*asinh(im*(s?1:-1)*sqrt(1/z))
        end
    elseif !isapprox(α,0) && isapprox(β,0.)
        return (-1)^k*stieltjesjacobimoment(β,α,k,-z,!s)
    end
    error("stieltjesmoment not implemented for JacobiWeight "*string(α)*string(β))
end

stieltjesmoment{DD}(S::JacobiWeight{Chebyshev{DD},DD},k::Integer,z,s...)=stieltjesjacobimoment(S.α,S.β,k,tocanonical(S,z),s...)

function stieltjesmoment{T,DD}(S::JacobiWeight{Jacobi{T,DD},DD},k::Integer,z,s...)
    z=tocanonical(S,z)
    J=S.space

    if k==1
        return stieltjesjacobimoment(S.α,S.β,k,z,s...)
    elseif k==2
        if J.a==J.b
            return (1+J.a)*stieltjesjacobimoment(S.α,S.β,k,z,s...)
        else
            return (J.a-J.b)/2*stieltjesjacobimoment(S.α,S.β,1,z,s...) + (2+J.a+J.b)/2*stieltjesjacobimoment(S.α,S.β,2,z,s...)
        end
    end
    error("stieltjesmoment not implemented for JacobiWeight "*string(S.α)*string(S.β))
end
=#

function hilbertmoment(S::JacobiWeight,k::Integer,x)
    x=tocanonical(S,x)

    if S.α == 0 && S.β == 0.5
        if k==1
            if isapprox(x,1.0)
                return 2*sqrt(2)/π
            else
                y=sqrt(-2/(x-1))
                return (2*sqrt(2)-sqrt(2)*(log((1+y)/(y-1))/y))/π
            end
        elseif k==2

        end
    elseif S.α == 0.5 && S.β == 0.

    end
    (stieltjesmoment(true,S,k,x)+stieltjesmoment(false,S,k,x))/(2π)::Float64
end


text:
using FastTransforms, Compat
using Compat.Test

import FastTransforms: δ

@testset "Gaunt coefficients" begin
    println("Testing Table 2 of Y.-l. Xu, JCAM 85:53–65, 1997.")
    for (m,n) in ((0,2),(1,2),(1,8),(6,8),(3,18),
                  (10,18),(5,25),(-23,25),(2,40),(-35,40),
                  (28,62),(-42,62),(1,99),(90,99),(10,120),
                  (80,120),(23,150),(88,150))
        @test norm(gaunt(m,n,-m,n)[end]./(big(-1.0)^m/(2n+1))-1, Inf) < 400eps()
    end
    println("Testing Table 3 of Y.-l. Xu, JCAM 85:53–65, 1997.")
    for (m,n,μ,ν) in ((0,1,0,5),(0,5,0,10),(0,9,0,10),(0,10,0,12),
                      (0,11,0,15),(0,12,0,20),(0,20,0,45),(0,40,0,80),
                      (0,45,0,100),(3,5,-3,6),(4,9,-4,15),(-8,18,8,23),
                      (-10,20,10,30),(5,25,-5,45),(15,50,-15,60),(-28,68,28,75),
                      (32,78,-32,88),(45,82,-45,100))
        @test norm(sum(gaunt(m,n,μ,ν))-δ(m,0), Inf) < 15000eps()
    end
end


text:
# SPDX-License-Identifier: MIT

#======================================================================
Interface of a generic FSM.
======================================================================#

"""
    abstract type AbstractFSM{T<:Semiring} end

Abstract base type for all the FSMs.
"""
abstract type AbstractFSM{T<:Semiring} end

"""
    states(fsm)

Return an iterator over the states of the fsm.
"""
states

"""
    arcs(fsm, state)

Return all arcs leaving `state`.
"""
arcs
@deprecate links(fsm, state) arcs(fsm, state)

"""
    length(fsm)

Return the number of states in the FSM.
"""
Base.length(fsm::AbstractFSM) = length(states(fsm))

#======================================================================
Interface of a mutable FSM.
======================================================================#

"""
    abstract type AbstractFSM{T<:Semiring} end

Abstract base type for FSM objects that can be mutated.
"""
abstract type AbstractMutableFSM{T} <: AbstractFSM{T} end

"""
    addstate!(fsm::AbstractMutableFSM{T}, label; initweight = zero(T),
              finalweight = zero(T))

Create a state and add it to `fsm`.
"""
addstate!

"""
    addarc!(fsm::AbstractFSM, src, dest, weight)

Add a directed weighted arc from state `src` to state `dest` with
weight `weight`.
"""
addarc!

@deprecate link!(fsm, src, dest) addarc!(fsm, src, dest)
@deprecate link!(fsm, src, dest, weight) addarc!(fsm, src, dest, weight)

#======================================================================
State of a FSM.
======================================================================#

struct State{T<:Semiring, Ti, Tl}
    id::Ti
    label::Tl
    initweight::T
    finalweight::T
end

Base.iterate(s::State) = (s, nothing)
Base.iterate(s::State, ::Any) = nothing

"""
    isinit(state)

Return `true` is the state is a starting state.
"""
isinit(state::State{T}) where T = state.initweight ≠ zero(T)

"""
    isfinal(state)

Return `true` is the state is a final state.
"""
isfinal(state::State{T}) where T = state.finalweight ≠ zero(T)

#======================================================================
Arc of a FSM.
======================================================================#

struct Arc{T<:Semiring}
    dest::State
    weight::T
end

#======================================================================
Pretty display of a FSM.
======================================================================#

function Base.show(io::IO, fsm::AbstractFSM)
    nstates = 0
    narcs = 0
    for s in states(fsm)
        nstates += 1
        for arc in arcs(fsm, s)
            narcs += 1
        end
    end
    print(io, "$(typeof(fsm)) # states: $nstates # arcs: $narcs")
end

function showlabel(label)
    if typeof(label) <: Tuple
        return "($(join([showlabel(l) for l in label], ",")))"
    end
    "$(label)"
end

function Base.show(io::IO, ::MIME"image/svg+xml", fsm::AbstractFSM)
    dotpath, dotfile = mktemp()
    svgpath, svgfile = mktemp()

    write(dotfile, "Digraph {\n")
    write(dotfile, "rankdir=LR;")

    smap = Dict()
    for (i, s) in enumerate(states(fsm))
        smap[s] = i
    end

    for s in states(fsm)
        name = "$(smap[s])"
        label = showlabel(s.label)
        if isinit(s)
            weight = round(convert(Float64, s.initweight), digits = 3)
            label *= "/$(weight)"
        end
        if isfinal(s)
            weight = round(convert(Float64, s.finalweight), digits = 3)
            label *= "/$(weight)"
        end
        attrs = "shape=" * (isfinal(s) ? "doublecircle" : "circle")
        attrs *= " penwidth=" * (isinit(s) ? "2" : "1")
        attrs *= " label=\"" * label * "\""
        write(dotfile, "$name [ $attrs ];\n")
    end

    for src in states(fsm)
        for arc in arcs(fsm, src)
            weight = round(convert(Float64, arc.weight), digits = 3)
            srcname = "$(smap[src])"
            destname = "$(smap[arc.dest])"
            write(dotfile, "$srcname -> $destname [ label=\"$(weight)\" ];\n")
        end
    end
    write(dotfile, "}\n")
    close(dotfile)
    run(`dot -Tsvg $(dotpath) -o $(svgpath)`)

    xml = read(svgfile, String)
    write(io, xml)

    close(svgfile)

    rm(dotpath)
    rm(svgpath)
end


text:
@testset "Conv-lx" begin
    cd(td)
    # Exception instead of ArgumentError as may fail with system error
    @test_throws Exception J.check_input_rpath("aldjfk")
end

@testset "Conv-html" begin
    @test_throws J.HTMLFunctionError J.convert_html("{{fill bb cc}}", J.PageVars())
    @test_throws J.HTMLFunctionError J.convert_html("{{insert bb cc}}", J.PageVars())
    @test_throws J.HTMLFunctionError J.convert_html("{{href aa}}", J.PageVars())
    @test (@test_logs (:warn, "Unknown dictionary name aa in {{href ...}}. Ignoring") J.convert_html("{{href aa bb}}", J.PageVars())) == "??"
    @test_throws J.HTMLBlockError J.convert_html("{{if asdf}}{{end}}", J.PageVars())
    @test_throws J.HTMLBlockError J.convert_html("{{if asdf}}", J.PageVars())
    @test_throws J.HTMLBlockError J.convert_html("{{isdef asdf}}", J.PageVars())
    @test_throws J.HTMLBlockError J.convert_html("{{ispage asdf}}", J.PageVars())
end

@testset "Conv-md" begin
    s = """
        @def blah
        """
    @test (@test_logs (:warn, "Found delimiters for an @def environment but it didn't have the right @def var = ... format. Verify (ignoring for now).") (s |> jd2html_td)) == ""

    s = """
        Blah
        [^1]: hello
        """ |> jd2html_td
    @test isapproxstr(s, "
Blah
")
end

@testset "Judoc" begin
    cd(td); mkpath("foo"); cd("foo");
    @test_throws ArgumentError serve(single=true)
    cd(td)
end

@testset "RSS" begin
    J.set_var!(J.GLOBAL_PAGE_VARS, "website_descr", "")
    J.RSS_DICT["hello"] = J.RSSItem("","","","","","","",Date(1))
    @test (@test_logs (:warn, """
              I found RSS items but the RSS feed is not properly described:
              at least one of the following variables has not been defined in
              your config.md: `website_title`, `website_descr`, `website_url`.
              The feed will not be (re)generated.""") J.rss_generator()) === nothing
end

@testset "parser-lx" begin
    s = raw"""
        \newcommand{hello}{hello}
        """
    @test_throws J.LxDefError (s |> jd2html)
    s = raw"""
        \foo
        """
    @test_throws J.LxComError (s |> jd2html)
    s = raw"""
        \newcommand{\foo}[2]{hello #1 #2}
        \foo{a} {}
        """
    @test_throws J.LxComError (s |> jd2html)
end

@testset "ocblocks" begin
    s = raw"""
        @@foo
        """
    @test_throws J.OCBlockError (s |> jd2html)
end


text:
@everywhere function fit_ising(T; L=7, nconfs=10_000)
    L2 = L^2
    # * Allocate space for the dataset
    # confs = Array{Float64}(undef, L, L, nconfs)
    c = Matrix{Int32}(undef, L, L)
    confs = Matrix{Int32}[]

    # * Load the datasets from their files
    @inbounds for i in 1:nconfs
        fileising = @sprintf "ising_%.2f_%d.jld2" T i
        JLD2.@load joinpath("data", fileising) c
        push!(confs, copy(c))
    end

    # * Change -1 to zeros
    broadcast!(x -> x == -1 ? 0 : 1, confs, confs)

    # * Flatten the datasets, make them a 1-D array
    newconfs = reshape(confs, (L2, nconfs))

    # for i = axes(newconfs, 2)
    #     for j in eachindex(newconfs[:, i])
    #         if newconfs[j, i] == -1
    #             newconfs[j, i] = 0
    #         end
    #     end
    # end

    # * Train the RBM for the corresponding temperature
    rbm = BernoulliRBM(L2, 32)
    fit(rbm, newconfs, n_epochs = 1000, batch_size = 50, randomize = true, lr = 0.01)

    # * Serialize and save the model
    model_name = @sprintf "model_%.2f.jld2" T
    JLD2.@save joinpath("models", model_name) rbm
end

max_stars_repo_name: edwinb-ai/IsingRBM

max_issues_repo_name: edwinb-ai/IsingRBM

max_forks_repo_name: edwinb-ai/IsingRBM


text:
sz_half = [0.5 0.0; 0.0 -.5]
@testset "MPS" begin
    d, lx = 2, 8

    # A uniformly chosen random state as MPS
    ketstate = complex.(randn(d^lx), randn(d^lx))
    randmps = MPS(lx, d, ketstate)

    # The Bethe chain ground state as MPS
    H = xxz_hamiltonian(lx)
    eheis, vheis = eigsolve(H, 1, :SR)
    mps = MPS(lx, 2, vheis[1])

    @testset "ketstate to MPS to ketstate" begin
        @test norm(ketstate) ≈ norm(randmps)
        @test ketstate ≈ mps2ketstate(randmps)
    end

    @testset "overlap of two MPS" begin
        @test sqrt(overlap(mps, mps)) ≈ norm(mps)
        @test overlap(mps, randmps) ≈ conj.(overlap(randmps, mps))
    end

    @testset "operator measurements" begin
        sz = spinoperators(1/2)[1]
        szdata = measure(mps, sz)
        @test all([isapprox(szdata[i], 0.0, atol=1.e-12) for i=1:lx])

        mpo = xxz_mpo(Float64, lx, 2)
        @test measure(mps, mpo)[1] ≈ eheis[1]
    end

    @testset "MPS center manipulations" begin
        ### TODO: write more tests for here!
        center_at!(randmps, lx-1)
        center_at!(randmps, lx-2)
        center_at!(randmps, 1)
        center_at!(randmps, lx)
        @test mps2ketstate(randmps) ≈ ketstate

        center_at!(randmps, lx)
        @test mps2ketstate(randmps) ≈ ketstate
    end

end

@testset "apply" begin
    @testset "twosite" begin
        lx, l = 8, 3

        H = xxz_hamiltonian(lx)
        eheis, vheis = eigsolve(H, 1, :SR)
        mps = MPS(lx, 2, vheis[1])

        mpscopy = deepcopy(mps)
        center_at!(mps, l)
        sz = spinoperators(1/2)[1]
        M = reshape(kron(sz, sz), 2, 2, 2, 2)
        apply!(mps, M, l)

        #@test mps_dims_are_consistent(mps)
        @test overlap(mpscopy, mps) ≈ measure(mpscopy, sz, sz, l, l+1)
    end
end

max_stars_repo_name: amiragha/mps

max_issues_repo_name: amiragha/mps

max_forks_repo_name: amiragha/mps


text:
### A Pluto.jl notebook ###
# v0.14.4

using Markdown
using InteractiveUtils

# ╔═╡ 9770ad46-aa7b-11eb-3d74-3f168a479469
begin

	# what is my current DIR -
	_PATH_TO_ROOT = pwd()
	_PATH_TO_CONFIG = joinpath(_PATH_TO_ROOT, "config")
	_PATH_TO_DATA = joinpath(_PATH_TO_ROOT, "data")
	_PATH_TO_WEEKLY = joinpath(_PATH_TO_ROOT, "weekly_adjusted")

	# activate me -
	import Pkg
	Pkg.activate(_PATH_TO_ROOT)
	Pkg.instantiate()

	# load external packages that I depend upon -
	using DataFrames
	using CSV
	using TOML
	using JSON
	using Dates
	using PooksoftAlphaVantageDataStore

	# should the download execute?
	_SHOULD_DOWNLOAD_EXECUTE = false

	# this block returns ...
	nothing
end

# ╔═╡ 08e679b8-2e69-42b2-8531-ca95ab8bf6b5
html"""<style>
main {
    max-width: 1200px;
  	width: 70%;
	margin: auto;
}
"""

# ╔═╡ f77c4667-5d7a-4f74-8f06-555424505957
md""" #### Create an Alphavantage USER object

We use the [Alphavantage.co](https://www.alphavantage.co) financial data API for pricing information. To use this API,
we need to create a user object which holds the credentials needed for the Alphavantage API call. We create this object using the `build_api_user_model` method which takes the path to a JSON configuration file e.g. `Configuration.json` of the form:

	{
		\"user_data\":{
			"alpha_vantage_api_email":"<YOUR EMAIL HERE>",
	 		"alpha_vantage_api_key":"<YOUR API KEY HERE>"
		}
	}

"""

# ╔═╡ 5a561360-4a64-4323-89ab-582d6bd45d36
begin

	# setup path to the Configuration.json -
	path_to_config_file = joinpath(_PATH_TO_CONFIG,"Configuration.json")

	# build the api user model -
	user_model_result = build_api_user_model(path_to_config_file)
	user_model = checkresult(user_model_result)

	# return nothing from this block -
	nothing
end

# ╔═╡ 03936477-775a-4246-b196-49d875513f67
md""" #### Make Alphavantage.co API calls

Now that we have the user model, we can make Alphavantage.co API calls to download the price data using the [PooksoftAlphaVantageDataStore](https://github.com/Pooksoft/PooksoftAlphaVantageDataStore.jl) package. The list of stocks that compose the SP500 was prepared by [Rufus Pollock](http://rufuspollock.com/) and the [Open Knowledge Foundation](http://okfn.org/) and can be found [here](https://github.com/datasets/s-and-p-500-companies). We load this list, and make an API call for each element, then save the data to disk:

"""

# ╔═╡ c312c22e-0478-4559-9f81-c639bfa4f22e
begin

	# load the data containing the information about the SP500 companies -
	path_to_SP500_file = joinpath(_PATH_TO_DATA, "constituents.csv")
	df_SP500 = CSV.read(path_to_SP500_file, DataFrame)

	# let's download the weekly adjusted values for the ticker symbols in the constituents.csv file -
	# we wrap the download operation so it doesn't always execute -
	if (_SHOULD_DOWNLOAD_EXECUTE == true)

		# what is the current date?
		timestamp_value = "2021-05-01"

		# setup call -
		data_type = :csv
		outputsize = :full

		# get list of ticker symbols -
		list_of_ticker_symbols = df_SP500[:,1]
		for ticker_symbol in list_of_ticker_symbols

			# make a call to the API -
			df = execute_sts_adjusted_weekly_api_call(user_model, ticker_symbol;
				data_type = data_type, outputsize = outputsize, logger=nothing) |> checkresult

			# dump to the result to disk -
			results_file_name = "WA-$(ticker_symbol)-T-$(timestamp_value).csv"
			path_to_results_file = joinpath(_PATH_TO_WEEKLY,results_file_name)
			CSV.write(path_to_results_file, df)

			# because of API limitations - we need to sleep a little bit here -
			sleep(30)
		end
	end
end

# ╔═╡ Cell order:
# ╟─08e679b8-2e69-42b2-8531-ca95ab8bf6b5
# ╠═9770ad46-aa7b-11eb-3d74-3f168a479469
# ╟─f77c4667-5d7a-4f74-8f06-555424505957
# ╠═5a561360-4a64-4323-89ab-582d6bd45d36
# ╟─03936477-775a-4246-b196-49d875513f67
# ╠═c312c22e-0478-4559-9f81-c639bfa4f22e


text:
url = "http://wiki.puzzlers.org/pub/wordlists/unixdict.txt"
wordlist = open(readlines, download(url))

wsort(word::AbstractString) = join(sort(collect(word)))

function anagram(wordlist::Vector{<:AbstractString})
    dict = Dict{String, Set{String}}()
    for word in wordlist
        sorted = wsort(word)
        push!(get!(dict, sorted, Set{String}()), word)
    end
    wcnt = maximum(length, values(dict))
    return collect(Iterators.filter((y) -> length(y) == wcnt, values(dict)))
end

println.(anagram(wordlist))

max_stars_repo_name: ethansaxenian/RosettaDecode

max_issues_repo_name: ethansaxenian/RosettaDecode

max_forks_repo_name: ethansaxenian/RosettaDecode


text:
@testset "Test recorder" begin
    filename = "test.log"
    try
        # Post event while disabled.
        IS.@record :test InfrastructureSystems.TestEvent("a", 1, 2.0)
        @test !isfile(filename)

        IS.register_recorder!(:test)
        IS.@record :test InfrastructureSystems.TestEvent("a", 1, 2.0)
        IS.unregister_recorder!(:test)

        @test isfile(filename)
        lines = readlines(filename)
        @test length(lines) == 1
        data = JSON2.read(lines[1])
        @test data.name == "TestEvent"
        @test data.val1 == "a"
        @test data.val2 == 1
        @test data.val3 == 2.0

        rm(filename)
        IS.@record :test InfrastructureSystems.TestEvent("a", 1, 2.0)
        @test !isfile(filename)
    finally
        IS.unregister_recorder!(:test)
        isfile(filename) && rm(filename)
    end
end

@testset "Test list_recorder_events" begin
    filename = "test.log"
    try
        IS.register_recorder!(:test)
        for _ in 1:5
            IS.@record :test InfrastructureSystems.TestEvent("a", 1, 2.0)
            IS.@record :test InfrastructureSystems.TestEvent2(5)
        end
        IS.unregister_recorder!(:test)
        @test isfile(filename)

        events = IS.list_recorder_events(InfrastructureSystems.TestEvent, filename)
        @test length(events) == 5
        @test events[1] isa InfrastructureSystems.TestEvent

        IS.show_recorder_events(devnull, InfrastructureSystems.TestEvent, filename)
    finally
        IS.unregister_recorder!(:test)
        isfile(filename) && rm(filename)
    end
end

@testset "Test list_recorder_events filter" begin
    filename = "test.log"
    try
        IS.register_recorder!(:test)
        for i in 1:5
            IS.@record :test InfrastructureSystems.TestEvent("a", i, 2.0)
            IS.@record :test InfrastructureSystems.TestEvent2(3)
        end
        IS.unregister_recorder!(:test)
        @test isfile(filename)

        events = IS.list_recorder_events(
            InfrastructureSystems.TestEvent,
            filename,
            x -> x.val2 == 3,
        )
        @test length(events) == 1

        IS.show_recorder_events(
            devnull,
            InfrastructureSystems.TestEvent,
            filename,
            x -> x.val2 > 0,
        )
    finally
        IS.unregister_recorder!(:test)
        isfile(filename) && rm(filename)
    end
end


text:
using Documenter, WGLMakie

makedocs(; modules=[WGLMakie], format=Documenter.HTML(), pages=["Home" => "index.md"],
         repo="https://github.com/SimonDanisch/WGLMakie.jl/blob/{commit}{path}#L{line}",
         sitename="WGLMakie.jl", authors="Simon Danisch", assets=[])

deploydocs(; repo="github.com/SimonDanisch/WGLMakie.jl")


text:
###########################################################################################
# Operators
###########################################################################################
const DerivStencilScratch = Dict{NTuple{2,Int}, Matrix{Float64}}()

function derivStencil(order::Integer, ptsleft::Integer, ptsright::Integer; gridspacing::Real = 1, kw...)

	invM = get!(DerivStencilScratch, (ptsleft, ptsright)) do
		M = zeros(ptsleft + ptsright + 1, ptsleft + ptsright + 1)
		for (m_ind, m) in enumerate(-ptsleft:ptsright), (p_ind, p) in enumerate(0:(ptsleft + ptsright))
			M[p_ind, m_ind] = m^p/factorial(p)
		end
		inv(M)
	end

	M_order = invM[:, order + 1]
	if (ptsleft == ptsright) && isodd(order)
		# Fix loss of numerical precision
		M_order[ptsleft + 1] = 0
	end

	den = gridspacing^order
	M_order ./= den
	return M_order
end

function derivStencil(order::Integer, N::Integer; kwargs...)
	# given N points decide on the number of points to the left and to the right
	if isodd(N)
		return derivStencil(order, div(N, 2), div(N, 2); kwargs...)
	else
		return derivStencil(order, div(N, 2)-1, div(N, 2); kwargs...)
	end
end

function derivStencil!(v, order::Integer, ptsleft::Integer, ptsright::Integer; gridspacing::Real = 1, kw...)

	invM = get!(DerivStencilScratch, (ptsleft, ptsright)) do
		M = zeros(ptsleft + ptsright + 1, ptsleft + ptsright + 1)
		for (m_ind, m) in enumerate(-ptsleft:ptsright), (p_ind, p) in enumerate(0:(ptsleft + ptsright))
			M[p_ind, m_ind] = m^p/factorial(p)
		end
		inv(M)
	end

	v .= @view invM[:, order + 1]
	if (ptsleft == ptsright) && isodd(order)
		# Fix loss of numerical precision
		v[ptsleft + 1] = 0
	end

	den = gridspacing^order
	v ./= den
	return v
end

function derivStencil!(v, order::Integer, N::Integer; kwargs...)
	# given N points decide on the number of points to the left and to the right
	if isodd(N)
		return derivStencil!(v, order, div(N, 2), div(N, 2); kwargs...)
	else
		return derivStencil!(v, order, div(N, 2)-1, div(N, 2); kwargs...)
	end
end

nextodd(n::Integer) = n + 1 + isodd(n)

function ceilsearch(arr, n)
	# assume sorted
	for elem in arr
		if elem>=n
			return elem
		end
	end
	return n
end

function D(N; stencil_gridpts = Dict(6=>3, 42=>5), #= Dictionary of gridpt => order of derivative upto that gridpt =#
	left_edge_npts = 2, left_edge_ghost = false,
	right_edge_npts = 2, right_edge_ghost = false,
	left_Dirichlet = false,
	right_Dirichlet = false,
	kwargs...)

	@assert(N≥2,"Need at least 2 points to compute the derivative")

	S = zeros(N, N)

	D!(S, N; stencil_gridpts = stencil_gridpts,
		left_edge_npts = left_edge_npts,
		right_edge_npts = right_edge_npts,
		left_Dirichlet = left_Dirichlet,
		right_Dirichlet = right_Dirichlet,
		kwargs...)

	return S
end

function D!(S, N; stencil_gridpts = Dict(6=>3, 42=>5), #= Dictionary of gridpt => order of derivative upto that gridpt =#
	left_edge_npts = 2,
	right_edge_npts = 2,
	left_Dirichlet = false,
	right_Dirichlet = false,
	kwargs...)

	@assert(N≥2,"Need at least 2 points to compute the derivative")
	@assert size(S) == (N, N) "Size of matrix must be $((N, N))"

	gridpt_cutoffs_list = sort(collect(keys(stencil_gridpts)))
	maxorder = maximum(values(stencil_gridpts))

	S .= 0

	# Derivatives on the boundary

	v = @view S[1, 1:left_edge_npts]
	derivStencil!(v, 1, 0, left_edge_npts-1; kwargs...)

	v = @view S[end, (N-right_edge_npts + 1):N]
	derivStencil!(v, 1, right_edge_npts-1, 0; kwargs...)

	for gridpt in 2:N-1

		gridpt_cutoff =  ceilsearch(gridpt_cutoffs_list, min(gridpt, N-gridpt + 1))
		if haskey(stencil_gridpts, gridpt_cutoff)
			npts = stencil_gridpts[gridpt_cutoff]
		else
			npts = nextodd(maxorder)
		end

		diagpt = gridpt
		startpt = max(1, diagpt - div(npts, 2)-(diagpt + div(npts, 2) > N ? diagpt + div(npts, 2) - N_cols : 0 ))
		endpt = min(N, startpt + npts -1)
		npts_left = diagpt - startpt
		npts_right = endpt - diagpt

		v = @view S[gridpt, startpt:endpt]
		derivStencil!(v, 1, npts_left, npts_right; kwargs...)
	end

	if left_Dirichlet
		S[1:2, 1] .= 0
	end
	if right_Dirichlet
		S[end-1:end, end] .= 0
	end

	return S
end

function D²(N; stencil_gridpts = Dict(6=>3, 42=>5),
	left_edge_npts = 3, left_edge_ghost = false, right_edge_npts = 3, right_edge_ghost = false, kwargs...)

	@assert(N≥3,"Need at least 3 points to compute the second derivative")

	N_cols = N
	if left_edge_ghost
		N_cols += 1
	end
	if right_edge_ghost
		N_cols += 1
	end

	S = zeros(N, N_cols)

	gridpt_cutoffs_list = sort(collect(keys(stencil_gridpts)))
	maxorder = maximum(values(stencil_gridpts))

	# Derivatives on the boundary

	if left_edge_ghost
		startpt = 2 - div(left_edge_npts, 2)
		endpt = startpt + left_edge_npts-1
		if startpt >=1
			v = @view S[1, startpt:endpt]
			derivStencil!(v, 2, left_edge_npts; kwargs...)
		else
			v = @view S[1, 1:left_edge_npts]
			derivStencil!(v, 2, 1, left_edge_npts-2; kwargs...)
		end
	else
		v = @view S[1, 1:left_edge_npts]
		derivStencil!(v, 2, 0, left_edge_npts-1; kwargs...)
	end

	if right_edge_ghost
		endpt = N +  div(right_edge_npts, 2) + ( left_edge_ghost ? 1 : 0 )
		startpt = endpt - right_edge_npts + 1
		if endpt<=N_cols
			v = @view S[end, startpt:endpt]
			derivStencil!(v, 2, right_edge_npts; kwargs...)
		else
			v = @view S[end, N_cols-right_edge_npts + 1:N_cols]
			derivStencil!(v, 2, right_edge_npts-2, 1; kwargs...)
		end
	else
		v = @view S[end, N_cols-right_edge_npts + 1:N_cols]
		derivStencil!(v, 2, right_edge_npts-1, 0; kwargs...)
	end

	for gridpt in 2:N-1

		gridpt_cutoff =  ceilsearch(gridpt_cutoffs_list, min(gridpt, N-gridpt + 1))
		if haskey(stencil_gridpts, gridpt_cutoff)
			npts = stencil_gridpts[gridpt_cutoff]
		else
			npts = nextodd(maxorder)
		end

		diagpt = gridpt + ( left_edge_ghost ? 1 : 0 )
		startpt = max(1, diagpt - div(npts, 2)-(diagpt + div(npts, 2) > N_cols ? diagpt + div(npts, 2) - N_cols : 0 ))
		endpt = min(N_cols, startpt + npts -1)
		npts_left = diagpt - startpt
		npts_right = endpt - diagpt

		v = @view S[gridpt, startpt:endpt]
		derivStencil!(v, 2, npts_left, npts_right; kwargs...)
	end

	return S
end

function Dⁿ(order, N; stencil_gridpts = Dict(6=>order + 1, 42=>order + 3),
	left_edge_npts = order + 1, left_edge_ghost = false, right_edge_npts = order + 1, right_edge_ghost = false, kwargs...)

	@assert(N≥order + 1,"Need at least $(order + 1) points to compute the derivative")

	N_cols = N
	if left_edge_ghost
		N_cols += 1
	end
	if right_edge_ghost
		N_cols += 1
	end

	S = zeros(N, N_cols)

	gridpt_cutoffs_list = sort(collect(keys(stencil_gridpts)))
	gridpt_cutoff_maxorder = maximum(keys(stencil_gridpts))
	maxorder = stencil_gridpts[gridpt_cutoff_maxorder]

	# Derivatives on the boundary

	if left_edge_ghost
		startpt = 2 - div(left_edge_npts, 2)
		endpt = startpt + left_edge_npts-1
		if startpt >=1
			v = @view S[1, startpt:endpt]
			derivStencil!(v, order, left_edge_npts; kwargs...)
		else
			v = @view S[1, 1:left_edge_npts]
			derivStencil!(v, order, 1, left_edge_npts-2; kwargs...)
		end
	else
		v = @view S[1, 1:left_edge_npts]
		derivStencil!(v, order, 0, left_edge_npts-1; kwargs...)
	end

	if right_edge_ghost
		endpt = N +  div(right_edge_npts, 2) + ( left_edge_ghost ? 1 : 0 )
		startpt = endpt - right_edge_npts + 1
		if endpt<=N_cols
			v = @view S[end, startpt:endpt]
			derivStencil!(v, order, right_edge_npts; kwargs...)
		else
			v = @view S[end, (N_cols-right_edge_npts + 1):N_cols]
			derivStencil!(v, order, right_edge_npts-2, 1; kwargs...)
		end
	else
		v = @view S[end, (N_cols-right_edge_npts + 1):N_cols]
		derivStencil!(v, order, right_edge_npts-1, 0; kwargs...)
	end

	for gridpt in 2:N-1

		gridpt_cutoff =  ceilsearch(gridpt_cutoffs_list, min(gridpt, N-gridpt + 1))
		if haskey(stencil_gridpts, gridpt_cutoff)
			npts = stencil_gridpts[gridpt_cutoff]
		else
			npts = nextodd(maxorder)
		end

		diagpt = gridpt + ( left_edge_ghost ? 1 : 0 )
		startpt = max(1, diagpt - div(npts, 2)-(diagpt + div(npts, 2) > N_cols ? diagpt + div(npts, 2) - N_cols : 0 ))
		endpt = min(N_cols, startpt + npts -1)
		npts_left = diagpt - startpt
		npts_right = endpt - diagpt

		v = @view S[gridpt, startpt:endpt]
		derivStencil!(v, order, npts_left, npts_right; kwargs...)
	end

	return S
end

function dbydr(dr::AbstractArray; stencil_gridpts = Dict(6=>3, 42=>5), kwargs...)
	D(length(dr); stencil_gridpts = stencil_gridpts, kwargs...) ./ dr
end

function dbydr!(S::AbstractMatrix, dr::AbstractArray; stencil_gridpts = Dict(6=>3, 42=>5), kwargs...)
	D!(S, length(dr); stencil_gridpts = stencil_gridpts, kwargs...)
	S ./= dr
	return S
end

function dbydr(N::Integer, dx::Real, stencil_gridpts = Dict(6=>3, 42=>5), kwargs...)
	D(N; stencil_gridpts = stencil_gridpts, gridspacing = dx)
end

function dbydr!(S::AbstractMatrix, N::Integer, dx::Real, stencil_gridpts = Dict(6=>3, 42=>5); kwargs...)
	D!(S, N; stencil_gridpts = stencil_gridpts, gridspacing = dx)
end

function roll(arr, shift)
	# translation with wrap
	# A′ = T(shift) A
	# translation: A′(i) = A(i-shift)
	# translation with wrap: A′(i + 1) = A(mod(i-shift, N)+1), i = 0...N-1

	N = size(arr, ndims(arr))
	shift = mod(shift, N)
	if shift==0
		return arr
	end

	newarr = similar(arr)

	for i in 0:N-1
		newarr[:,:, i + 1] .= @view arr[:,:, mod(i-shift, N)+1]
	end
	return newarr
end


text:
const FASTDICT = Dict{Symbol,Expr}([
    :(+) => :(Base.FastMath.add_fast),
    :(-) => :(Base.FastMath.sub_fast),
    :(*) => :(Base.FastMath.mul_fast),
    :(/) => :(Base.FastMath.div_fast),
    :(÷) => :(VectorizationBase.vdiv_fast), # VectorizationBase.vdiv == integer, VectorizationBase.vfdiv == float
    :(%) => :(Base.FastMath.rem_fast),
    :abs2 => :(Base.FastMath.abs2_fast),
    # :inv => :(Base.FastMath.inv_fast), # this is slower in most benchmarks
    :hypot => :(Base.FastMath.hypot_fast),
    :max => :(Base.FastMath.max_fast),
    :min => :(Base.FastMath.min_fast),
    :muladd => :(VectorizationBase.vmuladd_fast),
    :fma => :(VectorizationBase.vfma_fast),
    :vfmadd => :(VectorizationBase.vfmadd_fast),
    :vfnmadd => :(VectorizationBase.vfnmadd_fast),
    :vfmsub => :(VectorizationBase.vfmsub_fast),
    :vfnmsub => :(VectorizationBase.vfnmsub_fast),
    :log => :(Base.FastMath.log_fast),
    :log2 => :(Base.FastMath.log2_fast),
    :log10 => :(Base.FastMath.log10_fast)
])

for (op,f) ∈ [
    (:(Base.:-),:vsub),
    (:(Base.FastMath.sub_fast), :vsub_fast),
    # (:(Base.FastMath.abs2_fast),:vabs2_fast),
    (:(Base.inv),:vinv),
    (:(Base.FastMath.inv_fast),:vinv_fast),
    (:(Base.abs),:vabs),
    (:(Base.round),:vround),
    (:(Base.floor),:vfloor),
    (:(Base.ceil),:vceil),
    (:(Base.trunc),:vtrunc),
    (:(Base.unsafe_trunc),:vtrunc),
    (:(Base.signed),:vsigned),
    (:(Base.unsigned),:vunsigned),
    (:(Base.float),:vfloat),
    (:(Base.sqrt),:vsqrt),
    (:(Base.leading_zeros),:vleading_zeros),
    (:(Base.trailing_zeros),:vtrailing_zeros),
    (:(Base.count_ones),:vcount_ones),
]
    @eval begin
        @inline $op(a::AbstractSIMD) = $f(a)
    end
end
@inline Base.:(~)(v::AbstractSIMD{W,T}) where {W,T<:IntegerTypesHW} = v ⊻ vbroadcast(Val(W), -1 % T)
@inline Base.FastMath.abs2_fast(v::AbstractSIMD) = vmul_fast(v,v)

@inline no_promote(a,b) = (a,b)
for (op, f, promote) ∈ [
    (:(Base.:+),:vadd,:promote),
    (:(Base.FastMath.add_fast), :vadd_fast,:promote),
    (:(Base.:-),:vsub,:promote),
    (:(Base.FastMath.sub_fast), :vsub_fast,:promote),
    (:(Base.:*),:vmul,:promote),
    (:(Base.FastMath.mul_fast), :vmul_fast,:promote),
    (:(Base.:/),:vfdiv,:promote_div),
    (:(Base.FastMath.div_fast),:vfdiv_fast,:promote_div),
    (:(Base.:%),:vrem,:promote_div),
    (:(Base.FastMath.rem_fast),:vrem_fast,:promote_div),
    (:(Base.:÷),:vdiv,:promote_div),
    (:(Base.:<<),:vshl,:promote_div),
    (:(Base.:>>),:vashr,:promote_div),
    (:(Base.:>>>),:vlshr,:promote_div),
    (:(Base.:&),:vand,:promote),
    (:(Base.:|),:vor,:promote),
    (:(Base.:⊻),:vxor,:promote),
    (:(Base.max),:vmax,:no_promote),
    (:(Base.min),:vmin,:no_promote),
    (:(Base.FastMath.max_fast),:vmax_fast,:no_promote),
    (:(Base.FastMath.min_fast),:vmin_fast,:no_promote),
    # (:(Base.copysign),:vcopysign,:no_promote),
    (:(Base.:(==)), :veq, :no_promote),
    (:(Base.:(≠)), :vne, :no_promote),
    (:(Base.:(>)), :vgt, :no_promote),
    (:(Base.:(≥)), :vge, :no_promote),
    (:(Base.:(<)), :vlt, :no_promote),
    (:(Base.:(≤)), :vle, :no_promote),
]
    @eval begin
        # @inline $op(a::AbstractSIMD,b::AbstractSIMD) = ((c,d) = $promote(a,b); $f(c,d))
        @inline $op(a::AbstractSIMD,b::AbstractSIMD) = ((c,d) = $promote(a,b); $f(c,d))
        @inline $op(a::NativeTypes,b::AbstractSIMDVector) = ((c,d) = $promote(a,b); $f(c,d))
        @inline $op(a::AbstractSIMDVector,b::NativeTypes) = ((c,d) = $promote(a,b); $f(c,d))
        @inline $op(a::NativeTypes,b::VecUnroll{N,W}) where {N,W} = ((c,d) = $promote(a,b); $f(c,d))
        @inline $op(a::VecUnroll{N,W},b::NativeTypes) where {N,W} = ((c,d) = $promote(a,b); $f(c,d))
        @inline $op(a::IntegerTypesHW,b::VecUnroll{N,W,T,MM{W,X,T}}) where {N,W,T<:IntegerTypesHW,X} = VecUnroll(fmap($op, a, getfield(b,:data)))
        @inline $op(a::VecUnroll{N,W,T,MM{W,X,T}},b::IntegerTypesHW) where {N,W,T<:IntegerTypesHW,X} = VecUnroll(fmap($op, getfield(a,:data), b))
    end
end
for op ∈ [:(Base.:(*)), :(Base.FastMath.mul_fast)]
    @eval begin
        @inline $op(m::AbstractSIMD{W,B1}, v::AbstractSIMD{W,B2}) where {W,B1<:Union{Bool,Bit},B2<:Union{Bool,Bit}} = m & v
        @inline $op(m::AbstractSIMD{W,B}, v::AbstractSIMD{W}) where {W,B<:Union{Bool,Bit}} = ifelse(m, v, zero(v))
        @inline $op(v::AbstractSIMD{W}, m::AbstractSIMD{W,B}) where {W,B<:Union{Bool,Bit}} = ifelse(m, v, zero(v))
    end
end
# copysign needs a heavy hand to avoid ambiguities
@inline Base.copysign(a::VecUnroll,b::AbstractSIMDVector) = VecUnroll(fmap(vcopysign, getfield(a, :data), b))
@inline Base.copysign(a::VecUnroll,b::VecUnroll) = VecUnroll(fmap(vcopysign, getfield(a, :data), getfield(b, :data)))
@inline Base.copysign(a::AbstractSIMDVector,b::VecUnroll) = VecUnroll(fmap(vcopysign, a, getfield(b, :data)))
@inline Base.copysign(a::AbstractSIMDVector,b::AbstractSIMDVector) = vcopysign(a,b)
@inline Base.copysign(a::NativeTypes,b::VecUnroll{N,W}) where {N,W} = VecUnroll(fmap(vcopysign, vbroadcast(Val{W}(), a), getfield(b, :data)))
@inline Base.copysign(a::VecUnroll{N,W},b::Base.HWReal) where {N,W} = VecUnroll(fmap(vcopysign, getfield(a, :data), vbroadcast(Val{W}(), b)))
@inline Base.copysign(a::IntegerTypesHW,b::AbstractSIMDVector) = vcopysign(a,b)
@inline Base.copysign(a::AbstractSIMDVector,b::Base.HWReal) = vcopysign(a,b)
for T ∈ [:Rational, :SignedHW, :Float32, :Float64]
    @eval begin
        @inline function Base.copysign(a::$T, b::AbstractSIMDVector{W,T}) where {W,T <: Union{Float32,Float64,SignedHW}}
            v1, v2 = promote(a, b)
            vcopysign(v1, v2)
        end
        @inline Base.copysign(a::$T, b::AbstractSIMDVector{W,T}) where {W,T <: UnsignedHW} = vbroadcast(Val{W}(), abs(a))
        @inline Base.copysign(a::$T, b::VecUnroll) = VecUnroll(fmap(copysign, a, getfield(b, :data)))
    end
end
for (op, f) ∈ [
    (:(Base.:+),:vadd),
    (:(Base.FastMath.add_fast), :vadd_fast),
    (:(Base.:-),:vsub),
    (:(Base.FastMath.sub_fast), :vsub_fast),
    (:(Base.:*),:vmul),
    (:(Base.FastMath.mul_fast), :vmul_fast)
]
    @eval begin
        @inline $op(m::MM, j::NativeTypes) = ($f(m,j))
        @inline $op(j::NativeTypes, m::MM) = ($f(j,m))
        @inline $op(m::MM, ::StaticInt{N}) where {N} = $f(m, StaticInt{N}())
        @inline $op(::StaticInt{N}, m::MM) where {N} = $f(StaticInt{N}(), m)
    end
end

for (op,c) ∈ [(:(>), :(&)), (:(≥), :(&)), (:(<), :(|)), (:(≤), :(|))]
    @eval begin
        @inline function Base.$op(v1::AbstractSIMD{W,I}, v2::AbstractSIMD{W,U}) where {W,I<:Signed,U<:Unsigned}
            m1 = $op(v1,  zero(I))
            m2 = $op(v1 % U,  v2)
            $c(m1, m2)
        end
    end
end
for (f,vf) ∈ [
    (:convert,:vconvert),(:reinterpret,:vreinterpret),(:trunc,:vtrunc),(:unsafe_trunc,:vtrunc),(:round,:vround),(:floor,:vfloor),(:ceil,:vceil)
]
    @eval begin
        @inline Base.$f(::Type{T}, x::NativeTypes) where {T <: AbstractSIMD} = $vf(T, x)
        @inline Base.$f(::Type{T}, v::AbstractSIMD) where {T <: NativeTypes} = $vf(T, v)
        @inline Base.$f(::Type{T}, v::AbstractSIMD) where {T <: AbstractSIMD} = $vf(T, v)
    end
end
for (f,vf) ∈ [
    (:(Base.rem),:vrem), (:(Base.FastMath.rem_fast),:vrem_fast)
]
    @eval begin
        @inline $f(x::NativeTypes, ::Type{T}) where {T <: AbstractSIMD} = $vf(x, T)
        @inline $f(v::AbstractSIMD, ::Type{T}) where {T <: NativeTypes} = $vf(v, T)
        @inline $f(v::AbstractSIMD, ::Type{T}) where {T <: AbstractSIMD} = $vf(v, T)
    end
end

# These are defined here on `Base` functions to avoid `promote`
@inline function Base.:(<<)(v1::AbstractSIMD{W,T1}, v2::AbstractSIMD{W,T2}) where {W,T1<:SignedHW,T2<:UnsignedHW}
    convert(T1, vshl(convert(T2, v1), v2))
end
@inline function Base.:(<<)(v1::AbstractSIMD{W,T1}, v2::AbstractSIMD{W,T2}) where {W,T1<:UnsignedHW,T2<:SignedHW}
    convert(T1, vshl(convert(T2, v1), v2))
end
@inline function Base.:(>>)(v1::AbstractSIMD{W,T1}, v2::AbstractSIMD{W,T2}) where {W,T1<:SignedHW,T2<:UnsignedHW}
    vashr(v1, (v2 % T1))
end
@inline function Base.:(>>)(v1::AbstractSIMD{W,T1}, v2::AbstractSIMD{W,T2}) where {W,T1<:UnsignedHW,T2<:SignedHW}
    vashr(v1, (v2 % T1))
end
@inline function Base.:(>>>)(v1::AbstractSIMD{W,T1}, v2::AbstractSIMD{W,T2}) where {W,T1<:SignedHW,T2<:UnsignedHW}
    convert(T1, vlshr(convert(T2, v1), v2))
end
@inline function Base.:(>>>)(v1::AbstractSIMD{W,T1}, v2::AbstractSIMD{W,T2}) where {W,T1<:UnsignedHW,T2<:SignedHW}
    convert(T2, vlshr(v1, convert(T1, v2)))
end

@inline function promote_except_first(a,b,c)
    d, e = promote(b, c)
    a, d, e
end
for (op, f, promotef) ∈ [
    (:(Base.fma), :vfma, :promote),
    (:(Base.muladd), :vmuladd, :promote),
    (:(IfElse.ifelse), :vifelse, :promote_except_first)
]
    @eval begin
        @inline function $op(a::AbstractSIMD, b::AbstractSIMD, c::AbstractSIMD)
            x, y, z = $promotef(a, b, c)
            $f(x, y, z)
        end
        @inline function $op(a::AbstractSIMD, b::AbstractSIMD, c::NativeTypes)
            x, y, z = $promotef(a, b, c)
            $f(x, y, z)
        end
        @inline function $op(a::AbstractSIMD, b::NativeTypes, c::AbstractSIMD)
            x, y, z = $promotef(a, b, c)
            $f(x, y, z)
        end
        @inline function $op(a::NativeTypes, b::AbstractSIMD, c::AbstractSIMD)
            x, y, z = $promotef(a, b, c)
            $f(x, y, z)
        end
        @inline function $op(a::AbstractSIMD, b::NativeTypes, c::NativeTypes)
            x, y, z = $promotef(a, b, c)
            $f(x, y, z)
        end
        @inline function $op(a::NativeTypes, b::AbstractSIMD, c::NativeTypes)
            x, y, z = $promotef(a, b, c)
            $f(x, y, z)
        end
        @inline function $op(a::NativeTypes, b::NativeTypes, c::AbstractSIMD)
            x, y, z = $promotef(a, b, c)
            $f(x, y, z)
        end
    end
end
@inline IfElse.ifelse(f::Function, m::AbstractSIMD{W,B}, args::Vararg{NativeTypesV,K}) where {W,K,B<:Union{Bool,Bit}} = vifelse(f, m, args...)
@inline IfElse.ifelse(f::Function, m::Bool, args::Vararg{NativeTypesV,K}) where {K} = vifelse(f, m, args...)
for (f,add,mul) ∈ [
    (:fma,:(+),:(*)), (:muladd,:(+),:(*)),
    (:vfma,:(+),:(*)), (:vmuladd,:(+),:(*)),
    (:vfma_fast,:(Base.FastMath.add_fast),:(Base.FastMath.mul_fast)),
    (:vmuladd_fast,:(Base.FastMath.add_fast),:(Base.FastMath.mul_fast))
]
    if (f !== :fma) && (f !== :muladd)
        @eval begin
            @inline function $f(a::NativeTypesV, b::NativeTypesV, c::NativeTypesV)
                x, y, z = promote(a, b, c)
                $f(x, y, z)
            end
        end
    else
        f = :(Base.$f)
    end
    @eval begin
        @inline $f(a::Zero, b::NativeTypesV, c::NativeTypesV) = c
        @inline $f(a::NativeTypesV, b::Zero, c::NativeTypesV) = c
        @inline $f(a::Zero, b::Zero, c::NativeTypesV) = c
        @inline $f(a::One, b::Zero, c::NativeTypesV) = c
        @inline $f(a::Zero, b::One, c::NativeTypesV) = c

        @inline $f(a::One, b::NativeTypesV, c::NativeTypesV) = $add(b, c)
        @inline $f(a::NativeTypesV, b::One, c::NativeTypesV) = $add(a, c)
        @inline $f(a::One, b::One, c::NativeTypesV) = $add(one(c), c)

        @inline $f(a::NativeTypesV, b::NativeTypesV, c::Zero) = $mul(a,b)
        @inline $f(a::Zero, b::NativeTypesV, c::Zero) = Zero()
        @inline $f(a::NativeTypesV, b::Zero, c::Zero) = Zero()
        @inline $f(a::Zero, b::Zero, c::Zero) = Zero()
        @inline $f(a::One, b::Zero, c::Zero) = Zero()
        @inline $f(a::Zero, b::One, c::Zero) = Zero()

        @inline $f(a::One, b::NativeTypesV, c::Zero) = b
        @inline $f(a::NativeTypesV, b::One, c::Zero) = a
        @inline $f(a::One, b::One, c::Zero) = One()
    end
end

# masks
for (vf,f) ∈ [
    (:vnot,:(!)),
    (:vnot,:(~)),
]
    @eval begin
        @inline Base.$f(m::AbstractSIMD{<:Any,<:Union{Bool,Bit}}) = $vf(m)
    end
end

for f ∈ [:typemin, :typemax, :floatmin, :floatmax]
    @eval @inline Base.$f(::Type{V}) where {W,T,V <: AbstractSIMD{W,T}} = convert(V, $f(T))
end


text:
using ITensors
using ITensorNetworkAD
using ITensorNetworkAD.ITensorNetworks: Models

@testset "test local hamiltonian builder" begin
  Nx = 2
  Ny = 3
  sites = siteinds("S=1/2", Ny, Nx)
  H = Models.mpo(Models.Model("tfim"), sites; h=1.0)
  H_local = Models.localham(Models.Model("tfim"), sites; h=1.0)
  H_line = Models.lineham(Models.Model("tfim"), sites; h=1.0)
  @test Models.checkham(H_local, H, sites)
  @test Models.checkham(H_line, H, sites)
end