@testset "Util" begin

    # Test `get_window_bounds()`.
    let A = Gray.(zeros(5,5))
        for index in CartesianIndices(A)
            bounds = ImageBinarization.get_window_bounds(A, index, 1)
            for bound = bounds
                @test 1 <= bound <= 5
            end
        end
    end

    # Test `μ_in_window()`.
    let A = [(x-1) * 5 + y for x = 1:5, y = 1:5]
        I = ImageBinarization.integral_image(A)
        for r₀ in 1:5, c₀ in 1:5
            for r₁ in r₀:5, c₁ in c₀:5
                @test ImageBinarization.μ_in_window(I, r₀, c₀, r₁, c₁) == mean(vec(A[r₀:r₁,c₀:c₁]))
            end
        end
    end

    # Test `σ²_in_window()`.
    let A = [(x-1) * 5 + y for x = 1:5, y = 1:5]
        I² = ImageBinarization.integral_image(A.^2)
        for r₀ in 1:5, c₀ in 1:5
            for r₁ in r₀:5, c₁ in c₀:5
                μ = mean(vec(A[r₀:r₁,c₀:c₁]))
                δ = ImageBinarization.σ²_in_window(I², μ, r₀, c₀, r₁, c₁) - var(vec(A[r₀:r₁,c₀:c₁]), corrected=false)
                @test δ == NaN || δ < 1e-10
            end
        end
    end

    # Test `σ_in_window()`.
    let A = [(x-1) * 5 + y for x = 1:5, y = 1:5]
        I² = ImageBinarization.integral_image(A.^2)
        for r₀ in 1:5, c₀ in 1:5
            for r₁ in r₀:5, c₁ in c₀:5
                μ = mean(vec(A[r₀:r₁,c₀:c₁]))
                δ = ImageBinarization.σ_in_window(I², μ, r₀, c₀, r₁, c₁) - std(vec(A[r₀:r₁,c₀:c₁]), corrected=false)
                @test δ == NaN || δ < 1e-10
            end
        end
    end
end


export LabeledTensor

# abstractions
abstract type AbstractTensor{T, N} end

"""
    LabeledTensor{T,N,AT<:AbstractArray{T,N}, LT, MT} <: AbstractTensor{T, N}
    LabeledTensor(array::AbstractArray, labels::AbstractVector[, meta])

Tensor with labeled legs.
When multiplying two tensors, legs with same labels will be treated as inner degree of freedom and get contracted.

The optional argument `meta` is the meta information, default is `nothing`. One can use it to store additional information like plotting layout.
"""
struct LabeledTensor{T,N,AT<:AbstractArray{T,N}, LT, MT} <: AbstractTensor{T, N}
    array::AT
    labels::Vector{LT}
    meta::MT
end
Base.ndims(t::LabeledTensor) = ndims(t.array)
LinearAlgebra.norm(t::LabeledTensor, p::Real=2) = norm(t.array, p)

function LabeledTensor(tensor::AbstractArray, labels::AbstractVector)
    @assert ndims(tensor) == length(labels) "dimension of tensor $(ndims(tensor)) != number of labels $(length(labels))"
    LabeledTensor(tensor, labels, nothing)
end

function Base.:(*)(A::LabeledTensor, B::LabeledTensor)
    labels_AB = setdiff(A.labels, B.labels) ∪ setdiff(B.labels, A.labels)
    LabeledTensor(tensorcontract(A.labels, A.array, B.labels, B.array, labels_AB), labels_AB, merge_meta(A.meta, B.meta))
end

function Base.isapprox(a::LabeledTensor, b::LabeledTensor; kwargs...)
    isapprox(a.array, b.array; kwargs...) && a.labels == b.labels
end
Base.size(t::LabeledTensor) = Base.size(t.array)
Base.copy(t::LabeledTensor) = LabeledTensor(copy(t.array), t.labels)
Base.similar(::Type{<:LabeledTensor}, arr::AbstractArray, labels::AbstractVector, meta=nothing) = LabeledTensor(arr, labels, meta)
Base.similar(::LabeledTensor, arr::AbstractArray, labels::AbstractVector, meta=nothing) = LabeledTensor(arr, labels, meta)
LinearAlgebra.rmul!(t::LabeledTensor, factor) = (rmul!(t.array, factor); t)

function Base.show(io::IO, lt::LabeledTensor)
    print(io, "$(typeof(lt).name.name){$(eltype(lt.array))}($(join(lt.labels, ", ")))")
end

function Base.show(io::IO, ::MIME"plain/text", lt::LabeledTensor)
    Base.show(io, lt)
end

function mul_dim(t::LabeledTensor, m::AbstractMatrix; dim::Int)
    data = t.array
    iA = ntuple(i->i, ndims(data))
    iB = (dim, -dim)
    iC = ntuple(i->i==dim ? -dim : i, ndims(data))
    LabeledTensor(tensorcontract(iA, data, iB, m, iC), t.labels)
end

struct PlotMeta
    loc::Tuple{Float64, Float64}
    name::String
end
merge_meta(m1::PlotMeta, m2::PlotMeta) = PlotMeta((m1.loc .+ m2.loc) ./ 2, m1.name*m2.name)
merge_meta(m1::Nothing, m2::Nothing) = nothing
dispmeta(m::PlotMeta) = m.name


#################################################################################
#               Genera gráfica de las funciones de activación
#################################################################################
# Paquetes
using Plots
using TOML
include("../../OptimizedNeuralNetwork.jl/src/activation_functions.jl")
using .ActivationFunctions

FICHERO_CONFIGURACION = "Experimentos/.config.toml"
config = TOML.parsefile(FICHERO_CONFIGURACION)["visualizacion-funciones-activacion"]
img_path = config["DIRECTORIO_IMAGENES"]

### Funciones que se van a analizar  y auxiliares
p(x) = 2x +1
id(x)=x
# Estructura caracteríticas a mostrar
struct FuncionActivacion
    nombre::String
    funcion_activacion # función propiamente dicho
    cota_inferior :: Real
    cota_superior :: Real
end

# Declaración de funciones
funciones_a_mostrar = [
    FuncionActivacion("Identidad", id, -1,1),
    FuncionActivacion("Threshold de 2x+1", @ThresholdFunction(p,0), -2, 2),
    FuncionActivacion("Cosine CosineSquasher", CosineSquasher, -2pi, 2pi),
    FuncionActivacion("Indicadora de 0", @IndicatorFunction(0), -2, 2),
    FuncionActivacion("Rampa", RampFunction, -2, 2),
    FuncionActivacion("ReLU", ReLU, -2, 2),
    FuncionActivacion("Sigmoidea", Sigmoid, -4,4),
    FuncionActivacion("Tangente hiperbolica", tanh, -4,4),
    FuncionActivacion("Valor absoluto", abs, -2,2),
    FuncionActivacion("Coseno", cos, -2pi,2pi),
    FuncionActivacion("Hardtanh", HardTanh, -2pi,2pi),
    FuncionActivacion("LReLU", @LReLU(0.01), -2pi,2pi),
]

"""
    MuestraYGuarda(f :: FuncionActivacion, img_path:: String)
Muestra por pantalla una función de activación con la estructura FuncionActivacion
y la guarda en el directorio indicado en `img_path`
"""
function MuestraYGuarda(fa :: FuncionActivacion, img_path:: String)
    plot( fa.funcion_activacion, fa.cota_inferior, fa.cota_superior, title="", label="",linewidth=3)
    savefig(img_path*fa.nombre)
end

# Mostramos y guardamos para las funciones definidas
for fa in funciones_a_mostrar
   MuestraYGuarda(fa, img_path)
end


#=--------------------------------------------------------------------

frequencyfilt - Functions for constructing image filters in the
                frequency domain.

Copyright (c) Peter Kovesi
peterkovesi.com

MIT License:

Permission is hereby granted, free of charge, to any person obtaining
a copy of this software and associated documentation files (the
"Software"), to deal in the Software without restriction, including
without limitation the rights to use, copy, modify, merge, publish,
distribute, sublicense, and/or sell copies of the Software, and to
permit persons to whom the Software is furnished to do so, subject to
the following conditions:

The above copyright notice and this permission notice shall be
included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

PK August 2015   Original porting from MATLAB to Julia
   October 2017  Updates for 0.6
   October 2018  Julia 0.7/1.0
---------------------------------------------------------------------=#

export filtergrid, filtergrids, gridangles
export cosineangularfilter, gaussianangularfilter
export lowpassfilter, highpassfilter, bandpassfilter, highboostfilter
export loggabor, monogenicfilters, packedmonogenicfilters
export perfft2
export geoseries
# export homomorphic

#--------------------------------------------------------------------
# filtergrids
"""
Generate grids for constructing frequency domain filters.
```
Usage:  (f, fx, fy) = filtergrids(rows, cols)
        (f, fx, fy) = filtergrids((rows, cols))

Arguments:  rows, cols - Size of image/filter

Returns:             f - Grid of size (rows, cols) containing frequency
                         values from 0 to 0.5,  where f =
                         sqrt(fx^2 + fy^2). The grid is quadrant
                         shifted so that 0 frequency is at f[1,1]

                fx, fy - Grids containing normalised frequency values
                         ranging from -0.5 to 0.5 in x and y directions
                         respectively. fx and fy are quadrant shifted.
```

See also: [`filtergrid`](@ref)  where you are only needing radius

"""
function filtergrids(rows::Integer, cols::Integer)

    # Set up X and Y spatial frequency matrices, fx and fy, with ranges
    # normalised to +/- 0.5 The following code adjusts things appropriately for
    # odd and even values of rows and columns so that the 0 frequency point is
    # placed appropriately.
    if isodd(cols)
        fxrange = (-(cols-1)/2:(cols-1)/2)/cols
    else
        fxrange = (-cols/2:(cols/2-1))/cols
    end

    if isodd(rows)
        fyrange = (-(rows-1)/2:(rows-1)/2)/rows
    else
        fyrange = (-rows/2:(rows/2-1))/rows
    end

    fx = [c for r = fyrange, c = fxrange]
    fy = [r for r = fyrange, c = fxrange]

    # Quadrant shift so that filters are constructed with 0 frequency at
    # the corners
    fx = ifftshift(fx)
    fy = ifftshift(fy)

    # Construct spatial frequency values in terms of normalised radius from
    # centre.
    f = sqrt.(fx.^2 .+ fy.^2)
    return f, fx, fy
end

# Tuple version
function filtergrids(sze::Tuple{Integer,Integer})
    return filtergrids(sze[1], sze[2])
end

#--------------------------------------------------------------------
# filtergrid
"""
Generate grid for constructing frequency domain filters.
```
Usage:  f = filtergrid(rows, cols)
        f = filtergrid((rows, cols))

Arguments:  rows, cols - Size of image/filter

Returns:             f - Grid of size (rows, cols) containing normalised
                         frequency values from 0 to 0.5.  Grid is quadrant
                         shifted so that 0 frequency is at f[1,1]

```
Used by [`phasecongmono`](@ref), [`phasecong3`](@ref), etc etc

See also: [`filtergrids`](@ref)   if you also want normalized frequency grids in
          the x and y directions as well.
"""
function filtergrid(rows::Integer, cols::Integer)

    # Set up X and Y spatial frequency ranges normalised to +/- 0.5
    # The following code adjusts things appropriately for odd and even
    # values of rows and columns so that the 0 frequency point is
    # placed appropriately.
    if isodd(cols)
        fxrange = (-(cols-1)/2:(cols-1)/2)/cols
    else
        fxrange = (-cols/2:(cols/2-1))/cols
    end

    if isodd(rows)
        fyrange = (-(rows-1)/2:(rows-1)/2)/rows
    else
        fyrange = (-rows/2:(rows/2-1))/rows
    end

    # Construct spatial frequency values in terms of normalised radius from
    # centre.
    f = [sqrt(fx^2 + fy^2) for fy in fyrange, fx in fxrange]
    return ifftshift(f)
end

# Tuple version
function filtergrid(sze::Tuple{Integer,Integer})
    return filtergrid(sze[1], sze[2])
end

#--------------------------------------------------------------------
# monogenicfilters
"""
Generate monogenic filter grids.

```
Usage: (H1, H2, f) = monogenicfilters(rows, cols)
       (H1, H2, f) = monogenicfilters((rows, cols))

Arguments:  rows,cols - Size of filters to generate

Returns: H1, H2 - The two monogenic filters.
              f - Frequency grid corresponding to the filters.

where:
       H1 = i*fx./f
       H2 = i*fy./f

```

Note that H1, H2, and f and quadrant shifted to that the 0 frequency
value is at coordinate [1,1].

See also: [`packedmonogenicfilters`](@ref)
"""
function monogenicfilters(rows::Integer, cols::Integer)

    (f, fx, fy) = filtergrids(rows, cols)
    f[1,1] = 1  # Set DC value to 1 to avoid divide by zero

    H1 = im.*fx./f
    H2 = im.*fy./f

    H1[1,1] = 0  # Restore 0 DC value
    H2[1,1] = 0
    f[1,1] = 0
    return H1, H2, f
end

# Tuple version
function monogenicfilters(sze::Tuple{Integer,Integer})
    return monogenicfilters(sze[1], sze[2])
end

#--------------------------------------------------------------------
# packedmonogenicfilters
"""
Monogenic filter where both filters are packed in the one Complex grid.

```
Usage: (H, f) = packedmonogenicfilters(rows, cols)
       (H, f) = packedmonogenicfilters((rows, cols))

Arguments:  rows,cols - Size of filters to generate

Returns:      H - The two monogenic filters packed into the
                  one Complex64 grid.
              f - Frequency grid corresponding to the filter.
```

The two monogenic filters are defined as

```
       H1 = i*fx./f
       H2 = i*fy./f
```

However the two filters can be packed together as a complex valued
matrix, one in the real part and one in the imaginary part.  Do this
by multiplying H2 by i and then adding it to H1.  When the convolution
is performed via the fft the real part of the result will correspond
to the convolution with H1 and the imaginary part with H2.  This
allows the two convolutions to be done as one in the frequency domain,
saving time and memory.

Note that H and f and quadrant shifted to that the 0 frequency
value is at coordinate [1,1].

See also: [`monogenicfilters`](@ref)
"""
function packedmonogenicfilters(rows::Integer, cols::Integer)
    (f, fx, fy) = filtergrids(rows, cols)
    f[1,1] = 1  # Set DC value to 1 to avoid divide by zero

    # Pack the two monogenic filters by multiplying H2 by i and then
    # adding it to H1 (note the subtraction because i*i = -1).
    H = (im.*fx .- fy)./f

    H[1,1] = 0   # Restore 0 DC value
    f[1,1] = 0
    return H, f
end

# Tuple version
function packedmonogenicfilters(sze::Tuple{Integer,Integer})
    return packedmonogenicfilters(sze[1], sze[2])
end

#--------------------------------------------------------------------
# lowpassfilter
"""
Construct a low-pass Butterworth filter.
```
Usage: f = lowpassfilter(sze, cutoff, n)

where: sze    is a two element tuple specifying the size of filter
              to construct (rows, cols).
       cutoff is the cutoff frequency of the filter 0 - 0.5
       n      is the order of the filter, the higher n is the sharper
              the transition is. (n must be an integer >= 1).
              Note that n is doubled so that it is always an even integer.

                      1
      f =    --------------------
                              2n
              1.0 + (w/cutoff)
```
The frequency origin of the returned filter is at the corners.

See also: [`highpassfilter`](@ref), [`highboostfilter`](@ref), [`bandpassfilter`](@ref)
"""
function lowpassfilter(sze::Tuple{Integer, Integer}, cutoff::Real, n::Integer)

    if cutoff < 0 || cutoff > 0.5
	error("cutoff frequency must be between 0 and 0.5")
    end

    f = filtergrid(sze)
    return 1.0 ./ (1.0 .+ (f ./ cutoff).^(2*n))
end

# Compute the low pass filter value at a specified normalised frequency
function lowpassfilter(f::Real, cutoff::Real, n::Integer)
    return 1.0 / (1.0 + (f / cutoff)^(2*n))
end
#--------------------------------------------------------------------
# bandpassfilter
"""
Construct a band-pass Butterworth filter.
```
Usage: f = bandpassfilter(sze, cutin, cutoff, n)

Arguments:
             sze - A 2 element tuple specifying the size of filter
                   to construct (rows, cols).
   cutin, cutoff - The frequencies defining the band pass 0 - 0.5
               n - The order of the filter, the higher n is the sharper
                   the transition is. (n must be an integer >= 1).
Returns:
               f - Frequency domain filter of size==sze, the frequency
                   origin is at the corners.
```
See also: [`lowpassfilter`](@ref), [`highpassfilter`](@ref), [`highboostfilter`](@ref)
"""
function bandpassfilter(sze::Tuple{Integer, Integer}, cutin::Real, cutoff::Real, n::Integer)

    if cutin < 0 || cutin > 0.5 || cutoff < 0 || cutoff > 0.5
	error("Frequencies must be between 0 and 0.5")
    end

    if n < 1
        error("Order of filter must be greater than 1")
    end

    return lowpassfilter(sze, cutoff, n) - lowpassfilter(sze, cutin, n)
end

#--------------------------------------------------------------------
# highboostfilter
"""
Construct a high-boost Butterworth filter.
```
Usage: f = highboostfilter(sze, cutoff, n, boost)

Arguments:
         sze - A 2 element tuple specifying the size of filter
               to construct (rows, cols).
      cutoff - The cutoff frequency of the filter 0 - 0.5
           n - The order of the filter, the higher n is the sharper
               the transition is. (n must be an integer >= 1).
       boost - The ratio that high frequency values are boosted
               relative to the low frequency values.  If boost is less
               than one then a 'lowboost' filter is generated
Returns:
           f - Frequency domain filter of size==sze, the frequency
               origin is at the corners.
```
See also: [`lowpassfilter`](@ref), [`highpassfilter`](@ref), [`bandpassfilter`](@ref)
"""
function highboostfilter(sze::Tuple{Integer, Integer}, cutoff::Real, n::Integer, boost::Real)

    if cutoff < 0 || cutoff > 0.5
	error("cutoff frequency must be between 0 and 0.5")
    end

    if boost >= 1     # high-boost filter
	f = (1.0 - 1.0/boost)*highpassfilter(sze, cutoff, n) .+ 1.0/boost
    else              # low-boost filter
	f = (1.0 - boost)*lowpassfilter(sze, cutoff, n) .+ boost
    end

    return f
end

#--------------------------------------------------------------------
# highpassfilter
"""
Construct a high-pass Butterworth filter.
```
Usage: f = highpassfilter(sze, cutoff, n)

         sze - A 2 element tuple specifying the size of filter
               to construct (rows, cols).
      cutoff - The cutoff frequency of the filter 0 - 0.5
           n - The order of the filter, the higher n is the sharper
               the transition is. (n must be an integer >= 1).
Returns:
           f - Frequency domain filter of size==sze, the frequency
               origin is at the corners.
```
See also: [`lowpassfilter`](@ref), [`highboostfilter`](@ref), [`bandpassfilter`](@ref)
"""
function highpassfilter(sze::Tuple{Integer, Integer}, cutoff::Real, n::Integer)

    if cutoff < 0 || cutoff > 0.5
	error("cutoff frequency must be between 0 and 0.5")
    end

    return 1.0 .- lowpassfilter(sze, cutoff, n)
end
#--------------------------------------------------------------------
# loggabor
"""
The logarithmic Gabor function in the frequency domain.

```
Usage: v = loggabor(f::Real, fo::Real, sigmaOnf::Real)

Arguments:
            f - Frequency to evaluate the function at.
           fo - Centre frequency of filter.
     sigmaOnf - Ratio of the standard deviation of the Gaussian
                describing the log Gabor filter's transfer function
                in the frequency domain to the filter center frequency.

sigmaOnf = 0.75 gives a filter bandwidth of about 1 octave.
sigmaOnf = 0.55 gives a filter bandwidth of about 2 octaves.

```
"""
function loggabor(f::Real, fo::Real, sigmaOnf::Real)
    if f < eps()
        return 0.0
    else
        return exp((-(log(f/fo))^2) / (2 * log(sigmaOnf)^2))
    end
end

#-------------------------------------------------------------
# gridangles
"""
Generate arrays of filter grid angles.
```
Usage: (sintheta, costheta) = gridangles(freq, fx, fy)

Arguments: freq, fx, fy - The output of filtergrids()

Returns:       sintheta - The sine and cosine of the angles in the filtergrid
               costheta

```
See also [`filtergrids`](@ref)

"""
function gridangles(freq::AbstractArray{T1,2},
                    fx::AbstractArray{T2,2}, fy::AbstractArray{T3,2}) where {T1 <: Real, T2 <: Real, T3 <: Real}

    freq[1,1] = 1         # Avoid divide by 0
    sintheta = fx./freq   # sine and cosine of filter grid angles
    costheta = fy./freq
    freq[1,1] = 0         # Restore 0 DC

    return sintheta, costheta
end

#--------------------------------------------------------------------
# cosineangularfilter
"""
Orientation selective frequency domain filter with cosine windowing function.

```
Usage: filter = cosineangularfilter(angl, wavelen, sintheta, costheta)

Arguments:
               angl - Orientation of the filter (radians)
            wavelen - Wavelength of the angular cosine window function.
 sintheta, costheta - Grids as returned by gridangles()
```
See also: [`gaussianangularfilter`](@ref), [`filtergrids`](@ref)
"""
function cosineangularfilter(angl::Real, wavelen::Real,
                             sintheta::Array{T1,2}, costheta::Array{T2,2}) where {T1 <: Real, T2 <: Real}

    sinangl = sin(angl); cosangl = cos(angl)
    fltr = zeros(size(sintheta))

    # For each point in the filter matrix calculate the angular
    # distance from the specified filter orientation.  To overcome
    # the angular wrap-around problem sine difference and cosine
    # difference values are first computed and then the atan2
    # function is used to determine angular distance.
    for n in eachindex(sintheta)
        ds = sintheta[n] * cosangl - costheta[n] * sinangl  # Difference in sine.
        dc = costheta[n] * cosangl + sintheta[n] * sinangl  # Difference in cosine.
        dtheta = abs(atan(ds, dc))                   # Absolute angular distance.

        # Scale theta so that cosine spread function has the right
        # wavelength and clamp to pi.  dtheta has a wavelength of
        # 2pi. If desired wavelength of cosine window function is
        # wavelen we need to multiply dtheta by 2*pi/wavelen.
        dtheta = min(dtheta*2*pi/wavelen, pi)
        # The spread function is cos(dtheta) between -pi and pi.  We add 1,
        # and then divide by 2 so that the value ranges 0-1
        fltr[n] = (cos(dtheta)+1)/2
    end

    return fltr
end

#--------------------------------------------------------------------
# gaussianangularfilter
"""
Orientation selective frequency domain filter with Gaussian windowing function.

```
Usage: filter = gaussianangularfilter(angl, thetaSigma, sintheta, costheta)

Arguments:
               angl - Orientation of the filter (radians)
         thetasigma - Standard deviation of angular Gaussian window function.
 sintheta, costheta - Grids as returned by gridangles()
```

See also: [`cosineangularfilter`](@ref), [`gridangles`](@ref), [`filtergrids`](@ref)
"""
function gaussianangularfilter(angl::Real, thetaSigma::Real,
                               sintheta::Array{T1,2}, costheta::Array{T2,2}) where {T1 <: Real, T2 <: Real}

    sinangl = sin(angl); cosangl = cos(angl)
    fltr = zeros(size(sintheta))

    # For each point in the filter matrix calculate the angular
    # distance from the specified filter orientation.  To overcome
    # the angular wrap-around problem sine difference and cosine
    # difference values are first computed and then the atan2
    # function is used to determine angular distance.
    for n in eachindex(sintheta)
        ds = sintheta[n] * cosangl - costheta[n] * sinangl  # Difference in sine.
        dc = costheta[n] * cosangl + sintheta[n] * sinangl  # Difference in cosine.
        dtheta = atan(ds, dc)                               # Angular distance.

        fltr[n] = exp((-dtheta.^2) / (2 * thetaSigma^2))
    end
    return fltr
end

#--------------------------------------------------------------------
#=
"""
homomorphic - Performs homomorphic filtering on an image.

Function performs homomorphic filtering on an image. This form of
filtering sharpens features and flattens lighting variantions in an
image.  It usually is very effective on images which have large
variations in lighting, for example when a subject appears against
strong backlighting.
```
Usage: newim =
homomorphic(inimage,boost,CutOff,order,lhistogram_cut,uhistogram_cut, hndl)
homomorphic(inimage,boost,CutOff,order,lhistogram_cut,uhistogram_cut)
homomorphic(inimage,boost,CutOff,order,hndl)
homomorphic(inimage,boost,CutOff,order)

Parameters:  (suggested values are in brackets)
        boost    - The ratio that high frequency values are boosted
                   relative to the low frequency values (2).
        CutOff   - Cutoff frequency of the filter (0 - 0.5)
        order    - Order of the modified Butterworth style filter that
                   is used, this must be an integer > 1 (2)
        lhistogram_cut - Percentage of the lower end of the filtered image's
                         histogram to be truncated, this eliminates extreme
                         values in the image from distorting the final result. (0)
        uhistogram_cut - Percentage of upper end of histogram to truncate. (5)
        hndl           - Optional handle to text box for updating
                         messages to be sent to a GUI interface.
```

If lhistogram_cut and uhistogram_cut are not specified no histogram
truncation will be applied.

Suggested values: newim = homomorphic(im, 2, .25, 2, 0, 5)
"""

# June 1999
# December 2001 cleaned up and modified to work with colour images

function him = homomorphic(im, boost, CutOff, order, varargin)

	if ndims(im) == 2  # Greyscale image
	    him = Ihomomorphic(im, boost, CutOff, order, varargin)

	else               # Assume colour image in RGB format
	    hsv = rgb2hsv(im)   # Convert to HSV and apply homomorphic
				 # filtering to just the intensity component.
            hsv(:,:,3) = Ihomomorphic(hsv(:,:,3), boost, CutOff, order, varargin)
	    him = hsv2rgb(hsv)  # Convert back to RGB
	end

#------------------------------------------------------------------------
# Internal function that does the real work
#------------------------------------------------------------------------

function him = Ihomomorphic(im, boost, CutOff, order, varargin)

    # The possible elements in varargin are:
    # {lhistogram_cut, uhistogram_cut, hndl}

    varargin = varargin{:}

    if nargin == 5
	nopparams  = length(varargin)
    end

    if (nopparams == 3)
	dispStatus = 1
	truncate = 1
	lhistogram_cut = varargin{1}
	uhistogram_cut = varargin{2}
	hndl = varargin{3}
    elseif (nopparams == 2)
	dispStatus = 0
	truncate = 1
	lhistogram_cut = varargin{1}
	uhistogram_cut = varargin{2}
    elseif (nopparams == 1)
	dispStatus = 1
	truncate = 0
	hndl = varargin{1}
    elseif (nopparams == 0)
	dispStatus = 0
	truncate = 0
    else
	disp('Usage: newim = homomorphic(inimage,LowGain,HighGain,CutOff,order,lhistogram_cut,uhistogram_cut)')
	error('or    newim = homomorphic(inimage,LowGain,HighGain,CutOff,order)')
    end

    [rows,cols] = size(im)

    im = normalise(im)                        # Rescale values 0-1 (and cast
					       # to `double' if needed).
    FFTlogIm = fft2(log(im+.01))              # Take FFT of log (with offset
                                               # to avoid log of 0).
    h = highboostfilter([rows cols], CutOff, order, boost)
    him = exp(real(ifft2(FFTlogIm.*h)))       # Apply the filter, invert
					       # fft, and invert the log.

    if truncate

	# Problem:
	# The extreme bright values in the image are exaggerated by the filtering.
	# These (now very) bright values have the overall effect of darkening the
	# whole image when we rescale values to 0-255.
	#
	# Solution:
	# Construct a histogram of the image.  Find the level below which a high
	# percentage of the image lies (say 95#).  Saturate the grey levels in
	# the image to this level.

	if dispStatus
	    set(hndl,'String','Calculating histogram and truncating...')
	    drawnow
	else
	    disp('Calculating histogram and truncating...')
	end

	him = histtruncate(him, lhistogram_cut, uhistogram_cut)

    else
	him = normalise(him)  # No truncation, but fix range 0-1
    end

end
=#

#--------------------------------------------------------------------
#  perfft2
"""
2D Fourier transform of Moisan's periodic image component.
```
Usage: (P, S, p, s) = perfft2(img)

Argument: img - Image to be transformed
Returns:    P - 2D fft of periodic image component
            S - 2D fft of smooth component
            p - Periodic component (spatial domain)
            s - Smooth component (spatial domain)
```
Moisan's "Periodic plus Smooth Image Decomposition" decomposes an image
into two components

        img = p + s

where s is the 'smooth' component with mean 0 and p is the 'periodic'
component which has no sharp discontinuities when one moves cyclically
across the image boundaries.

This decomposition is very useful when one wants to obtain an FFT of
an image with minimal artifacts introduced from the boundary
discontinuities.  The image p gathers most of the image information
but avoids periodization artifacts.

The typical use of this function is to obtain a 'periodic only' fft of an
image

      P = perfft2(img)

Displaying the amplitude spectrum of P will yield a clean spectrum without the
typical vertical-horizontal 'cross' arising from the image boundaries that you
would normally see.

Note if you are using the function to perform filtering in the frequency
domain you may want to retain s (the smooth component in the spatial domain)
and add it back to the filtered result at the end.

The computational cost of obtaining the 'periodic only' FFT involves taking an
additional FFT.
"""
function perfft2(img::Array{T,2}) where T <: Real
    #=
    Reference:
    This code is adapted from Lionel Moisan's Scilab function 'perdecomp.sci'
    "Periodic plus Smooth Image Decomposition" 07/2012 available at

    http://www.mi.parisdescartes.fr/~moisan/p+s

    Paper:
    L. Moisan, "Periodic plus Smooth Image Decomposition", Journal of
    Mathematical Imaging and Vision, vol 39:2, pp. 161-179, 2011.
    =#

    (rows,cols) = size(img)

    # Compute the boundary image which is equal to the image discontinuity
    # values across the boundaries at the edges and is 0 elsewhere
    s = zeros(rows, cols)
    s[1,:]   = img[1,:] - img[end,:]
    s[end,:] = -s[1,:]
    s[:,1]   = s[:,1]   + img[:,1] - img[:,end]
    s[:,end] = s[:,end] - img[:,1] + img[:,end]

    # Generate grid upon which to compute the filter for the boundary image in
    # the frequency domain.  Note that cos() is cyclic hence the grid values can
    # range from 0 .. 2*pi rather than 0 .. pi and then pi .. 0

    # Generate FFT of smooth component
    cxrange = 2*pi*(0:cols-1)/cols
    cyrange = 2*pi*(0:rows-1)/rows
    denom = [2*(2 - cos(cx) - cos(cy)) for cy in cyrange, cx in cxrange]
    S = fft(s)./denom

    # The [1,1] element of the filter will be 0 so S[1,1] may be Inf or NaN
    S[1,1] = 0.0         # Enforce 0 mean

    P = fft(img) .- S    # FFT of periodic component

# ** ? Perhaps have a separate version or a flag to request the
# ** spatial versions of p and s
    s = real(ifft(S))
    p = img .- s

    return P, S, p, s
end

#----------------------------------------------------------------------
# geoseries
"""
Generate geometric series.

Useful for generating geometrically scaled wavelengths for specifying
filter banks.

```
Usage 1: s = geoseries(s1, mult, n)

Arguments:      s1 - The starting value in the series.
              mult - The scaling factor between succesive values.
                 n - The desired number of elements in the series.

Usage 2: s = geoseries((s1, sn), n)

Arguments: (s1, sn) - Tuple specifying the 1st and last values
                      in the the series.
                  n - The desired number of elements in the series.
```

Example:
```
      s = geoseries(0.5, 2, 4)
      s =  [0.5000,    1.0000,    2.0000,    4.0000]
```
 Alternatively obtain the same series using
```
           s = geoseries((0.5, 4), 4)
```
"""
function geoseries(s1::Real, mult::Real, n::Integer)
    @assert n > 0  "Number of elements must be a +ve integer"
    @assert s1 > 0  "Starting value must be > 0"
    return s = s1 * mult.^(0:(n-1))
end

function geoseries(s1sn::Tuple{Real, Real}, n::Int)

    # Compute the multiplier from the desired number of elements.
    # max_val = s1*mult^(n-1)
    s1 = s1sn[1]
    sn = s1sn[2]
    @assert s1 > 0  "Starting value must be > 0"
    mult = exp(log(sn/s1)/(n-1))
    return  geoseries(s1, mult, n)
end

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


using DelimitedFiles

function read_numbers(input_filename)
    number_string = readlines(input_filename)[1]
    number_strings = split(number_string, ",")
    numbers = map(x->parse(Int64, x), number_strings)
    return numbers
end

function read_bingo_cards(input_filename)
    bigmatrix = readdlm(input_filename, Int, skipstart=1)
    cards = []
    # Assume cards are square
    card_size = size(bigmatrix)[2]
    for i in 1:Int64((size(bigmatrix)[1]/card_size))
        push!(cards, bigmatrix[(i-1)*card_size+1:i*card_size, :])
    end
    return cards
end

function check_row_or_column(row_or_column, nums)
    return sum(in(nums).(row_or_column)) == length(row_or_column)
end

function check_card_win(card, nums)
    card_size = size(card)[1]
    for i in 1:card_size
        if (check_row_or_column(card[i,:],nums) ||
            check_row_or_column(card[:,i],nums))
            return true
        end
    end
    return false
end

function test_card_win()
    card = [1 2 3
            4 5 6
            7 8 9]
    @assert check_card_win(card, [1,2]) == false
    @assert check_card_win(card, [1,2,3]) == true
    @assert check_card_win(card, [1,4,7]) == true
    @assert check_card_win(card, [9,5,8,1,6,4]) == true
end

function get_card_score(card, nums)
    last_number_called = nums[length(nums)]
    # flatten the card
    card_flat = vec(card)
    # get a bit vector of numbers that are not called
    matches = (!).(in(nums).(card_flat))
    # multiply by the numbers in the card
    total = sum(matches .* card_flat)
    return total*last_number_called
end

function test_card_score()
    card = [1 2 3
            4 5 6
            7 8 9]
    # should be (1+2+3+4)*9 = 90
    @assert get_card_score(card, [5,6,7,8,9]) == 90
    # should be (2+7)*4 = 36
    @assert get_card_score(card, [9,8,1,3,6,23,5,4]) == 36
end

test_card_win()
test_card_score()

numbers = read_numbers("input.txt")
cards = read_bingo_cards("input.txt")

num_cards = length(cards)
card_size = size(cards[1])[1]

have_winning_card = false
for i in card_size:length(numbers)
    local nums_so_far = numbers[1:i]
    for icard in 1:num_cards
        if check_card_win(cards[icard], nums_so_far)
            global have_winning_card = true
            println("Card ",icard, " wins ", nums_so_far)
            println("Winning score ",get_card_score(cards[icard], nums_so_far))
            break
        end
    end
    if have_winning_card
        break
    end
end


# Add tests to the quote for functions with  FastMath varients.
function jacobian_via_frule(f,z)
    du_dx, dv_dx = reim(frule((Zero(), 1),f,z)[2])
    du_dy, dv_dy = reim(frule((Zero(),im),f,z)[2])
    return [
        du_dx  du_dy
        dv_dx  dv_dy
    ]
end
function jacobian_via_rrule(f,z)
    _, pullback = rrule(f,z)
    du_dx, du_dy = reim(pullback( 1)[2])
    dv_dx, dv_dy = reim(pullback(im)[2])
    return [
        du_dx  du_dy
        dv_dx  dv_dy
    ]
end

function jacobian_via_fdm(f, z::Union{Real, Complex})
    fR2((x, y)) = (collect ∘ reim ∘ f)(x + im*y)
    v = float([real(z)
               imag(z)])
    j = jacobian(central_fdm(5,1), fR2, v)[1]
    if size(j) == (2,2)
        j
    elseif size(j) == (1, 2)
        [j
         false false]
    else
        error("Invalid Jacobian size $(size(j))")
    end
end

function complex_jacobian_test(f, z)
    @test jacobian_via_fdm(f, z) ≈ jacobian_via_frule(f, z)
    @test jacobian_via_fdm(f, z) ≈ jacobian_via_rrule(f, z)
end

const FASTABLE_AST = quote
    is_fastmath_mode = sin === Base.FastMath.sin_fast

    @testset "Trig" begin
        @testset "Basics" for x = (Float64(π)-0.01, Complex(π, π/2))
            test_scalar(sin, x)
            test_scalar(cos, x)
            test_scalar(tan, x)
        end
        @testset "Hyperbolic" for x = (Float64(π)-0.01, Complex(π-0.01, π/2))
            test_scalar(sinh, x)
            test_scalar(cosh, x)
            test_scalar(tanh, x)
        end
        @testset "Inverses" for x = (0.5, Complex(0.5, 0.25))
            test_scalar(asin, x)
            test_scalar(acos, x)
            test_scalar(atan, x)
        end
        @testset "Multivariate" begin
            @testset "sincos(x::$T)" for T in (Float64, ComplexF64)
                x, Δx, x̄ = randn(T, 3)
                Δz = (randn(T), randn(T))

                frule_test(sincos, (x, Δx))
                rrule_test(sincos, Δz, (x, x̄))
            end
        end
    end

    @testset "exponents" begin
        for x in (-0.1, 7.9, 0.5 + 0.25im)
            test_scalar(inv, x)

            test_scalar(exp, x)
            test_scalar(exp2, x)
            test_scalar(exp10, x)
            test_scalar(expm1, x)

            if x isa Real
                test_scalar(cbrt, x)
            end

            if x isa Complex || x >= 0
                test_scalar(sqrt, x)
                test_scalar(log, x)
                test_scalar(log2, x)
                test_scalar(log10, x)
                test_scalar(log1p, x)
            end
        end
    end

    @testset "rounding" begin
        for x in (-0.6, -0.2, 0.1, 0.6)
            # thanks to RoundNearest
            if 0 > x % 1 > 0.5 || -1 < x % 1 <= 0.5
                test_scalar(round, x; fdm=backward_fdm(5,1))
            else
                test_scalar(round, x; fdm=forward_fdm(5,1))
            end
            test_scalar(floor, x; fdm=backward_fdm(5, 1))
            test_scalar(ceil, x; fdm=forward_fdm(5, 1))
        end
    end

    @testset "Unary complex functions" begin
        for f ∈ (abs, abs2, conj), z ∈ (-4.1-0.02im, 6.4, 3 + im)
            @testset "Unary complex functions f = $f, z = $z" begin
                complex_jacobian_test(f, z)
            end
        end
        # As per PR #196, angle gives a Zero() pullback for Real z and ΔΩ, rather than
        # the one you'd get from considering the reals as embedded in the complex plane
        # so we need to special case it's tests
        for z ∈ (-4.1-0.02im, 6.4 + 0im, 3 + im)
            complex_jacobian_test(angle, z)
        end
        @test frule((Zero(), randn()), angle, randn())[2] === Zero()
        @test rrule(angle, randn())[2](randn())[2]        === Zero()

        # test that real primal with complex tangent gives complex tangent
        ΔΩ = randn(ComplexF64)
        for x in (-0.5, 2.0)
            @test isapprox(
                frule((Zero(), ΔΩ), angle, x)[2],
                frule((Zero(), ΔΩ), angle, complex(x))[2],
            )
        end
    end

    @testset "Unary functions" begin
        for x in (-4.1, 6.4, 0.0, 0.0 + 0.0im, 0.5 + 0.25im)
            test_scalar(+, x)
            test_scalar(-, x)
            test_scalar(atan, x; rtol=(is_fastmath_mode ? 1e-7 : 1e-9))
        end
    end

    @testset "binary functions" begin
        @testset "$f(x, y)" for f in (atan, rem, max, min)
            x, Δx, x̄ = 100rand(3)
            y, Δy, ȳ = 10rand(3)
            Δz = rand()

            frule_test(f, (x, Δx), (y, Δy))
            rrule_test(f, Δz, (x, x̄), (y, ȳ))
        end

        @testset "$f(x::$T, y::$T)" for f in (/, +, -, hypot), T in (Float64, ComplexF64)
            x, Δx, x̄ = 10rand(T, 3)
            y, Δy, ȳ = rand(T, 3)
            Δz = randn(typeof(f(x, y)))

            # some tests struggle in fast_math mode to get accurasy so we relax it some.
            frule_test(f, (x, Δx), (y, Δy); rtol=(is_fastmath_mode ? 1e-5 : 1e-7))
            rrule_test(f, Δz, (x, x̄), (y, ȳ); rtol=(is_fastmath_mode ? 1e-5 : 1e-7))
        end

        @testset "$f(x::$T, y::$T) type check" for f in (/, +, -,\, hypot, ^), T in (Float32, Float64)
            x, Δx, x̄ = 10rand(T, 3)
            y, Δy, ȳ = rand(T, 3)
            @assert T == typeof(f(x, y))
            Δz = randn(typeof(f(x, y)))

            @test frule((Zero(), Δx, Δy), f, x, y) isa Tuple{T, T}
            _, ∂x, ∂y = rrule(f, x, y)[2](Δz)
            @test extern.((∂x, ∂y)) isa Tuple{T, T}

            if f != hypot
                # Issue #233
                @test frule((Zero(), Δx, Δy), f, x, 2) isa Tuple{T, T}
                _, ∂x, ∂y = rrule(f, x, 2)[2](Δz)
                @test extern.((∂x, ∂y)) isa Tuple{T, T}

                @test frule((Zero(), Δx, Δy), f, 2, y) isa Tuple{T, T}
                _, ∂x, ∂y = rrule(f, 2, y)[2](Δz)
                @test extern.((∂x, ∂y)) isa Tuple{T, T}
            end
        end

        @testset "^(x::$T, n::$T)" for T in (Float64, ComplexF64)
            # for real x and n, x must be >0
            x, Δx, x̄ = rand(T, 3) .+ 3
            y, Δy, ȳ = rand(T, 3) .+ 3
            Δz = rand(T)

            frule_test(^, (x, Δx), (y, Δy))
            rrule_test(^, Δz, (x, x̄), (y, ȳ))

            T <: Real && @testset "discontinuity for ^(x::Real, n::Int) when x ≤ 0" begin
                # finite differences doesn't work for x < 0, so we check manually
                x = -rand(T) .- 3
                y = 3
                Δx = randn(T)
                Δy = randn(T)
                Δz = randn(T)

                @test frule((Zero(), Δx, Δy), ^, x, y)[2] ≈ Δx * y * x^(y - 1)
                @test frule((Zero(), Δx, Δy), ^, zero(x), y)[2] ≈ 0
                _, ∂x, ∂y = rrule(^, x, y)[2](Δz)
                @test ∂x ≈ Δz * y * x^(y - 1)
                @test ∂y ≈ 0
                _, ∂x, ∂y = rrule(^, zero(x), y)[2](Δz)
                @test ∂x ≈ 0
                @test ∂y ≈ 0
            end
        end
    end

    @testset "sign" begin
        @testset "real" begin
            @testset "at $x" for x in (-1.1, -1.1, 0.5, 100.0)
                test_scalar(sign, x)
            end

            @testset "Zero over the point discontinuity" begin
                # Can't do finite differencing because we are lying
                # following the subgradient convention.

                _, pb = rrule(sign, 0.0)
                _, x̄ = pb(10.5)
                @test extern(x̄) == 0

                _, ẏ = frule((Zero(), 10.5), sign, 0.0)
                @test extern(ẏ) == 0
            end
        end
        @testset "complex" begin
            @testset "at $z" for z in (-1.1 + randn() * im, 0.5 + randn() * im)
                test_scalar(sign, z)

                # test that complex (co)tangents with real primal gives same result as
                # complex primal with zero imaginary part

                ż, ΔΩ = randn(ComplexF64, 2)
                Ω, ∂Ω = frule((Zero(), ż), sign, real(z))
                @test Ω == sign(real(z))
                @test ∂Ω ≈ frule((Zero(), ż), sign, real(z) + 0im)[2]

                Ω, pb = rrule(sign, real(z))
                @test Ω == sign(real(z))
                @test pb(ΔΩ)[2] ≈ rrule(sign, real(z) + 0im)[2](ΔΩ)[2]
            end

            @testset "zero over the point discontinuity" begin
                # Can't do finite differencing because we are lying
                # following the subgradient convention.

                _, pb = rrule(sign, 0.0 + 0.0im)
                _, z̄ = pb(randn(ComplexF64))
                @test extern(z̄) == 0.0 + 0.0im

                _, Ω̇ = frule((Zero(), randn(ComplexF64)), sign, 0.0 + 0.0im)
                @test extern(Ω̇) == 0.0 + 0.0im
            end
        end
    end
end

# Now we generate tests for fast and nonfast versions
@eval @testset "fastmath_able Base functions" begin
    $FASTABLE_AST
end

@eval @testset "fastmath_able FastMath functions" begin
    $(Base.FastMath.make_fastmath(FASTABLE_AST))
end


@recipe function f(curve::AbstractManifold{1, C}; npoints=1000) where C
    pts = [curve(t) for t in range(0.0, 0.1, length=npoints)]

    linewidth --> 2
    color     --> :orange
    legend    --> false

    xs = get.(pts, 1, 0.0)
    ys = get.(pts, 2, 0.0)
    zs = get.(pts, 3, 0.0)

    xs, ys, zs
end

@recipe function f(surface::AbstractManifold{2, C}; npoints = 100) where C
    rng = range(0.0, 1.0, length=npoints)

    xs = zeros(npoints, npoints)
    ys = zeros(npoints, npoints)
    zs = zeros(npoints, npoints)
    ws = zeros(npoints, npoints)
    for (i, θ) in enumerate(rng), (j, φ) in enumerate(rng)
        v = surface(θ, φ)
        xs[i, j] = get(v, 1, 0.0)
        ys[i, j] = get(v, 2, 0.0)
        zs[i, j] = get(v, 3, 0.0)
        ws[i, j] = get(v, 4, 0.0)
    end
    seriestype := :surface
    if !iszero(ws)
        surfacecolor := ws
    end

    xs, ys, zs
end

@userplot Points3d
@recipe function f(p::Points3d)
    if length(p.args) != 1 || !(typeof(first(p.args)) <: AbstractVector{<:AbstractVector})
        error("points3d is expecting a vector of vectors. " *
              "Got: $(typeof(p.args))")
    end
    pts = first(p.args)

    seriestype := :scatter
    markersize --> 0.5
    get.(pts, 1, 0.0), get.(pts, 2, 0.0), get.(pts, 3, 0.0)
end


const data_gmm = [
    -1.6518947325446898,
    -0.9012309098240023,
    -0.443220750663359,
]
const s1_gmm = 5.0
const s2_gmm = 1.0

const data_hmm = [
    0.8207426252019316,
    2.081341502733774,
    1.0565299991122412,
]
const s1_hmm = 0.1
const s2_hmm = 0.5

# data for IMM from R. Neal paper:
const data_neal = [-1.48, -1.40, -1.16, -1.08, -1.02, 0.14, 0.51, 0.53, 0.78]
const α_neal = 10.0
const s1_imm = 5.0
const s2_imm = 1.0

####################################################################################################
####################################################################################################
@model function hierarchical_gaussian(x)
    λ ~ Gamma(2.0, inv(3.0))
    m ~ Normal(0, sqrt(1 / λ))
    x ~ Normal(m, sqrt(1 / λ))
end

@model function bernoulli_mixture(x)
    # Mixture prior.
    w ~ Dirichlet(2, 1/2)

    # Latent probability.
    π ~ DiscreteNonParametric([0.3, 0.7], w)

    # Observation.
    x ~ Bernoulli(π)
end

bernoulli_example(;x = false) = bernoulli_mixture(x)

@model function gmm(x, K)
    N = length(x)

    # Cluster association prior.
    w ~ Dirichlet(K, 1/K)

    # Cluster assignments.
    z ~ filldist(DiscreteNonParametric(1:K, w), N)

    # Cluster centers.
    μ ~ filldist(Normal(0.0, s1_gmm), K)

    # Observations.
    for n = 1:N
        x[n] ~ Normal(μ[z[n]], s2_gmm)
    end

    return x
end

gmm_example(;x = data_gmm, K = 2) = gmm(x, K)
gmm_generate(N; kwargs...) = gmm_generate(Random.GLOBAL_RNG, N; kwargs...)
gmm_generate(rng, N; K = 2) =
    :x => identity.(gmm(fill(missing, N), K)(rng, VarInfo(), SampleFromPrior(), DefaultContext()))

@model function gmm_loopy(x, K, ::Type{T}=Float64) where {T<:Real}
    N = length(x)

    # Cluster centers.
    μ = Vector{T}(undef, K)
    for k = 1:K
        μ[k] ~ Normal(0.0, s2_gmm)
    end

    # Cluster association prior.
    w ~ Dirichlet(K, 1/K)

    # Cluster assignments & observations.
    z = Vector{Int}(undef, N)
    for n = 1:N
        z[n] ~ DiscreteNonParametric(1:K, w)
        x[n] ~ Normal(μ[z[n]], s2_gmm)
    end

    return x
end

gmm_loopy_example(;x = data_gmm, K = 2) = gmm_loopy(x, K)

# same as gmm_loopy, but with an affine transformation on μ.
@model function gmm_shifted(x, K, ::Type{T}=Float64) where {T<:Real}
    N = length(x)

    # Cluster centers.
    μ = Vector{T}(undef, K)
    for k = 1:K
        μ[k] ~ Normal(0.0, s2_gmm)
    end

    # Cluster association prior.
    w ~ Dirichlet(K, 1/K)

    # Cluster assignments & observations.
    z = Vector{Int}(undef, N)
    for n = 1:N
        z[n] ~ DiscreteNonParametric(1:K, w)
        x[n] ~ Normal(4μ[z[n]] - 1, s2_gmm)
    end

    return x
end

gmm_shifted_example(;x = data_gmm, K = 2) = gmm_shifted(x, K)

@model function gmm_marginalized(x, K)
    N = length(x)

    # Cluster association prior.
    w ~ Dirichlet(K, 1/K)

    # Cluster centers.
    μ ~ filldist(Normal(0.0, s1_gmm), K)

    # Observations.
    for n = 1:N
        x[n] ~ MixtureModel(Normal, μ, w)
    end

    return x
end

gmm_marginalized_example(;x = data_gmm, K = 2) = gmm_marginalized(x, K)
gmm_marginalized_generate(N; kwargs...) = gmm_marginalized_generate(Random.GLOBAL_RNG, N; kwargs...)
gmm_marginalized_generate(rng, N; K = 2) =
    :x => identity.(gmm_marginalized(fill(missing, N), K)(rng, VarInfo(), SampleFromPrior(), DefaultContext()))

# K clusters, each one around i for i = 1:K with fixed variance
@model function hmm(x, K, ::Type{X}=Float64) where {X<:Real}
    N = length(x)

    # Transition matrix.
    T = Vector{Vector{X}}(undef, K)
    for i = 1:K
        T[i] ~ Dirichlet(K, 1/K)
    end

    # State sequence.
    s = zeros(Int, N)
    s[1] ~ Categorical(K)
    for i = 2:N
        s[i] ~ Categorical(T[s[i-1]])
    end

    # Emission matrix.
    m = Vector{X}(undef, K)
    for i = 1:K
        m[i] ~ Normal(i, s1_hmm)
    end

    # Observe each point of the input.
    x[1] ~ Normal(m[s[1]], s2_hmm)
    for i = 2:N
        x[i] ~ Normal(m[s[i]], s2_hmm)
    end

    return x
end

hmm_example(;x = data_hmm, K = 2) = hmm(x, K)
hmm_generate(N; kwargs...) = hmm_generate(Random.GLOBAL_RNG, N; kwargs...)
hmm_generate(rng, N; K = 2) =
    :x => identity.(hmm(fill(missing, N), K)(rng, VarInfo(), SampleFromPrior(), DefaultContext()))

@model function changepoint(y)
    N = length(y)
    α = 1 / mean(y)
    λ₁ ~ Exponential(α)
    λ₂ ~ Exponential(α)
    τ ~ DiscreteUniform(1, N)

    for n in 1:N
        y[n] ~ Poisson(τ > n ? λ₁ : λ₂)
    end

    return y
end

# @model function reverse_deps(x)
#     m = Vector{Float64}(undef, 2)
#     m[1] ~ Normal()
#     m[2] ~ Normal()
#     x ~ MvNormal(m)
# end

function stickbreak(v)
    K = length(v) + 1
    cumprod_one_minus_v = cumprod(1 .- v)

    return map(1:K) do k
        if k == 1
            v[1]
        elseif k == K
            cumprod_one_minus_v[K - 1]
        else
            v[k] * cumprod_one_minus_v[k - 1]
        end
    end
end

# from https://luiarthur.github.io/TuringBnpBenchmarks/dpsbgmm
@model function imm_stick(y, α, K)
    N = length(y)
    crm = DirichletProcess(α)

    v ~ filldist(StickBreakingProcess(crm), K - 1)
    w = stickbreak(v)

    # Cluster assignments
    z = zeros(Int, N)
    for n = 1:N
        z[n] ~ Categorical(w)
    end

    # Cluster centers
    L = identity(K)
    μ ~ filldist(Normal(0.0, s1_imm), L)

    # Observations
    for n = 1:N
        y[n] ~ Normal(μ[z[n]], s2_imm)
    end

    return y
end

imm_stick_example(;y = data_neal, α = α_neal, K = 10) = imm_stick(y, α, K)
imm_stick_generate(N; kwargs...) = imm_stick_generate(Random.GLOBAL_RNG, N; kwargs...)
imm_stick_generate(rng, N; α = α_neal, K = 10) =
    :y => identity.(imm_stick(fill(missing, N), α, K)(rng, VarInfo(), SampleFromPrior(), DefaultContext()))

####################################################################################################
####################################################################################################
@model function gmm_tarray(x, K)
    N = length(x)

    # Cluster association prior.
    w ~ Dirichlet(K, 1/K)

    # Cluster assignments.
    z = tzeros(Int, N)
    for n = 1:N
        z[n] ~ DiscreteNonParametric(1:K, w)
    end

    # Cluster centers.
    μ ~ filldist(Normal(0.0, s1_gmm), K)

    # Observations.
    for n = 1:N
        x[n] ~ Normal(μ[z[n]], s2_gmm)
    end

    return x
end

gmm_tarray_example(;x = data_gmm, K = 2) = gmm_tarray(x, K)

# K clusters, each one around i for i = 1:K with fixed variance
@model function hmm_tarray(x, K)
    N = length(x)

    # Transition matrix.
    T ~ filldist(Dirichlet(K, 1/K), K)

    # State sequence.
    s = tzeros(Int, N)
    s[1] ~ Categorical(K)
    for i = 2:N
        s[i] ~ Categorical(T[:, s[i-1]])
    end

    # Emission matrix.
    m ~ arraydist(Normal.(1:K, s1_hmm))

    # Observe each point of the input.
    x[1] ~ Normal(m[s[1]], s2_hmm)
    for i = 2:N
        x[i] ~ Normal(m[s[i]], s2_hmm)
    end

    return x
end

hmm_tarray_example(;x = data_hmm, K = 2) = hmm_tarray(x, K)

function stickbreak(v)
    K = length(v) + 1
    cumprod_one_minus_v = cumprod(1 .- v)

    return map(1:K) do k
        if k == 1
            v[1]
        elseif k == K
            cumprod_one_minus_v[K - 1]
        else
            v[k] * cumprod_one_minus_v[k - 1]
        end
    end
end

@model function imm_stick_tarray(y, α, K)
    N = length(y)
    crm = DirichletProcess(α)

    v ~ filldist(StickBreakingProcess(crm), K - 1)
    w = stickbreak(v)

    # Cluster assignments
    z = tzeros(Int, N)
    for n = 1:N
        z[n] ~ Categorical(w)
    end

    # Cluster centers
    L = identity(K)
    μ ~ filldist(Normal(0.0, s1_imm), L)

    # Observations
    for n = 1:N
        y[n] ~ Normal(μ[z[n]], s2_imm)
    end

    return y
end

imm_stick_tarray_example(;y = data_neal, α = α_neal, K = 10) = imm_stick_tarray(y, α, K)


export Polynomial, exponents, coefficients, variables, nvariables, nterms,
    ishomogenized, ishomogenous, homogenize, dehomogenize, degree, evaluate, substitute,
    differentiate, ∇, weylnorm, weyldot, scale_coefficients!

"""
    Polynomial(p::MultivariatePolynomials.AbstractPolynomial [, variables [, homogenized=false]])

A structure for fast evaluation of multivariate polynomials.
The terms are sorted first by total degree, then lexicographically.
`Polynomial` has first class support for [homogenous polynomials](https://en.wikipedia.org/wiki/Homogeneous_polynomial).
This field indicates whether the first variable should be considered as the homogenization variable.

    Polynomial{T}(p::MultivariatePolynomials.AbstractPolynomial [, variables [, homogenized=false]])

You can force a coefficient type `T`. For optimal performance `T` should be same type as
the input to with which it will be evaluated.

    Polynomial(exponents::Matrix{Int}, coefficients::Vector{T}, variables, [, homogenized=false])

You can also create a polynomial directly. Note that in exponents each column represents the exponent of a term.

### Example
```
Poly([3 1; 1 1; 0 2 ], [-2.0, 3.0], [:x, :y, :z]) == 3.0x^2yz^2 - 2x^3y
```
"""
mutable struct Polynomial{T<:Number}
    exponents::Matrix{Int}
    coefficients::Vector{T}
    variables::Vector{Symbol}
    homogenized::Bool

    function Polynomial{T}(exponents::Matrix{Int}, coefficients::Vector{T}, variables::Vector{Symbol}, homogenized::Bool) where {T<:Number}
        sorted_cols = sort!([1:size(exponents,2);], lt=((i, j) -> lt_total_degree(exponents[:,i], exponents[:,j])), rev=true)
        exps = exponents[:, sorted_cols]
        coeffs = coefficients[sorted_cols]
        new(exps, coeffs, copy(variables), homogenized)
    end
end

function Polynomial(exponents::Matrix{Int}, coefficients::Vector{T}, variables::Vector{Symbol}, homogenized::Bool=false) where {T<:Number}
    Polynomial{T}(exponents, coefficients, variables, homogenized)
end

function Polynomial(exponents::Matrix{Int}, coefficients::Vector{T}, homogenized::Bool=false) where {T<:Number}
    Polynomial{T}(exponents, coefficients, [Symbol("x$(i)") for i=1:size(exponents, 1)], homogenized)
end
function Polynomial(p::MP.AbstractPolynomialLike, variables::Vector{<:MP.AbstractVariable}, homogenized::Bool=false)
    exps, coefficients = _coefficients_exponents(p, variables)
    Polynomial(exps, coefficients, Symbol.(variables), homogenized)
end
Polynomial(p::MP.AbstractPolynomialLike, homogenized::Bool=false) = Polynomial(p, MP.variables(p), homogenized)

function Polynomial{T}(p::MP.AbstractPolynomialLike, variables::Vector{<:MP.AbstractVariable}, homogenized::Bool=false) where {T<:Number}
    exps, coefficients = _coefficients_exponents(p, variables)
    Polynomial{T}(exps, convert(Vector{T}, coefficients), Symbol.(variables), homogenized)
end
Polynomial{T}(p::MP.AbstractPolynomialLike, homogenized::Bool=false) where {T<:Number} = Polynomial{T}(p, MP.variables(p), homogenized)

function _coefficients_exponents(poly::MP.AbstractPolynomialLike{T}, vars) where {T}
    terms = MP.terms(poly)
    nterms = length(terms)
    nvars = length(vars)
    exps = Matrix{Int}(undef, nvars, nterms)
    coefficients = Vector{T}(undef, nterms)
    for j = 1:nterms
        term = terms[j]
        coefficients[j] = MP.coefficient(term)
        for i = 1:nvars
            exps[i,j] = MP.degree(term, vars[i])
        end
    end
    exps, coefficients
end

"Sorts two vectory by total degree"
function lt_total_degree(a::Vector{T}, b::Vector{T}) where {T<:Real}
    sum_a = sum(a)
    sum_b = sum(b)
    if sum_a < sum_b
        return true
    elseif sum_a > sum_b
        return false
    else
        for i in eachindex(a)
            if a[i] < b[i]
                return true
            elseif a[i] > b[i]
                return false
            end
        end
    end
    false
end

"""
    exponents(p::Polynomial)

Returns the exponents matrix of `p`. Each column represents the exponents of a term of `p`.
"""
exponents(p::Polynomial) = p.exponents

"""
    coefficients(p::Polynomial)

Returns the coefficient vector of `p`.
"""
coefficients(p::Polynomial) = p.coefficients

"""
    variables(p::Polynomial)

Returns the variables of `p`.
"""
variables(p::Polynomial) = p.variables

"""
    ishomogenized(p::Polynomial)

Checks whether `p` was homogenized.
"""
ishomogenized(p::Polynomial) = p.homogenized

"""
    nterms(p::Polynomial)

Returns the number of terms of p
"""
nterms(p::Polynomial) = size(exponents(p), 2)

"""
    nvariables(p::Polynomial)

Returns the number of variables of p
"""
nvariables(p::Polynomial) = size(exponents(p), 1)

"""
    degree(p::Polynomial)

Returns the (total) degree of p.
"""
degree(p::Polynomial) = sum(exponents(p)[:,1])

==(p::Polynomial, q::Polynomial) = exponents(p) == exponents(q) && coefficients(p) == coefficients(q)
Base.isequal(p::Polynomial, q::Polynomial) = exponents(p) == exponents(q) && coefficients(p) == coefficients(q)

# ITERATOR
function Base.iterate(p::Polynomial, state...)
    n = nterms(p)
    istate = iterate(1:n, state...)
    istate === nothing && return nothing

    i, state = istate
    (coefficients(p)[i], exponents(p)[:,i]), state
end
Base.length(p::Polynomial) = nterms(p)
Base.eltype(p::Polynomial{T}) where {T} = T

Base.broadcastable(p::Polynomial) = Ref(p)

@inline pow(x::AbstractFloat, k::Integer) = Base.FastMath.pow_fast(x, k)
#@inline pow(x::Complex, k::Integer) = k == 1 ? x : x^k
# simplified from Base.power_by_squaring
@inline function pow(x::Number, p::Integer)
    if p == 1
        return copy(x)
    elseif p == 0
        return one(x)
    elseif p == 2
        return x*x
    end
    t = trailing_zeros(p) + 1
    p >>= t
    while (t -= 1) > 0
        x *= x
    end
    y = x
    while p > 0
        t = trailing_zeros(p) + 1
        p >>= t
        while (t -= 1) >= 0
            x *= x
        end
        y *= x
    end
    return y
end

"""
    evaluate(p::Polynomial{T}, x::AbstractVector{T})

Evaluates `p` at `x`, i.e. ``p(x)``.
`Polynomial` is also callable, i.e. you can also evaluate it via `p(x)`.
"""
function evaluate(p::Polynomial{S}, x::AbstractVector{T}) where {S<:Number, T<:Number}
    cfs = coefficients(p)
    exps = exponents(p)
    nvars, nterms = size(exps)
    res = zero(promote_type(S,T))
    for j = 1:nterms
        term = p.coefficients[j]
        for i = 1:nvars
            k = exps[i, j]
            term *= pow(x[i], k)
        end
        res += term
    end
    res
end

(p::Polynomial)(x) = evaluate(p, x)

evaluate(F::Vector{<:Polynomial}, x::AbstractVector) = map(f -> evaluate(f, x), F)
"""
    substitute(p::Polynomial, i, x)

In `p` substitute for the variable with index `i` the value `x`. You can use this for partial
evaluation of polynomial.

### Example
```julia-repl
julia> substitute(x^2+3y, 2, 5)
x^2+15
```
"""
function substitute(p::Polynomial{S}, varindex, x::T) where {S<:Number, T<:Number}
    cfs = coefficients(p)
    exps = exponents(p)
    nvars, nterms = size(exps)

    new_coefficients = Vector{promote_type(S,T)}()
    new_exps = Vector{Vector{Int}}()

    for j = 1:nterms
        coeff = cfs[j]
        exp = Vector{Int}(undef, nvars - 1)
        # first we calculate the new coefficient and remove the varindex-th row
        for i = 1:nvars
            if i == varindex
                coeff *= x^(exps[i, j])
            elseif i > varindex
                exp[i-1] = exps[i, j]
            else
                exp[i] = exps[i, j]
            end
        end
        # now we have to delete possible duplicates
        found_duplicate = false
        for k = 1:length(new_exps)
            if new_exps[k] == exp
                new_coefficients[k] += coeff
                found_duplicate = true
                break
            end
        end
        if !found_duplicate
            push!(new_coefficients, coeff)
            push!(new_exps, exp)
        end
    end

    # now we have to create a new matrix and return the poly
    Polynomial(hcat(new_exps...), new_coefficients, [variables(p)[1:varindex-1]; variables(p)[varindex+1:end]], p.homogenized)
end

"""
    differentiate(p::Polynomial, varindex::Int)

Differentiate `p` w.r.t the `varindex`th variable.

    differentiate(p::Polynomial)

Differentiate `p` w.r.t. all variables.
"""
function differentiate(p::Polynomial{T}, i_var) where T
    exps = copy(exponents(p))
    cfs = copy(coefficients(p))
    n_vars, n_terms = size(exps)

    zerocolumns = Int[]
    for j=1:n_terms
        k = exps[i_var, j]
        if k > 0
            exps[i_var, j] = max(0, k - 1)
            cfs[j] *= k
        else
            push!(zerocolumns , j)
        end
    end

    # now we have to get rid of all zeros
    nzeros = length(zerocolumns)
    if nzeros == 0
        return Polynomial(exps, cfs, variables(p), p.homogenized)
    end

    skipped_cols = 0
    new_exps = zeros(Int, n_vars, n_terms - nzeros)
    new_coefficients = zeros(T, n_terms - nzeros)

    for j = 1:n_terms
        # if we not yet have skipped all zero columns
        if skipped_cols < nzeros && j == zerocolumns[skipped_cols+1]
            skipped_cols += 1
            continue
        end

        new_exps[:, j - skipped_cols] = exps[:, j]
        new_coefficients[j - skipped_cols] = cfs[j]
    end

    Polynomial(new_exps, new_coefficients, variables(p), p.homogenized)
end
differentiate(poly::Polynomial) = map(i -> differentiate(poly, i), 1:nvariables(poly))

"""
    ∇(p::Polynomial)

Returns the gradient vector of `p`. This is the same as [`differentiate`](@ref).
"""
@inline ∇(poly::Polynomial) = differentiate(poly)

"""
    ishomogenous(p::Polynomial)

Checks whether `p` is a homogenous polynomial. Note that this is unaffected from the
value of `homogenized(p)`.
"""
function ishomogenous(p::Polynomial)
    monomials_degree = sum(exponents(p), dims=1)
    max_deg = monomials_degree[1]
    all(x -> x == max_deg, monomials_degree)
end

"""
    homogenize(p::Polynomial [, variable = :x0])

Makes `p` homogenous, if `ishomogenized(p)` is `true` this is just the identity.
The homogenization variable will always be considered as the first variable of the polynomial.
"""
function homogenize(p::Polynomial, variable::Symbol=:x0; respect_homogenous=true)
    if p.homogenized || (respect_homogenous && ishomogenous(p))
        p
    else
        monomials_degree = sum(exponents(p), dims=1)
        max_deg = monomials_degree[1]
        Polynomial([max_deg .- monomials_degree; exponents(p)], coefficients(p), [variable; variables(p)], true)
    end
end

"""
    dehomogenize(p::Polynomial)

Substitute `1` as for the first variable `p`, if `ishomogenized(p)` is `false` this is just
the identity.
"""
function dehomogenize(p::Polynomial)
    if !p.homogenized
        p
    else
        Polynomial(exponents(p)[2:end,:], coefficients(p), variables(p)[2:end], false)
    end
end

"Computes the multinomial coefficient (|k| \\over k)"
function multinomial(k::Vector{Int})
    s = 0
    result = 1
    @inbounds for i in k
        s += i
        result *= binomial(s, i)
    end
    result
end

"""
    weyldot(f::Polynomial, g::Polynomial)

Compute the [Bombieri-Weyl dot product](https://en.wikipedia.org/wiki/Bombieri_norm).
Note that this is only properly defined if `f` and `g` are homogenous.

    weyldot(f::Vector{Polynomial}, g::Vector{Polynomial})

Compute the dot product for vectors of polynomials.
"""
function weyldot(f::Polynomial,g::Polynomial)
    if (f === g)
        return sum(x -> abs2(x[1]) / multinomial(x[2]), f)
    end
    result = 0
    for (c_f, exp_f) in f
        normalizer = multinomial(exp_f)
        for (c_g, exp_g) in g
            if exp_f == exp_g
                result += (c_f * conj(c_g)) / normalizer
                break
            end
        end
    end
    result
end

function weyldot(F::Vector{Polynomial{T}}, G::Vector{Polynomial{S}}) where {T,S}
    res = zero(promote_type(S, T))
    for (f, g) in zip(F, G)
        res += weyldot(f, g)
    end
    res
end

"""
    weylnorm(f::Polynomial)

Compute the [Bombieri-Weyl norm](https://en.wikipedia.org/wiki/Bombieri_norm).
Note that this is only properly defined if `f` is homogenous.
"""
weylnorm(f::Polynomial) = √weyldot(f,f)
weylnorm(f::Vector{Polynomial{T}}) where T = √weyldot(f,f)

"""
    scale_coefficients!(f::Polynomial, λ)

Scale the coefficients of `f` with the factor `λ`.
"""
function scale_coefficients!(f::Polynomial, λ)
    f.coefficients .*= λ
end


### A Pluto.jl notebook ###
# v0.16.0

using Markdown
using InteractiveUtils

# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error).
macro bind(def, element)
    quote
        local el = $(esc(element))
        global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : missing
        el
    end
end

# ╔═╡ 5b918be6-0003-11ec-1e1f-534faed7dfe5
begin
    using Pkg
    Pkg.activate("..")
    using Revise,
        Javis, Distributions, StatsBase, JUtils, Random, Colors, FileIO, PlutoUI, Plots
    import Luxor
end

# ╔═╡ 46fd9cf5-eae0-4dec-91df-ebd892cbf4d6
converter = Dict(["bernoulli" => Bernoulli, "pareto" => Pareto]);

# ╔═╡ b6e3b06f-a521-40fb-9d5a-9268302cd17b
md"""
Select the distribution $(@bind dist Select(collect(keys(converter)), default="bernoulli"))
"""

# ╔═╡ d3ed0756-5fdf-4eba-a8ed-7d944acc079d
md"""
Select the distribution Parameters $(@bind param NumberField(0.0:0.1:5.0, default=0.3))
"""

# ╔═╡ e8fef37b-e0a7-42f3-89a7-414d61092c6e
md"""
Check to start rendering gif$(@bind start CheckBox())
"""

# ╔═╡ e725b87f-fb23-472d-84d7-8e6a8dd01aae
param_dist = converter[dist](param)

# ╔═╡ 47bcc4fe-a44f-4715-8ed3-bbe00c02f30f
begin
    start && JUtils.CLT(
        param_dist,
        boundingbox = BoundingBox(O + (-250, -300), O + (250, 200)),
        pathname = "../output/central_limit_theorem.gif",
    )
end

# ╔═╡ Cell order:
# ╟─5b918be6-0003-11ec-1e1f-534faed7dfe5
# ╟─46fd9cf5-eae0-4dec-91df-ebd892cbf4d6
# ╟─b6e3b06f-a521-40fb-9d5a-9268302cd17b
# ╟─d3ed0756-5fdf-4eba-a8ed-7d944acc079d
# ╟─e8fef37b-e0a7-42f3-89a7-414d61092c6e
# ╟─e725b87f-fb23-472d-84d7-8e6a8dd01aae
# ╟─47bcc4fe-a44f-4715-8ed3-bbe00c02f30f


@testset "Test LAStopping" begin
###############################################################################
#
# The Stopping structure eases the implementation of algorithms and the
# stopping criterion.
#
# The following examples illustrate a specialized Stopping for the
# LinearOperators type.
# https://github.com/JuliaSmoothOptimizers/LinearOperators.jl
#
using LinearAlgebra, SparseArrays, LinearOperators #v.1.0.1

##############################################################
#TEST

n = 5
A = rand(n,n)
xref = 100 * rand(n)
x0 = zeros(n)
b    = A * xref
sA  = sparse(A)
opA = LinearOperator(A)

mLO = LinearSystem(A, b)
sLO = LLSModel(sA, b)
opLO = LinearSystem(opA, b)
meta = StoppingMeta(max_cntrs = init_max_counters_linear_operators())
mLOstp_meta = LAStopping(mLO, meta, GenericState(x0)) #this is different because of optimality_check
mLOstp = LAStopping(mLO, GenericState(x0), max_cntrs = init_max_counters_linear_operators())
sLOstp = LAStopping(sLO, GenericState(x0))
short_stop = LAStopping(A,b, sparse = true) #note that sparse is true by default
maxcn = init_max_counters_linear_operators(nprod = 1)
opLOstp = LAStopping(opLO, GenericState(x0), max_cntrs = maxcn)

opLOstp2 = LAStopping(opLO, GenericState(x0), optimality_check = linear_system_check)
@test opLOstp2.meta.max_cntrs == init_max_counters_linear_operators()
@test opLOstp2.meta.optimality_check == opLOstp.meta.optimality_check

mLOstp_src = LAStopping(mLO, meta, StopRemoteControl(), GenericState(x0))

@test typeof(mLOstp.pb.counters) <: LACounters
@test typeof(sLOstp.pb.counters) <: NLSCounters
@test typeof(opLOstp.pb.counters) <: LACounters
@test typeof(short_stop.pb.counters) <: NLSCounters
@test typeof(short_stop.pb) <: LLSModel

@test start!(mLOstp) == false
@test start!(sLOstp) == false
@test start!(opLOstp) == false

update_and_stop!(mLOstp, x = xref, res = Float64[])
@test status(mLOstp) == :Optimal
@test mLOstp.meta.resources == false
update_and_stop!(sLOstp, x = xref, res = Float64[])
@test status(sLOstp) == :Optimal
@test sLOstp.meta.resources == false
update_and_stop!(opLOstp, x = xref, res = Float64[])
@test status(opLOstp) == :Optimal
@test opLOstp.meta.resources == true

include("_additional_linear_operator_algebra.jl")

"""
Randomized block Kaczmarz
"""
function RandomizedBlockKaczmarz(stp :: AbstractStopping; kwargs...)

    b   = stp.pb.b
    x0  = stp.current_state.x
    m, n = length(b), length(x0)
    A   = try
           stp.pb.A #LSSModel.A does not exist
       catch
           sparse(stp.pb.Arows, stp.pb.Acols, stp.pb.Avals, m, n)
       end

    xk  = x0

    OK = start!(stp)

    while !OK

     i  = Int(floor(rand() * m)+1) #rand a number between 1 and m
     Ai = A[i,:]
     xk = Ai == 0 ? x0 : x0 - (dot(Ai,x0)-b[i])/dot(Ai,Ai) * Ai
     #xk  = Ai == 0 ? x0 : x0 - (Ai' * x0-b[i])/(Ai' * Ai) * Ai

     if xtype(stp.current_state) <: SparseVector && !(typeof(xk) <: SparseVector)
        OK = update_and_stop!(stp, x = sparse(xk), res = sparse(Float64[]))
     elseif xtype(stp.current_state) <: Vector && typeof(xk) <: SparseVector
        OK = update_and_stop!(stp, x = convert(Vector, xk), res = Float64[])
     else
        OK = update_and_stop!(stp, x = xk, res = Stopping._init_field(typeof(stp.current_state.res)))
     end
     x0  = xk

    end

 return stp
end

m, n = 200, 100 #size of A: m x n
A    = 100 * rand(m, n) #It's a dense matrix :)
xref = 100 * rand(n)
b    = A * xref

#Our initial guess
x0 = zeros(n)

la_stop = LAStopping(A, b, GenericState(x0),
                     max_iter = 150000,
                     rtol = 1e-6,
                     max_cntrs = init_max_counters_NLS(residual = 150000))
#Be careful using GenericState(x0) would not work here without forcing convert = true
#in the update function. As the iterate will be a SparseVector to the contrary of initial guess.
#Tangi: maybe start! should send a Warning for such problem !?
sa_stop = LAStopping(sparse(A), b, GenericState(sparse(x0)),
                     max_iter = 150000,
                     rtol = 1e-6,
                     max_cntrs = init_max_counters_NLS(residual = 150000))
op_stop = LAStopping(LinearSystem(LinearOperator(A), b),
                     GenericState(x0),
                     max_iter = 150000,
                     rtol = 1e-6,
                     max_cntrs = init_max_counters_linear_operators(nprod = 150000))
opbis_stop = LAStopping(LinearOperator(A), b)

try
 @timed RandomizedBlockKaczmarz(la_stop)
 @test status(la_stop) == :Optimal
 @timed RandomizedBlockKaczmarz(sa_stop)
 @test status(sa_stop) == :Optimal
catch
    @warn "If LSSModel.A does not exist consider [la_stop.pb.Avals[i,j] for (i) in la_stop.pb.Arows, j in la_stop.pb.Acols]"
    #https://github.com/JuliaSmoothOptimizers/NLPModels.jl/blob/master/src/lls_model.jl
end

@timed RandomizedBlockKaczmarz(op_stop)
@test status(op_stop) == :Optimal

update!(la_stop.current_state, x = xref)
@test normal_equation_check(la_stop.pb, la_stop.current_state) <= la_stop.meta.atol
update!(op_stop.current_state, x = xref)
@test normal_equation_check(op_stop.pb, op_stop.current_state) <= la_stop.meta.atol

end


@testset "Normalized Cut" begin

    gx = SimpleGraph(6)
    add_edge!(gx, 1, 2)
    add_edge!(gx, 2, 3)
    add_edge!(gx, 1, 3)
    add_edge!(gx, 4, 5)
    add_edge!(gx, 4, 6)
    add_edge!(gx, 5, 6)
    add_edge!(gx, 1, 6)
    add_edge!(gx, 3, 4)

    w = zeros(6, 6)
    w[2, 1] = 1.0
    w[3, 1] = 1.0
    w[6, 1] = 0.1
    w[1, 2] = 1.0
    w[3, 2] = 1.0
    w[1, 3] = 1.0
    w[2, 3] = 1.0
    w[4, 3] = 0.2
    w[3, 4] = 0.2
    w[5, 4] = 1.0
    w[6, 4] = 1.0
    w[4, 5] = 1.0
    w[6, 5] = 1.0
    w[1, 6] = 0.1
    w[4, 6] = 1.0
    w[5, 6] = 1.0

    for g in testgraphs(gx)
      labels = @inferred(normalized_cut(g, 1, w))
      @test labels == [1, 1, 1, 2, 2, 2] || labels == [2, 2, 2, 1, 1, 1]
    end

    w = SparseMatrixCSC(w)
    for g in testgraphs(gx)
      labels = @inferred(normalized_cut(g, 1, w))
      @test labels == [1, 1, 1, 2, 2, 2] || labels == [2, 2, 2, 1, 1, 1]
    end

    gx = SimpleGraph(4)
    add_edge!(gx, 1, 2)
    add_edge!(gx, 3, 4)

    w = zeros(4, 4)
    w[2, 1] = 1.0
    w[1, 2] = 1.0
    w[4, 3] = 1.0
    w[3, 4] = 1.0
    for g in testgraphs(gx)
      labels = @inferred(normalized_cut(g, 0.1, w))
      @test labels == [1, 1, 2, 2] || labels == [2, 2, 1, 1]
    end

    w = SparseMatrixCSC(w)
    for g in testgraphs(gx)
      labels = @inferred(normalized_cut(g, 0.1, w))
      @test labels == [1, 1, 2, 2] || labels == [2, 2, 1, 1]
    end

    w = ones(12, 12)
    for g in testgraphs(gx)
      labels = @inferred(normalized_cut(g, 0.1, w))
      @test labels == [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
    end

    w = ones(12, 12)
    for g in testgraphs(gx)
      labels = @inferred(normalized_cut(g, 0.1, w))
      @test labels == [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
    end

    g = PathGraph(50)

    function contiguous(labels::Vector{Int})::Bool
        changes = 0
        for i in 1:length(labels)-1
            if labels[i] != labels[i+1]
                changes += 1
            end
        end
        return changes == length(unique(labels)) - 1
    end

    num_subgraphs = Vector{Int}(undef, 9)

    for t in [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
        labels = @inferred(normalized_cut(g, t))
        @test contiguous(labels) == true
        num_subgraphs[convert(Int, 10*t)] = size(unique(labels), 1)
    end

    @test issorted(num_subgraphs) == true

    @test any(length(unique(@show normalized_cut(g,t))) == 4 for t in [0.125, 0.15, 0.16, 0.175, 0.20])

end


using Test
using FinancialToolbox,DualNumbers

#Test Parameters
testToll=1e-14;

#print_colored("Starting Standard Test\n",:green)
spot=10;K=10;r=0.02;T=2.0;sigma=0.2;d=0.01;
spotDual=dual(spot,1.0);

toll=1e-8

#EuropeanCall Option
PriceCall=blsprice(spot,K,r,T,sigma,d);
PriceCallBlack=blkprice(spot,K,r,T,sigma);

df(f,x,dx)=(f(x+dx)-f(x))/dx
dx=1e-8;
assert_(value,toll)=@test abs(value)<toll

sigma1=blsimpv(spot, K, r, T, PriceCall, d);
ResDual=blsimpv(spotDual, K, r, T, PriceCall, d);
DerDF_=df(spot->blsimpv(spot, K, r, T, PriceCall, d),spot,dx);

assert_(sigma1-ResDual,toll)
assert_(DerDF_-ResDual.epsilon,toll)

sigma1=blkimpv(spot, K, r, T, PriceCallBlack);
ResDual=blkimpv(spotDual, K, r, T, PriceCallBlack);
DerDF_=df(spot->blkimpv(spot, K, r, T, PriceCallBlack),spot,dx);

assert_(sigma1-ResDual,toll)
assert_(DerDF_-ResDual.epsilon,toll)

print_colored("Starting DualNumbers Test\n",:green)

#Test Parameters
testToll=1e-14;
spot=10;K=10;r=0.02;T=2.0;sigma=0.2;d=0.01;

#EuropeanCall Option
PriceCall=blsprice(spot,K,r,T,sigma,d);
DeltaCall=blsdelta(spot,K,r,T,sigma,d);
ThetaCall=blstheta(spot,K,r,T,sigma,d);
RhoCall=blsrho(spot,K,r,T,sigma,d);
LambdaCall=blslambda(spot,K,r,T,sigma,d);
VannaCall=blsvanna(spot,K,r,T,sigma,d);

#EuropeanPut Option
PricePut=blsprice(spot,K,r,T,sigma,d,false);
DeltaPut=blsdelta(spot,K,r,T,sigma,d,false);
ThetaPut=blstheta(spot,K,r,T,sigma,d,false);
RhoPut=blsrho(spot,K,r,T,sigma,d,false);
LambdaPut=blslambda(spot,K,r,T,sigma,d,false);
VannaPut=blsvanna(spot,K,r,T,sigma,d,false);

#Equals for both Options
Gamma=blsgamma(spot,K,r,T,sigma,d);
Vega=blsvega(spot,K,r,T,sigma,d);

########DUAL NUMBERS
DerToll=1e-13;
#Function definition
#Call
FcallDual(spot)=blsprice(spot,K,r,T,sigma,d);
GcallDual(r)=blsprice(spot,K,r,T,sigma,d);
HcallDual(T)=blsprice(spot,K,r,T,sigma,d);
LcallDual(sigma)=blsprice(spot,K,r,T,sigma,d);
PcallDual(spot)=blsprice(spot,K,r,T,sigma,d)*spot.value/blsprice(spot.value,K,r,T,sigma,d);
VcallDual(sigma)=blsdelta(spot,K,r,T,sigma,d);
#Put
FputDual(spot)=blsprice(spot,K,r,T,sigma,d,false);
GputDual(r)=blsprice(spot,K,r,T,sigma,d,false);
HputDual(T)=blsprice(spot,K,r,T,sigma,d,false);
PputDual(spot)=blsprice(spot,K,r,T,sigma,d,false)*spot.value/blsprice(spot.value,K,r,T,sigma,d,false);
VPutDual(sigma)=blsdelta(spot,K,r,T,sigma,d,false);
#Input
SpotDual=Dual(spot,1.0);
rDual=Dual(r,1.0);
TDual=Dual(T,1.0);
SigmaDual=Dual(sigma,1.0);

#Automatic Differentiation Test
#TEST
print_colored("--- European Call Sensitivities: DualNumbers\n",:yellow)
print_colored("-----Testing Delta\n",:blue);
assert_(FcallDual(SpotDual).epsilon-DeltaCall,DerToll)
print_colored("-----Testing Rho\n",:blue);
assert_(GcallDual(rDual).epsilon-RhoCall,DerToll)
print_colored("-----Testing Theta\n",:blue);
assert_(-HcallDual(TDual).epsilon-ThetaCall,DerToll)
print_colored("-----Testing Lambda\n",:blue);
assert_(PcallDual(SpotDual).epsilon-LambdaCall,DerToll)
print_colored("-----Testing Vanna\n",:blue);
assert_(VcallDual(SigmaDual).epsilon-VannaCall,DerToll)
#TEST
print_colored("--- European Put Sensitivities: DualNumbers\n",:yellow)
print_colored("-----Testing Delta\n",:blue);
assert_(FputDual(SpotDual).epsilon-DeltaPut,DerToll)
print_colored("-----Testing Rho\n",:blue);
assert_(GputDual(rDual).epsilon-RhoPut,DerToll)
print_colored("-----Testing Theta\n",:blue);
assert_(-HputDual(TDual).epsilon-ThetaPut,DerToll)
print_colored("-----Testing Lambda\n",:blue);
assert_(PputDual(SpotDual).epsilon-LambdaPut,DerToll)
print_colored("-----Testing Vanna\n",:blue);
assert_(VPutDual(SigmaDual).epsilon-VannaPut,DerToll)

print_colored("-----Testing Vega\n",:blue);
assert_(LcallDual(SigmaDual).epsilon-Vega,DerToll)
print_colored("Dual Numbers Test Passed\n",:green)
println("")

#TEST OF INPUT VALIDATION
print_colored("Starting Input Validation Test Dual\n",:magenta)
print_colored("----Testing Negative  Spot Price \n",:cyan)
@test_throws(ErrorException, blsprice(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blkprice(-SpotDual,K,r,T,sigma))
@test_throws(ErrorException, blsdelta(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blsgamma(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blstheta(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blsrho(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blspsi(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blslambda(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blsvanna(-SpotDual,K,r,T,sigma,d))
@test_throws(ErrorException, blsvega(-SpotDual,K,r,T,sigma,d))

print_colored("----Testing Negative  Strike Price \n",:cyan)
KK=Dual(K,0);
@test_throws(ErrorException, blsprice(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blkprice(spot,-KK,r,T,sigma))
@test_throws(ErrorException, blsdelta(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blsgamma(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blstheta(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blsrho(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blspsi(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blslambda(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blsvanna(spot,-KK,r,T,sigma,d))
@test_throws(ErrorException, blsvega(spot,-KK,r,T,sigma,d))

print_colored("----Testing Negative  Time to Maturity \n",:cyan)
@test_throws(ErrorException, blsprice(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blkprice(spot,K,r,-TDual,sigma))
@test_throws(ErrorException, blsdelta(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blsgamma(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blstheta(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blsrho(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blspsi(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blslambda(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blsvanna(spot,K,r,-TDual,sigma,d))
@test_throws(ErrorException, blsvega(spot,K,r,-TDual,sigma,d))

print_colored("----Testing Negative  Volatility \n",:cyan)
@test_throws(ErrorException, blsprice(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blkprice(spot,K,r,T,-SigmaDual))
@test_throws(ErrorException, blsdelta(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blsgamma(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blstheta(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blsrho(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blspsi(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blslambda(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blsvanna(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blsvega(spot,K,r,T,-SigmaDual,d))
@test_throws(ErrorException, blsimpv(spotDual,K,r,T,-0.2,d))

print_colored("Dual Input Validation Test Passed\n",:magenta)


using MLJ
using PyPlot
using PrettyPrinting
using StableRNGs
import DataFrames: DataFrame, describe

X, y = @load_boston
sch = schema(X)
p = length(sch.names)
describe(y)  # From DataFrames

DecisionTreeRegressor = @load DecisionTreeRegressor pkg=DecisionTree

tree = machine(DecisionTreeRegressor(), X, y)
e = evaluate!(tree, resampling=Holdout(fraction_train=0.8),
              measure=[rms, rmslp1])
e

forest = EnsembleModel(model=DecisionTreeRegressor())
forest.model.n_subfeatures = 3

rng = StableRNG(5123) # for reproducibility
m = machine(forest, X, y)
r = range(forest, :n, lower=10, upper=1000)
curves = learning_curve!(m, resampling=Holdout(fraction_train=0.8, rng=rng),
                         range=r, measure=rms);

figure(figsize=(8,6))
plot(curves.parameter_values, curves.measurements)
ylabel("Root Mean Squared error", fontsize=16)
xlabel("Number of trees", fontsize=16)
xticks([10, 100, 250, 500, 750, 1000], fontsize=14)
yticks(fontsize=14)

forest.n = 150;

params(forest) |> pprint

r_sf = range(forest, :(model.n_subfeatures), lower=1, upper=12)
r_bf = range(forest, :bagging_fraction, lower=0.4, upper=1.0);

tuned_forest = TunedModel(model=forest,
                          tuning=Grid(resolution=3),
                          resampling=CV(nfolds=6, rng=StableRNG(32)),
                          ranges=[r_sf, r_bf],
                          measure=rms)
m = machine(tuned_forest, X, y)
e = evaluate!(m, resampling=Holdout(fraction_train=0.8),
              measure=[rms, rmslp1])
e

r = report(m)

figure(figsize=(8,6))

res = r.plotting

vals_sf = res.parameter_values[:, 1]
vals_bf = res.parameter_values[:, 2]

tricontourf(vals_sf, vals_bf, res.measurements)
xticks(1:3:12, fontsize=12)
xlabel("Number of sub-features", fontsize=14)
yticks(0.4:0.2:1, fontsize=12)
ylabel("Bagging fraction", fontsize=14)

ŷ = predict(m, X)
@show rms(ŷ, y)

# This file was generated using Literate.jl, https://github.com/fredrikekre/Literate.jl


"""
    closest(A, x)

Return the index of the array entry closest in value to `x`. If there are multiple closest entries,
then the first one will be returned.

# Example

```julia
julia> closest([-1, -0.5, 0.5], 0)
2

julia> closest([1 2 3
                4 5 6], 5.8)
CartesianIndex(2, 3)
```
"""
closest(A::AbstractArray, x) = findmin(abs.(A.-x))[2]

logrange(x1, x2, n=50) = (10^y for y in range(log10(x1), log10(x2), length=n))


"""
```
imgb = binarize(Polysegment(), img)
binarize!(Polysegment(), img)
```

Uses the *polynomial segmentation* technique to group the image pixels
into two categories (foreground and background).

# Details

The approach involves constructing a univariate second-degree polynomial
such that the two roots of the polynomial represent the graylevels
of two cluster centers (i.e the foreground and background). Pixels are then
assigned to the foreground or background depending on which cluster center
is closest.

# Arguments

The function arguments are described in more detail below.

##  `img`

An `AbstractGray` image which is to be binarized into background (`Gray(0)`)
and foreground (`Gray(1)`) pixels.

## Reference

[1] R. E. Vidal, "Generalized Principal Component Analysis (GPCA): An Algebraic Geometric Approach to Subspace Clustering and Motion Segmentation." Order No. 3121739, University of California, Berkeley, Ann Arbor, 2003.
"""
function binarize(algorithm::Polysegment,  img::AbstractArray{T,2}) where T <: AbstractGray
  binarize!(algorithm, copy(img))
end

"""
```
t = binarize!(Polysegment(), img)
```
Same as [`binarize`](@ref binarize(::Polysegment, ::AbstractArray{<:AbstractGray,2})) except that it modifies the image that was passed as an argument.
"""
function binarize!(algorithm::Polysegment,  img::AbstractArray{T,2}) where T <: AbstractGray
  # Construct data matrix for second-degree polynomial (Equation 2.3) in [1].
  x = vec(img)
  Lₙ = hcat(ones(length(x)), x, x.^2)
  F = svd(Lₙ)
  c = F.Vt[end,:]
  p = Poly(vec(c))
  μ₁, μ₂ = roots(p)
  # Binarize the image.
  map!(img,img) do val
    (val-μ₁)^2 < (val-μ₂)^2 ? Gray(0) : Gray(1)
  end
end


# This tutorial uses the TMAR Filter from
#
#    @article{doi:10.1175/MWR-D-16-0220.1,
#      author = {Light, Devin and Durran, Dale},
#      title = {Preserving Nonnegativity in Discontinuous Galerkin
#               Approximations to Scalar Transport via Truncation and Mass
#               Aware Rescaling (TMAR)},
#      journal = {Monthly Weather Review},
#      volume = {144},
#      number = {12},
#      pages = {4771-4786},
#      year = {2016},
#      doi = {10.1175/MWR-D-16-0220.1},
#    }
#
# to reproduce the tutorial in section 4b.  It is a shear swirling
# flow deformation of a transported quantity from LeVeque 1996.  The exact
# solution at the final time is the same as the initial condition.

using MPI
using Test
using ClimateMachine
ClimateMachine.init()
using Logging
using ClimateMachine.Mesh.Topologies
using ClimateMachine.Mesh.Grids
using ClimateMachine.Mesh.Filters
using ClimateMachine.DGMethods
using ClimateMachine.DGMethods.NumericalFluxes
using ClimateMachine.MPIStateArrays
using ClimateMachine.ODESolvers
using LinearAlgebra
using Printf
using Dates
using ClimateMachine.GenericCallbacks:
    EveryXWallTimeSeconds, EveryXSimulationSteps
using ClimateMachine.VTK: writevtk, writepvtu

using ClimateMachine
const clima_dir = dirname(dirname(pathof(ClimateMachine)));
include(joinpath(
    clima_dir,
    "test",
    "Numerics",
    "DGMethods",
    "advection_diffusion",
    "advection_diffusion_model.jl",
))

Base.@kwdef struct SwirlingFlow{FT} <: AdvectionDiffusionProblem
    period::FT = 5
end

init_velocity_diffusion!(::SwirlingFlow, aux::Vars, geom::LocalGeometry) =
    nothing

cosbell(τ, q) = τ ≤ 1 ? ((1 + cospi(τ)) / 2)^q : zero(τ)

function initial_condition!(::SwirlingFlow, state, aux, coord, t)
    FT = eltype(state)
    x, y, _ = aux.coord
    x0, y0 = FT(1 // 4), FT(1 // 4)
    τ = 4 * hypot(x - x0, y - y0)
    state.ρ = cosbell(τ, 3)
end;

has_variable_coefficients(::SwirlingFlow) = true
function update_velocity_diffusion!(
    problem::SwirlingFlow,
    ::AdvectionDiffusion,
    state::Vars,
    aux::Vars,
    t::Real,
)
    x, y, _ = aux.coord
    sx, cx = sinpi(x), cospi(x)
    sy, cy = sinpi(y), cospi(y)
    ct = cospi(t / problem.period)

    u = 2 * sx^2 * sy * cy * ct
    v = -2 * sy^2 * sx * cx * ct
    aux.u = SVector(u, v, 0)
end;

function do_output(mpicomm, vtkdir, vtkstep, dg, Q, model, testname)
    ## name of the file that this MPI rank will write
    filename = @sprintf(
        "%s/%s_mpirank%04d_step%04d",
        vtkdir,
        testname,
        MPI.Comm_rank(mpicomm),
        vtkstep
    )

    statenames = flattenednames(vars_state(model, Prognostic(), eltype(Q)))

    writevtk(filename, Q, dg, statenames)

    ## generate the pvtu file for these vtk files
    if MPI.Comm_rank(mpicomm) == 0
        ## name of the pvtu file
        pvtuprefix = @sprintf("%s/%s_step%04d", vtkdir, testname, vtkstep)

        ## name of each of the ranks vtk files
        prefixes = ntuple(MPI.Comm_size(mpicomm)) do i
            @sprintf("%s_mpirank%04d_step%04d", testname, i - 1, vtkstep)
        end

        writepvtu(pvtuprefix, prefixes, statenames, eltype(Q))

        @info "Done writing VTK: $pvtuprefix"
    end
end;

function test_run(
    mpicomm,
    ArrayType,
    topl,
    problem,
    dt,
    N,
    timeend,
    FT,
    vtkdir,
    outputtime,
)
    grid = DiscontinuousSpectralElementGrid(
        topl,
        FloatType = FT,
        DeviceArray = ArrayType,
        polynomialorder = N,
    )

    model = AdvectionDiffusion{2, false, true}(problem)

    dg = DGModel(
        model,
        grid,
        UpwindNumericalFlux(),
        CentralNumericalFluxSecondOrder(),
        CentralNumericalFluxGradient(),
    )

    Q = init_ode_state(dg, FT(0))

    initialsumQ = weightedsum(Q)

    ## We integrate so that the final solution is equal to the initial solution
    Qe = copy(Q)

    rhs! = function (dQdt, Q, ::Nothing, t; increment = false)
        Filters.apply!(Q, :, grid, TMARFilter())
        dg(dQdt, Q, nothing, t; increment = false)
    end

    odesolver = SSPRK33ShuOsher(rhs!, Q; dt = dt, t0 = 0)

    cbTMAR = EveryXSimulationSteps(1) do
        Filters.apply!(Q, :, grid, TMARFilter())
    end

    ## create output directory on first rank of communicator
    if MPI.Comm_rank(mpicomm) == 0
        mkpath(vtkdir)
    end
    MPI.Barrier(mpicomm)

    vtkstep = 0
    ## output initial step
    do_output(mpicomm, vtkdir, vtkstep, dg, Q, model, "nonnegative")
    ## setup the output callback
    cbvtk = EveryXSimulationSteps(floor(outputtime / dt)) do
        vtkstep += 1
        minQ, maxQ = minimum(Q), maximum(Q)
        sumQ = weightedsum(Q)

        sumerror = (initialsumQ - sumQ) / initialsumQ

        @info @sprintf """Step  = %d
          minimum(Q)  = %.16e
          maximum(Q)  = %.16e
          sum error   = %.16e
          """ vtkstep minQ maxQ sumerror

        do_output(mpicomm, vtkdir, vtkstep, dg, Q, model, "nonnegative")
    end

    callbacks = (cbTMAR, cbvtk)
    solve!(Q, odesolver; timeend = timeend, callbacks = callbacks)

    minQ, maxQ = minimum(Q), maximum(Q)
    finalsumQ = weightedsum(Q)
    sumerror = (initialsumQ - finalsumQ) / initialsumQ
    error = euclidean_distance(Q, Qe)

    @test minQ ≥ 0

    @info @sprintf """Finished
    minimum(Q) = %.16e
    maximum(Q) = %.16e
    L2 error   = %.16e
    sum error  = %.16e
    """ minQ maxQ error sumerror
end;

let
    ArrayType = ClimateMachine.array_type()

    mpicomm = MPI.COMM_WORLD

    FT = Float64
    dim = 2
    Ne = 20
    polynomialorder = 4

    problem = SwirlingFlow()

    brickrange = (
        range(FT(0); length = Ne + 1, stop = 1),
        range(FT(0); length = Ne + 1, stop = 1),
        range(FT(0); length = Ne + 1, stop = 1),
    )

    topology = BrickTopology(
        mpicomm,
        brickrange[1:dim],
        boundary = ntuple(d -> (3, 3), dim),
    )

    maxvelocity = 2
    elementsize = 1 / Ne
    dx = elementsize / polynomialorder^2
    CFL = 1
    dt = CFL * dx / maxvelocity

    vtkdir =
        abspath(joinpath(ClimateMachine.Settings.output_dir, "vtk_nonnegative"))
    outputtime = 0.0625
    dt = outputtime / ceil(Int64, outputtime / dt)

    timeend = problem.period

    @info @sprintf """Starting
    FT               = %s
    dim              = %d
    Ne               = %d
    polynomial order = %d
    final time       = %.16e
    time step        = %.16e
    """ FT dim Ne polynomialorder timeend dt

    test_run(
        mpicomm,
        ArrayType,
        topology,
        problem,
        dt,
        polynomialorder,
        timeend,
        FT,
        vtkdir,
        outputtime,
    )
end;


"""
    compute_scc(m::Model=get_model(); year::Union{Int,Nothing}=nothing, last_year::Int=model_years[end], prtp::Float64=0.015, eta=1.45)

Computes the social cost of CO2 for an emissions pulse in `year` for the provided MimiDICE2013 model.
If no model is provided, the default model from MimiDICE2013.get_model() is used.
Constant discounting is used from the specified pure rate of time preference `prtp`.
"""
function compute_scc(m::Model=get_model(); year::Union{Int,Nothing}=nothing, last_year::Int=model_years[end], prtp::Float64=0.015, eta=1.45)
    year === nothing ? error("Must specify an emission year. Try `compute_scc(m, year=2020)`.") : nothing
    !(last_year in model_years) ? error("Invlaid value of $last_year for last_year. last_year must be within the model's time index $model_years.") : nothing
    !(year in model_years[1]:5:last_year) ? error("Cannot compute the scc for year $year, year must be within the model's time index $(model_years[1]):5:$last_year.") : nothing

    mm = get_marginal_model(m; year=year)

    return _compute_scc(mm, year=year, last_year=last_year, prtp=prtp, eta=eta)
end

"""
    compute_scc_mm(m::Model=get_model(); year::Union{Int,Nothing}=nothing, last_year::Int=model_years[end], prtp::Float64=0.015, eta=1.45)

Returns a NamedTuple (scc=scc, mm=mm) of the social cost of carbon and the MarginalModel used to compute it.
Computes the social cost of CO2 for an emissions pulse in `year` for the provided MimiDICE2013 model.
If no model is provided, the default model from MimiDICE2013.get_model() is used.
Constant discounting is used from the specified pure rate of time preference `prtp`.
"""
function compute_scc_mm(m::Model=get_model(); year::Union{Int,Nothing}=nothing, last_year::Int=model_years[end], prtp::Float64=0.015, eta=1.45)
    year === nothing ? error("Must specify an emission year. Try `compute_scc_mm(m, year=2020)`.") : nothing
    !(last_year in model_years) ? error("Invlaid value of $last_year for last_year. last_year must be within the model's time index $model_years.") : nothing
    !(year in model_years[1]:5:last_year) ? error("Cannot compute the scc for year $year, year must be within the model's time index $(model_years[1]):5:$last_year.") : nothing

    mm = get_marginal_model(m; year=year)
    scc = _compute_scc(mm; year=year, last_year=last_year, prtp=prtp, eta=eta)

    return (scc = scc, mm = mm)
end

# helper function for computing SCC from a MarginalModel, not to be exported or advertised to users
function _compute_scc(mm::MarginalModel; year::Int, last_year::Int, prtp::Float64, eta::Float64)
    ntimesteps = findfirst(isequal(last_year), model_years)     # Will run through the timestep of the specified last_year
    run(mm, ntimesteps=ntimesteps)

    # Go from trillion$ to $; multiply by -1 so that damages are positive; pulse
    # was already converted to CO2 so we don't need to multiply by 12/44
    marginal_damages = -1 * mm[:neteconomy, :C][1:ntimesteps] * 1e12

    cpc = mm.base[:neteconomy, :CPC]

    year_index = findfirst(isequal(year), model_years)

    df = [zeros(year_index - 1)..., ((cpc[year_index] / cpc[i])^eta * 1 / (1 + prtp)^(t - year) for (i, t) in enumerate(model_years) if year <= t <= last_year)...]
    scc = sum(df .* marginal_damages * 5)  # currently implemented as a 5year step function; so each timestep of discounted marginal damages is multiplied by 5
    return scc
end

"""
    get_marginal_model(m::Model=get_model(); year::Union{Int,Nothing}=nothing)

Creates a Mimi MarginalModel where the provided m is the base model, and the marginal model has additional emissions of CO2 in year `year`.
If no Model m is provided, the default model from MimiDICE2013.get_model() is used as the base model.
"""
function get_marginal_model(m::Model=get_model(); year::Union{Int,Nothing}=nothing)
    year === nothing ? error("Must specify an emission year. Try `get_marginal_model(m, year=2015)`.") : nothing
    !(year in model_years) ? error("Cannot add marginal emissions in $year, year must be within the model's time index $(model_years[1]):10:$last_year.") : nothing

    # note here that the pulse size will be used as the `delta` parameter for
    # the `MarginalModel` and thus allow computation of the SCC to return units of
    # dollars per ton, as long as `pulse_size` is in tons
    mm = create_marginal_model(m, 5 * 1e9)    # 1 GtCO2 per year for 5 years, so 5 * 10^9
    add_marginal_emissions!(mm.modified, year)

    return mm
end

"""
    add_marginal_emissions!(m::Model, year::Int)

Adds a marginal emission component to year m which adds 1Gt of additional CO2 emissions per year for five years starting in the specified `year`.
"""
function add_marginal_emissions!(m::Model, year::Int)
    add_comp!(m, Mimi.adder, :marginalemission, before=:co2cycle)

    time = Mimi.dimension(m, :time)
    addem = zeros(length(time))
    addem[time[year]] = 1.0     # 1 GtCO2 per year for five years

    set_param!(m, :marginalemission, :add, addem)
    connect_param!(m, :marginalemission, :input, :emissions, :E)
    connect_param!(m, :co2cycle, :E, :marginalemission, :output)
end

# Old available marginal model function
function getmarginal_dice_models(;emissionyear=2010)

    DICE = constructdice()
    run(DICE)

    mm = MarginalModel(DICE)
    m1 = mm.base
    m2 = mm.modified

    add_comp!(m2, adder, :marginalemission, before=:co2cycle)

    time = dimension(m1, :time)
    addem = zeros(length(time))
    addem[time[emissionyear]] = 1.0

    set_param!(m2, :marginalemission, :add, addem)
    connect_param!(m2, :marginalemission, :input, :emissions, :E)
    connect_param!(m2, :co2cycle, :E, :marginalemission, :output)

    run(m1)
    run(m2)

    return m1, m2
end