push!(LOAD_PATH,"../src/") # load module from local directory
using PolaronMobility

# Ivan Biaggio emailed 7th Nov 2018, suggesting there may be an error in these
# codes, due to a typographic error in the equations (particularly 62c) in
# Hellwarth1999
#
# As part of investigating this, these tests reproduce Table III in
# Hellwarth1999, to play with different ways of specifying the electron-phonon
# action.

"Print out F(alpha,beta) for a specific v,w; as a test"
function test_fns()
    @printf("\t\t")
    for α in 1:5
        @printf("α=%d\t\t",α)
    end
    @printf("\n")

    for β in 1:0.25:3.0
        v=w=4
        print("β: $β\t|\t")
        for α in 1:5
            @printf("%f\t",F(v,w,β,α))
        end
        println()
    end
end

# Our original definition of B, directly following the equation as written in Hellwarth1999
# 62c
println("Original B (62c, electron-phonon free energy):")
PolaronMobility.B(v,w,β,α) = α*v/(sqrt(π)*(exp(β)-1)) * quadgk(x->PolaronMobility.f(x,v,w,β),0,β/2)[1]
test_fns()  # OK - very primitive!
# Following private communication with Ivan Biaggio, and crosschecking against Osaka1959

println("Corrected B (62c, electron-phonon free energy, but corrected as Osaka.")
PolaronMobility.B(v,w,β,α) = exp(β)* α*v/(sqrt(π)*(exp(β)-1)) * quadgk(x->PolaronMobility.f(x,v,w,β),0,β/2)[1]
test_fns()

# Reproduce HellwarthTableIII, v and w solutions with finite Beta
function HellwarthTableIII()
    @printf("\t\t")
    for α in 1:5
        @printf("α=%d\t\t",α)
    end
    @printf("\n")

    for β in 1:0.25:3.0
        print("β: $β\t|\t")
        for α in 1:5
            v,w = feynmanvw(α, β)
            @printf("%.2f %.2f\t",v,w)
        end
        println()
    end

    print("Athermal\n β=Inf\t|\t")
    for α in 1:5
        v,w = feynmanvw(α)
        @printf("%.2f %.2f\t",v,w)
    end
    println()
end

# Our original definition of B, directly following the equation as written in Hellwarth1999
# 62c
println("Original B (62c, electron-phonon free energy):")
PolaronMobility.B(v,w,β,α) = α*v/(sqrt(π)*(exp(β)-1)) * quadgk(x->PolaronMobility.f(x,v,w,β),0,β/2)[1]
HellwarthTableIII()  # OK - very primitive!
# Following private communication with Ivan Biaggio, and crosschecking against Osaka1959

println("Corrected B (62c, electron-phonon free energy, but corrected as Osaka.")
PolaronMobility.B(v,w,β,α) = exp(β)* α*v/(sqrt(π)*(exp(β)-1)) * quadgk(x->PolaronMobility.f(x,v,w,β),0,β/2)[1]
HellwarthTableIII()

end

# Taking it back to where it all started; directly doing the full integral in Osaka1959
OsakaIntegrand(t,v,w,β) = exp(-t) / sqrt( w^2*t*(1-t/β) + (v^2-w^2)/v  *(1-exp(-v*t)+(1-coth(v/2*β)*(cosh(v*t)-1))) )
PolaronMobility.B(v,w,β,α) = 1/sqrt(pi) * α*v* exp(β)/(exp(β)-1) * quadgk(x->OsakaIntegrand(x,v,w,β),0,β)
# CURRENTLY BROKEN (!)
# sqrt(x) in the integrand is taken negative; so Julia throws a hissy fit - not expecting to find Imaginary values.
HellwarthTableIII()


using Test

using CompScienceMeshes

for T in [Float32, Float64]
    e0 = point(T,0,0,0)
    e1 = point(T,1,0,0)
    e2 = point(T,0,1,0)
    e3 = point(T,0,0,1)

    local p = simplex(e0, e1)
    q = simplex((e0+e1)/2, (-e0+3*e1)/2)

    pq = intersection(p,q)
    @test length(pq) == 1
    pq = pq[1]
    @test volume(pq) == 1/2
    @test pq.vertices[1] == q.vertices[1]
    @test pq.vertices[2] == p.vertices[2]
    @test 2*pq.tangents[1] == p.tangents[1] == q.tangents[1]

    c = (e0+e1+e2)/3
    p = simplex(e0, e1, e2)
    q = simplex(e0, e1, c)

    qq = intersection(q,q)
    pq = intersection(p,q)
    qp = intersection(q,p)

    # For these case where points of the subject are on the clipper or
    # the other way around, the precise computationof the intersection
    # is subject to numerical rounding errors. For efficiency reasons
    # this something that needs fixing. It should not affect the computation
    # of the interaction elements, however.
    @test sum([volume(r) for r in qq]) ≈ volume(q)
    @test sum([volume(r) for r in pq]) ≈ volume(q)
    @test sum([volume(r) for r in qp]) ≈ volume(q)

    # This case caused NaNs in the RWG BC Gram matrix. Turned out to be caused
    # by computing the intersection between two parallel lines. This needs a
    # more robust solution in the long run.
    p = point(T,-0.819550109062379,-0.394674531365384,-0.415415013001886)
    q = point(T,-0.66680743687093,-0.618706884644307,-0.415415013001886)
    r = point(T,-0.832690784854388,-0.535137873473586,-0.142314838273285)

    local a = point(T,-0.66680743687093,-0.618706884644307,-0.415415013001886)
    b = point(T,-0.749749110862659,-0.5769223790589465,-0.27886492563758547)
    c = point(T,-0.7730161102625658,-0.5161730964944257,-0.32438162142568566)

    isct1 = CompScienceMeshes.sutherlandhodgman([p,q,r],[a,b,c])

    @test isct1 == [
        point(T,-0.7730161102625658,-0.5161730964944257,-0.32438162142568566),
        point(T,-0.66680743687093,-0.618706884644307,-0.415415013001886),
        point(T,-0.66680743687093,-0.618706884644307,-0.415415013001886),
        point(T,-0.749749110862659,-0.5769223790589465,-0.27886492563758547)]

    splx1 = simplex(
        point(T,0,0,0),
        point(T,1,0,0),
        point(T,0,1,0))
    splx2 = simplex(
        point(T,0,-1,0),
        point(T,1,1,0),
        point(T,0.5,1,0))

    # Verify that all parts of the intersection are oriented consistently
    # and equal to the orientation of the inputs
    isct2 = intersection(splx1, splx2)
    @test length(isct2) == 2

    for splx in isct2
        @test normal(splx) ≈ point(T,0,0,1)
    end

    # Repeat this test but now with the orientation of the second
    # operand flipped
    splx1 = simplex(
        point(T,0,0,0),
        point(T,1,0,0),
        point(T,0,1,0))
    splx2 = simplex(
        point(T,1,1,0),
        point(T,0,-1,0),
        point(T,0.5,1,0))

    isct2 = intersection(splx1, splx2)
    @test length(isct2) == 2

    for splx in isct2
        @test normal(splx) ≈ point(T,0,0,1)
    end

    # Repeat this test but now with the orientation of the first operand flipped
    splx1 = simplex(
        point(T,1,0,0),
        point(T,0,0,0),
        point(T,0,1,0))
    splx2 = simplex(
        point(T,0,-1,0),
        point(T,1,1,0),
        point(T,0.5,1,0))

    isct2 = intersection(splx1, splx2)
    @test length(isct2) == 2

    for splx in isct2
        @test normal(splx) ≈ point(T,0,0,-1)
    end
end


####################################################################
function gpmodelVI(x, flux, σ2, idx; maxiter = 10, seed = 1, show_trace = true)
####################################################################

    #--------------------------------------------------------
    # Auxiliary
    #--------------------------------------------------------

    rg = MersenneTwister(seed)

    JITTER = 1e-8

    Σdiag = Diagonal(σ2)

    numF  = length(unique(idx))

    @assert(all(sort(unique(idx)) .== collect(1:numF)))

    N = length(x)

    @assert(N == length(flux) == length(idx) == length(σ2))

    @printf("Running gpmodelVIexact with %d and %d filters\n", N, numF)

    #--------------------------------------------------------
    # Initialise GP hyperparameters with ML estimation
    #--------------------------------------------------------

    a, b, logs = gpmodel(x, flux, σ2, idx; maxiter = 10_000, seed = 1, maxrandom=1000,show_trace = show_trace)

    #--------------------------------------------------------
    # Calculate kernel matrix and keep fixed!
    #--------------------------------------------------------

    dist2 = pairwise(SqEuclidean(), reshape(x, 1, N), dims=2)

    K1 = zeros(N, N)

    calculatekernelmatrix!(K1, dist2, rbf, [0.0; logs])

    K1 = K1 + JITTER * I

    K = PDMat((K1 + K1') / 2) # This helps us speed calculations

    #--------------------------------------------------------------------
    # Define functions for calling lower bound.
    # Actual calculation takes place in function "complete_lower_bound"
    #--------------------------------------------------------------------

    # Auxiliary function that calls the complete lower bound

    function lb(μα, Σα, μβ, Σβ, μf, Σf)

        local v²ₐ = v²ᵦ = 1000.0

        complete_lower_bound(; y=flux, idx=idx, S=Σdiag, μα=μα, Σα=Σα, v²ₐ=v²ₐ, μβ=μβ, Σβ=Σβ, v²ᵦ=v²ᵦ, μf=μf, Σf=Σf, K=K)

    end

    # Unpacks vectorised parameters

    function unpack(param)

        @assert(length(param) == numF*4)

        local μα =               param[1+0*numF:1*numF]

        local Σα = Diagonal(exp.(param[1+1*numF:2*numF]))

        local μβ =               param[1+2*numF:3*numF]

        local Σβ = Diagonal(exp.(param[1+3*numF:4*numF]))

        return μα, Σα, μβ, Σβ

    end

    # Constructs necessary matrices

    function negative_lb_helper(param)

        local μα, Σα, μβ, Σβ = unpack(param)

        local Ā  = Diagonal(μα[idx])

        local Σa = Diagonal(Σα.diag[idx])

        local b̄  = μβ[idx]

        # Calculate optimal posterior q(f)

        # local Σf = ((Ā*inv(Σdiag)*Ā + inv(Σdiag)*Σa) + inv(K))\I

        local invDiagPart = (Ā*(Σdiag\Ā) + Σdiag\Σa)\I #
                                                       # This is faster/more stable
        local Σf = invDiagPart*((invDiagPart + K)\K)   #

        local μf = Σf * Ā * (Σdiag \ (flux - b̄))

        -1.0 * lb(μα, Σα, μβ, Σβ, μf, Σf)

    end

    # Initial solution given by point estimates

    initsol = [a; log.(1.0*ones(numF)); b; log.(1.0*ones(numF))]

    # Use the LBFGS optimiser

    # opt = Optim.Options(show_trace = show_trace, iterations = maxiter, show_every = 1)
    #
    # result = Optim.optimize(negative_lb_helper, initsol, LBFGS(), opt, autodiff=:forward)

    # result = bboptimize(negative_lb_helper; SearchRange=(-10.0,10.0), NumDimensions = numF*4, Method=:xnes)

    opt = NLopt.Opt(:LN_BOBYQA, numF*4)

    opt.xtol_rel = 1e-8

    COUNTER, BESTF = 0, Inf

    function myfunc(x, dummy)

        COUNTER +=1

        F = negative_lb_helper(x)

        BESTF = min(F, BESTF)

        mod(COUNTER, 100)==0 ? @printf("Iter %d f=%e\n", COUNTER, BESTF) : nothing

        return F

    end

    opt.min_objective = myfunc

    (_minf, minx, _ret) = NLopt.optimize(opt, initsol)

    # Retrieve result and return posterior distribution of type Distributions.MvNormal

    μα, Σα, μβ, Σβ = unpack(minx)#result.minimizer)

    return MvNormal([μα; μβ], Diagonal([Σα.diag; Σβ.diag]))

end


"""
```
    NonmaximaSuppression <: AbstractEdgeThinningAlgorithm
    NonmaximaSuppression(; threshold::Union{Number, Percentile} = Percentile(20))

    f = NonmaximaSuppression()
    f(out::AbstractArray, mag::AbstractArray, g₁::AbstractArray, g₂::AbstractArray, f::NonmaximaSuppression)
```

Isolates local maxima of gradient magnitude `mag` along the local gradient
direction and stores the result in `out`.  The arguments `g₁` and `g₂` represent
the  gradient in the first spatial dimension (y), and the second spatial
dimension (x), respectively.

# Details

TODO

# Example

```julia

using TestImages, FileIO, ImageView, ImageEdgeDetection, ImageFiltering

img =  testimage("mandril_gray")

# Gradient in the first and second spatial dimension
g₁, g₂ = imgradients(img, KernelFactors.scharr)

# Gradient magnitude
mag = hypot.(g₁, g₂)

nms = zeros(eltype(mag), axes(mag))
# Instantiate the NonmaximaSuppression functor.
f = NonmaximaSuppression()

# Suppress the non-maximal gradient magnitudes and store the result in `nms`.
f(nms, mag, g₁, g₂)

imshow(img)
imshow(mag)
imshow(nms)
```

# References
J. Canny, "A Computational Approach to Edge Detection," in IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. PAMI-8, no. 6, pp. 679-698, Nov. 1986, doi: 10.1109/TPAMI.1986.4767851.
"""
@with_kw struct NonmaximaSuppression{ T <: Union{Real,AbstractGray, Percentile}} <: AbstractEdgeThinningAlgorithm
    threshold::T = Percentile(20)
end

function (f::NonmaximaSuppression)(out::AbstractArray, mag::AbstractArray, g₁::AbstractArray, g₂::AbstractArray)
    @unpack threshold = f

    low_threshold =  typeof(threshold) <: Percentile ? StatsBase.percentile(vec(mag), threshold.p) : threshold

    # Isolate local maxima of gradient magnitude by “non-maximum suppression”
    # along the local gradient direction.
    suppress_non_maxima!(out, mag, g₁, g₂, low_threshold)

    return out
end

"""
```
    suppress_non_maxima!(out::AbstractArray, mag::AbstractArray, gx::AbstractArray, gy::AbstractArray, threshold::Number)
```

Isolates local maxima of gradient magnitude by “non-maximum suppression” along the local gradient direction.

"""
function suppress_non_maxima!(out::AbstractArray, mag::AbstractArray, g₁::AbstractArray, g₂::AbstractArray, threshold::Number)
    itp = interpolate(mag, BSpline(Linear()))
    etp = extrapolate(itp, Flat())
    m₁ = zero(eltype(mag))
    m₂ = zero(eltype(mag))
    @inbounds for i in CartesianIndices(mag)
        r, c = i.I
        d₁ = gray(g₁[i])
        d₂ = gray(g₂[i])
        mc = mag[i]
        if mc < threshold || mc == 0 || isnan(mc)
            out[r,c] = zero(eltype(mag))
        else
            # Ensure the vector 𝐝 = [d₁, d₂] has unit norm.
            d₁ = d₁ / mc
            d₂ = d₂ / mc
            m₁ = etp(r + d₁, c + d₂)
            m₂ = etp(r - d₁, c - d₂)
            out[r,c] = (m₁ <= mc) && (mc >= m₂) ? mag[r,c] : zero(eltype(mag))
        end
    end
    return nothing
end


#=

  Map coloring in Julia ConstraintSolver.jl

  Simple map coloring problem.

  Model created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Julia page: http://www.hakank.org/julia/

=#

using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function map_coloring(problem,print_solutions=true,all_solutions=true)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            # "branch_split"=>:Biggest,
                                                            "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>6,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    n, _ = size(problem)
    num_colors = 4
    # println("n:$n num_colors:$num_colors")
    @variable(model, 1 <= x[1:n] <= num_colors, Int)

    for c1 in 1:n, c2 in 1:n
        if problem[c1,c2] == 1
            @constraint(model,x[c1] != x[c2])
        end
    end

    # symmertry breaking
    for c in 1:min(n,num_colors)
        @constraint(model,x[c] <= c)
    end

    # @objective(model,Max,z)

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                # println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                println("x:$x_val")

            end
        end
    else
        println("status:$status")
    end

    return status
end

map_coloring_problems = Dict(
    # Connections between these countries:
    # [belgium, denmark, france, germany, netherlands, luxembourg]
    :1 => resize_matrix([[0, 0, 1, 1, 1, 1],
                         [0, 0, 0, 1, 0, 0],
                         [1, 0, 0, 1, 1, 0],
                         [1, 1, 1, 0, 1, 1],
                         [1, 0, 1, 1, 0, 0],
                         [1, 0, 0, 1, 0, 0]])

)

@time map_coloring(map_coloring_problems[:1])

# random instance
function random_map(n=10,lim=0.5)
    m = zeros(Int8,n,n)
    for i in 1:n, j in 1:i-1
       if rand() >= lim
        m[i,j] = 1
        m[j,i] = 1
       end
    end
    return m
end

for n in 10:20
    println("\nn:$n")
    m = random_map(n,0.5)
    # display(m)
    @time map_coloring(m,true,false)
end


# TODO: this file will be used for setting up similar tests for the rest of the samplers.
#
# This samples the funnel distribution, then tests that inference is correct.
#

using MCMC

#log  ~ N(0, 3), x_i ~ N(0, ), i=1,2,...,model.size

function funnel_loglik(x::Vector{Float64})
    # First element is log(sigma^2). Rest of the elements are x_d
    s2 = exp(x[1])
    nx = length(x) - 1
    -0.5*((x[1]-0.0)^2/9.0 + dot(x[2:end], x[2:end])/s2 + nx*log(s2))
end

mcmodel = model(funnel_loglik, init=[10, zeros(4)])
mcsampler = SliceSampler()
mcrunner = SerialMC(steps=100000, burnin=10000)
mcchain = run(mcmodel*mcsampler*mcrunner)

samples01 = mcchain.samples[:, 1]

ess01 = ess(samples01, maxlag=1000)

@printf("Monte Carlo samples of first parameter: ess=%.2f, out of raw %d\n", ess01, length(samples01))

mean01 = mean(samples01)
std01 = std(samples01)
meanse01 = std01/sqrt(ess01)
zscore01 = mean01/meanse01

@printf("Monte Carlo samples of first parameter: mean=%f, std=%f, meanse=%f, zscore=%f\n",
  mean01, std01, meanse01, zscore01)
if abs(zscore01) > 3
    warn("Log-posterior mean estimate for the first parameter differs from truth by more than 3 MC-stderr. Chance of this should be <0.3%.")
end

var01 = var(samples01)
varse01 = var01*sqrt(2/(ess01-1))
zscore01 = (var01-9)/varse01

@printf("Monte Carlo samples of first parameter: var=%f, varse=%f, zscore=%f\n", var01, varse01, zscore01)
if abs(zscore01) > 3
    warn("Log-posterior mean estimate for the first parameter differs from truth by more than 3 MC-stderr. Chance of this should be <0.3%.")
end


#=
Copyright (c) 2015, Intel Corporation
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
- Redistributions of source code must retain the above copyright notice,
  this list of conditions and the following disclaimer.
- Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
THE POSSIBILITY OF SUCH DAMAGE.
=#

# This code was ported from a MATLAB implementation available at
# http://www.albertostrumia.it/Fractals/FractalMatlab/Julia/juliaSH.m.
# Original MATLAB implementation is copyright (c) the author.

using DocOpt
if "--demo" in ARGS
    using Winston
end

include("ndgrid.jl") # provides the meshgrid function

function juliaset(demo::Bool)
    col = 128             # color depth
    m = 1000              # image size
    cx = 0                # center X
    cy = 0                # center Y
    l = 1.5               # span

    if demo
        iters = 50
        zoomAmount = 0.9
    else
        iters = 10
        zoomAmount = 0.6
    end

    # The c constant.
    c = -.745429 + .11308*im

    for zoom = 1:iters

        # `x` and `y` are two 1000-element arrays representing the x
        # and y axes: [-1.5, -1.497, ..., 0, ..., 1.497, 1.5] on the
        # first iteration of this loop.
        x = linspace(cx-l, cx+l, m)
        y = linspace(cy-l, cy+l, m)

        # `X` and `Y` are two arrays containing, respectively, the x-
        # and y-coordinates of each point on a 1000x1000 grid.
        (X, Y) = meshgrid(x, y)

        # Let `Z` represent the complex plane: a 1000x1000 array of
        # numbers each with a real and a complex part.
        Z = X + im*Y

        # Iterate the Julia set computation (squaring each element of
        # Z and adding c) for `col` steps.
        W = iterate(col, Z, c)

        # Mask out the NaN values (overflow).
        minval = minimum(W)
        W[isnan(W)] = minval - minval/10

        if demo
            # Plot the result as a "heat map"-style plot.  The values
            # of the elements of W determine the "color" of the
            # corresponding point.
            plot = Winston.imagesc(W)
            Winston.display(plot)
        end

        # Zoom into the next frame, shrinking the distance that `x`
        # and `y` will cover.
        l = l * zoomAmount
    end

end

function iterate(col, Z, c)
    for k = 1:col
        Z = Z.*Z .+ c
    end
    return exp(-abs(Z))
end

function main()
    doc = """juliaset.jl

Julia set fractal visualization.

Usage:
  juliaset.jl -h | --help
  juliaset.jl [--demo]

Options:
  -h --help  Show this screen.
  --demo     Use settings that look good for the purposes of a demo, and plot the results using Winston.
"""
    arguments = docopt(doc)
    demo = arguments["--demo"]

    tic()
    iterate(0, Array{Complex{Float64}}(0,2), Complex(0.0))
    time = toq()
    println("SELFPRIMED ", time)

    tic()
    juliaset(demo)
    time = toq()
    println("SELFTIMED ", time)
end

main()


using KNITRO, MathProgBase
using Compat.Test

mutable struct Myquad <: MathProgBase.AbstractNLPEvaluator
end

MathProgBase.features_available(d::Myquad) = [:Grad,:HessVec]
function MathProgBase.initialize(d::Myquad,requested_features::Vector{Symbol})
    for feat in requested_features
        if !(feat in [:Grad,:Jac,:HessVec])
            error("Unsupported feature $feat")
        end
    end
end

MathProgBase.eval_g(d::Myquad,g,x) = nothing
MathProgBase.eval_jac_g(d::Myquad,g1,g2) = nothing
MathProgBase.jac_structure(d::Myquad) = Int64[], Int64[]

function MathProgBase.eval_f(d::Myquad, x)
    return (x[1]-1.)^2+(x[2]-2.)^2
end

function MathProgBase.eval_grad_f(d::Myquad, grad_f, x)
    grad_f[1] = 2*(x[1]-1.)
    grad_f[2] = 2*(x[2]-2.)
end

function MathProgBase.eval_hesslag(d::Myquad,h,x,σ,μ)
    h[1] = 2*σ
    h[2] = 2*σ
end

m = MathProgBase.NonlinearModel(KnitroSolver(KTR_PARAM_OUTLEV=2))
MathProgBase.loadproblem!(m, 2, 0, -3. * ones(2), 3. * ones(2), Float64[], Float64[], :Min, Myquad())
MathProgBase.setwarmstart!(m,rand(2))
MathProgBase.optimize!(m)
x = MathProgBase.getsolution(m)

# --- test optimal solutions ---
@testset "Test optimal solutions" begin
    @test isapprox(x[1],1.,atol=1.0e-5)
    @test isapprox(x[2],2.,atol=1.0e-5)
end


## CurveSpace

Space{S<:Fourier}(d::PeriodicCurve{S})=Fourier(d)
Space{S<:Laurent}(d::PeriodicCurve{S})=Laurent(d)

#TODO: Make type stable
Base.convert(::Type{Curve},f::Fun)=isa(domain(f),IntervalDomain)?IntervalCurve(f):PeriodicCurve(f)

# function evaluate{C<:Curve,TT}(f::AbstractVector,S::Space{TT,C},x::Number)
#     rts=roots(domain(S).curve-x)
#     @assert length(rts)==1
#     evaluate(Fun(f,setdomain(S,canonicaldomain(S))),first(rts))
# end

identity_fun{C<:Curve,TT}(d::Space{TT,C})=Fun(domain(d).curve.coefficients,
                                              setdomain(space(domain(d).curve),domain(d)))

# Bernstein polynomials are given by:
#
# b_{ν,n}(x) = binomial(n,ν)/2^n*(1+x)^ν(1-x)^{n-v},  for x ∈ [-1,1], ν = 0,…,n.
#
export Bernstein, Bézier

immutable Bernstein{order,T} <: RealUnivariateSpace{T}
    domain::Interval{T}
    Bernstein(d) = new(d)
    Bernstein() = new(Interval{T}())
end

const Bézier = Bernstein # option+e e gives é

@compat (::Type{Bernstein{O}}){O}() = Bernstein{O,Float64}()
@compat (::Type{Bernstein{O}}){O}(d::Domain) = Bernstein{O,eltype(d)}(d)
@compat (::Type{Bernstein{O}}){O}(d::Vector) = Bernstein{O}(Interval(d))

order{O}(::Bernstein{O}) = O
order{O,T}(::Type{Bernstein{O,T}}) = O
dimension{O}(::Bernstein{O}) = O+1
dimension{O,T}(::Type{Bernstein{O,T}}) = O+1

canonicalspace(B::Bernstein) = Chebyshev(domain(B))
canonicaldomain{O,T}(B::Bernstein{O,T}) = Interval{T}()

spacescompatible{O,T}(a::Bernstein{O,T},b::Bernstein{O,T})=domainscompatible(a,b)

setdomain{O}(S::Bernstein{O},d::Domain)=Bernstein{O}(d)

identity_fun{order,T}(B::Bernstein{order,T})=Fun(collect(-one(T):2one(T)/order:one(T)),B)

evaluate(f::AbstractVector,S::Bernstein,z) = decasteljau(f,S,tocanonical(S,z))

# de Casteljau's numerically stable evaluation of Bernstein polynomials

function decasteljau(f::AbstractVector,S::Bernstein,z)
    β,omz,opz = copy(f),one(z)-z,one(z)+z
    for i=order(S):-1:1
        for j=1:i
            β[j] = (β[j]*omz+β[j+1]*opz)/2
        end
    end
    β[1]
end

function decasteljau(f::AbstractVector,S::Bernstein,z::AbstractVector)
    βinit = copy(f)
    β = copy(βinit)
    T = promote_type(eltype(f),eltype(z))
    ret = zeros(T,length(z))
    for k=1:length(z)
        omzk,opzk = one(z[k])-z[k],one(z[k])+z[k]
        for i=order(S):-1:1
            for j=1:i
                β[j] = (β[j]*omzk+β[j+1]*opzk)/2
            end
        end
        ret[k] = β[1]
        copy!(β,βinit)
    end
    ret
end

splitbernstein{B<:Bernstein,T}(f::Fun{B,T},z) = splitbernstein(coefficients(f),space(f),z)

function splitbernstein(f::AbstractVector,S::Bernstein,z)
    β,omz,opz = copy(f),one(z)-z,one(z)+z
    β1,β2 = copy(β),copy(β)
    for i=order(S):-1:1
        for j=1:i
            β[j] = (β[j]*omz+β[j+1]*opz)/2
        end
        β1[order(S)+2-i] = β[1]
        β2[i] = β[i]
    end
    Fun(interlace(β1,β2),PiecewiseSpace(Bernstein{order(S)}(Interval(first(domain(S)),z)),Bernstein{order(S)}(Interval(z,last(domain(S))))))
end

Fun(f::Function,S::Bernstein) = Fun(Fun(f,canonicalspace(S),dimension(S)),S)

# Connection coefficients with Chebyshev polynomials are derived in
#
#    A. Rababah, "Transformation of Chebyshev–Bernstein polynomial basis", Comp. Meth. Appl. Math. 3:608–622, 2003.
#
function coefficients(f::Vector,a::Chebyshev,b::Bernstein)
    if domain(a) == domain(b)
        n = length(f)-1
        @assert n == order(b)
        M = zeros(eltype(f),n+1,n+1)
        for j=0:n
            for k=0:n
                temp = zero(eltype(f))
                for i=max(0,j+k-n):min(j,k)
                    temp += (-1)^(k-i)*binomial(2k,2i)*binomial(n-k,j-i)
                end
                M[j+1,k+1] = temp/binomial(n,j)
            end
        end
        M*f
    else
        defaultcoefficients(f,a,b)
    end
end

function coefficients(f::Vector,a::Bernstein,b::Chebyshev)
    if domain(a) == domain(b)
        n = length(f)-1
        @assert n == order(a)
        Minv = zeros(eltype(f),n+1,n+1)
        for j=0:n
            for k=0:n
                temp = zero(eltype(f))
                for i=0:j
                    temp += (-1)^(j-i)*binomial(2j,2i)*binomial(2k+2i,k+i)*binomial(2n+2j-2k-2i,n+j-k-i)/binomial(n+j,k+i)
                end
                Minv[j+1,k+1] = (2-FastTransforms.δ(j,0))/4^(n+j)*binomial(n,k)*temp
            end
        end
        Minv*f
    else
        defaultcoefficients(f,a,b)
    end
end


#  Compute r = A * x + s - b in-place
function primal_kkt_condition!(r::AbstractVector{T}, A::AbstractMatrix{T}, x::AbstractVector{T}, s::SplitVector{T}, b::AbstractVector{T}) where {T <: AbstractFloat}
	# r = A * x
	mul!(r, A, x)
	@. r += s
	@. r -= b
	return nothing
end

# Compute r = P * x + q - A' * μ in-place
function dual_kkt_condition!(r::AbstractVector{T}, r_temp::AbstractVector{T}, P::AbstractMatrix{T}, x::AbstractVector{T}, q::AbstractVector{T}, A::AbstractMatrix{T}, μ::AbstractVector{T}) where {T <: AbstractFloat}
	mul!(r, P, x)
	@. r += q

	mul!(r_temp, A', μ)
	@. r -= r_temp
	return nothing
end

# Einv * x
unscale_primal!(x::AbstractVector{T}, Einv::AbstractMatrix{T}) where {T} = lmul!(Einv, x)

# cinv * Dinv * x
function unscale_dual!(x::AbstractVector{T}, Dinv::AbstractMatrix{T}, cinv) where {T}
		lmul!(Dinv, x)
		@. x *= cinv
		return nothing
end

function calculate_residuals!(ws::COSMO.Workspace,  IGNORESCALING_FLAG::Bool = false)
	# reuse vectors from temporary workspace
	r_prim = ws.utility_vars.vec_m
	r_dual = ws.utility_vars.vec_n
	r_temp = ws.utility_vars.vec_n2

	# r_prim = A * x + s - b
	primal_kkt_condition!(r_prim, ws.p.A, ws.vars.x, ws.vars.s, ws.p.b)

	# ∇f0 + ∑ νi ∇hi(x*) == 0 condition
	# r_dual = P * x + q - A' * μ
	dual_kkt_condition!(r_dual, r_temp, ws.p.P, ws.vars.x, ws.p.q, ws.p.A, ws.vars.μ)

	if (ws.settings.scaling != 0 && !IGNORESCALING_FLAG)
		# unscale primal residual: r_prim = Einv * ( A * x + s - b)
		unscale_primal!(r_prim, ws.sm.Einv)

		# unscale dual residual: r_dual = cinv * Dinv * (P * x + q - A' * μ)
		unscale_dual!(r_dual, ws.sm.Dinv, ws.sm.cinv)
	end

	return norm(r_prim, Inf), norm(r_dual, Inf)

end

function max_res_component_norm(ws::COSMO.Workspace, IGNORESCALING_FLAG::Bool = false)
		r_prim = ws.utility_vars.vec_m
		r_dual = ws.utility_vars.vec_n

		# Determine if the components inside max{ } have to be unscaled
		unscale = (ws.settings.scaling != 0 && !IGNORESCALING_FLAG)

		# Calculate: max { ||Einv * Ax ||_Inf, ||Einv * s||_Inf, ||Einv * b||_Inf }
		# Einv * A * x
		mul!(r_prim, ws.p.A, ws.vars.x)
		unscale && unscale_primal!(r_prim, ws.sm.Einv)
		max_norm_prim = norm(r_prim, Inf)

		# Einv * s
		@. r_prim = ws.vars.s
		unscale && unscale_primal!(r_prim, ws.sm.Einv)
		max_norm_prim = max.(max_norm_prim, norm(r_prim, Inf))

		# Einv * b
		@. r_prim = ws.p.b
		unscale && unscale_primal!(r_prim, ws.sm.Einv)
		max_norm_prim = max.(max_norm_prim, norm(r_prim, Inf))

		# Calculate: max { ||cinv * Dinv * P * x ||_Inf, ||cinv * Dinv * q||_Inf, ||cinv * Dinv * A' * μ||_Inf }
		# ||cinv * Dinv * P * x ||_Inf
		mul!(r_dual, ws.p.P, ws.vars.x)
		unscale && unscale_dual!(r_dual, ws.sm.Dinv, ws.sm.cinv)
		max_norm_dual = norm(r_dual, Inf)

		# ||cinv * Dinv * q||_Inf
		@. r_dual = ws.p.q
		unscale && unscale_dual!(r_dual, ws.sm.Dinv, ws.sm.cinv)
		max_norm_dual = max.(max_norm_dual, norm(r_dual, Inf))

		# ||cinv * Dinv * A' * μ||_Inf
		mul!(r_dual, ws.p.A', ws.vars.μ)
		unscale && unscale_dual!(r_dual, ws.sm.Dinv, ws.sm.cinv)
		max_norm_dual = max.(max_norm_dual, norm(r_dual, Inf))

	return max_norm_prim, max_norm_dual
end

function isapprox_feasible(residual, max_norm, eps_abs, eps_rel)
	ϵ = eps_abs + eps_rel * max_norm
	return residual < ϵ
end

function isapprox_primal_feasible(r::ResultInfo, eps_abs, eps_rel)
    return isapprox_feasible(r.r_prim, r.max_norm_prim, eps_abs, eps_rel)
end

function isapprox_dual_feasible(r::ResultInfo, eps_abs, eps_rel)
    return isapprox_feasible(r.r_dual, r.max_norm_dual, eps_abs, eps_rel)
end

function is_primal_feasible(r::ResultInfo, settings::Settings)
	return isapprox_primal_feasible(r, settings.eps_abs, settings.eps_rel)
end

function is_dual_feasible(r::ResultInfo, settings::Settings)
	return isapprox_dual_feasible(r, settings.eps_abs, settings.eps_rel)
end

function is_primal_nearly_feasible(r::ResultInfo, settings::Settings)
	return isapprox_primal_feasible(r, settings.nearly_ratio * settings.eps_abs, settings.nearly_ratio * settings.eps_rel)
end

function is_dual_nearly_feasible(r::ResultInfo, settings::Settings)
	return isapprox_dual_feasible(r, settings.nearly_ratio * settings.eps_abs, settings.nearly_ratio * settings.eps_rel)
end

function has_converged(ws::COSMO.Workspace{T}, r::ResultInfo) where {T <: AbstractFloat}
	# check activation of accelerator
	COSMO.check_activation!(ws, ws.activation_reason, r)

	# if an optimal objective value was specified for the problem check if current solution is within specified accuracy
	obj_true_FLAG = true
	settings = ws.settings
	if !isnan(settings.obj_true)
		current_cost = calculate_cost!(ws.utility_vars.n, ws.vars.x, ws.p.P, ws.p.q, ws.sm.cinv[])
		obj_true_FLAG = abs(settings.obj_true - current_cost) <= settings.obj_true_tol
	end

	return is_primal_feasible(r, settings) && is_dual_feasible(r, settings) && obj_true_FLAG
end

# cost = cinv *( 1/2 x' * P * x + q' x)
function calculate_cost!(temp::AbstractVector{T}, x::AbstractVector{T}, P::SparseMatrixCSC{T, Int}, q::AbstractVector{T}, cinv::T = one(T)) where {T <: AbstractFloat}
	# P x
	mul!(temp, P, x)
	return cinv * (T(0.5) * dot(temp, x) + dot(q, x))
end

function calculate_result_info!(ws::COSMO.Workspace)
    r_prim, r_dual = calculate_residuals!(ws)
    max_norm_prim, max_norm_dual = max_res_component_norm(ws)
    return ResultInfo(r_prim, r_dual, max_norm_prim, max_norm_dual, ws.rho_updates)
end


using Test, LinearAlgebra
using DiffEqSensitivity, SteadyStateDiffEq, DiffEqBase, NLsolve
using OrdinaryDiffEq
using ForwardDiff, Calculus
@testset "Adjoint sensitivities of steady state solver" begin
  function f!(du,u,p,t)
    du[1] = p[1] + p[2]*u[1]
    du[2] = p[3]*u[1] + p[4]*u[2]
  end

  function jac!(J,u,p,t) #df/dx
    J[1,1] = p[2]
    J[2,1] = p[3]
    J[1,2] = 0
    J[2,2] = p[4]
    nothing
  end

  function paramjac!(fp,u,p,t) #df/dp
    fp[1,1] = 1
    fp[2,1] = 0
    fp[1,2] = u[1]
    fp[2,2] = 0
    fp[1,3] = 0
    fp[2,3] = u[1]
    fp[1,4] = 0
    fp[2,4] = u[2]
    nothing
  end

  function dg!(out,u,p,t,i)
    (out.=-2.0.+u)
  end

  function g(u,p,t)
    sum((2.0.-u).^2)/2 + sum(p.^2)/2
  end

  u0 = zeros(2)
  p = [2.0,-2.0,1.0,-4.0]
  prob = SteadyStateProblem(f!,u0,p)
  abstol = 1e-10
  @testset "for p" begin
    println("Calculate adjoint sensitivities from Jacobians")

    sol_analytical = [-p[1]/p[2], p[1]*p[3]/(p[2]*p[4])]

    J = zeros(2,2)
    fp = zeros(2,4)
    gp = zeros(4)
    gx = zeros(1,2)
    delg_delp = copy(p)

    jac!(J,sol_analytical,p,nothing)
    dg!(vec(gx),sol_analytical,p,nothing,nothing)
    paramjac!(fp,sol_analytical,p,nothing)

    lambda = J' \gx'
    res_analytical = delg_delp'-lambda' * fp # = -gx*inv(J)*fp

    @info "Expected result" sol_analytical, res_analytical, delg_delp'-gx*inv(J)*fp

    @info "Calculate adjoint sensitivities from autodiff & numerical diff"
    function G(p)
      tmp_prob = remake(prob,u0=convert.(eltype(p),prob.u0),p=p)
      sol = solve(tmp_prob,
        SSRootfind(nlsolve = (f!,u0,abstol) -> (res=NLsolve.nlsolve(f!,u0,autodiff=:forward,method=:newton,iterations=Int(1e6),ftol=1e-14);res.zero))
        )
      A = convert(Array,sol)
      g(A,p,nothing)
    end
    res1 = ForwardDiff.gradient(G,p)
    res2 = Calculus.gradient(G,p)
    #@info res1, res2, res_analytical

    @test res1 ≈ res_analytical' rtol = 1e-7
    @test res2 ≈ res_analytical' rtol = 1e-7
    @test res1 ≈ res2 rtol = 1e-7

    @info "Adjoint sensitivities"

    # with jac, param_jac
    f1 = ODEFunction(f!;jac=jac!, paramjac=paramjac!)
    prob1 = SteadyStateProblem(f1,u0,p)
    sol1 = solve(prob1,DynamicSS(Rodas5()))
    res1a = adjoint_sensitivities(sol1,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(),g,dg!)
    res1b = adjoint_sensitivities(sol1,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(),g,nothing)
    res1c = adjoint_sensitivities(sol1,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autodiff=false),g,nothing)
    res1d = adjoint_sensitivities(sol1,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autojacvec=TrackerVJP()),g,nothing)
    res1e = adjoint_sensitivities(sol1,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autojacvec=ReverseDiffVJP()),g,nothing)
    res1f = adjoint_sensitivities(sol1,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autojacvec=ZygoteVJP()),g,nothing)
    res1g = adjoint_sensitivities(sol1,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autodiff=false,autojacvec=false),g,nothing)

    # with jac, without param_jac
    f2 = ODEFunction(f!;jac=jac!)
    prob2 = SteadyStateProblem(f2,u0,p)
    sol2 = solve(prob2,DynamicSS(Rodas5()))
    res2a = adjoint_sensitivities(sol2,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(),g,dg!)
    res2b = adjoint_sensitivities(sol2,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(),g,nothing)
    res2c = adjoint_sensitivities(sol2,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autodiff=false),g,nothing)
    res2d = adjoint_sensitivities(sol2,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autojacvec=TrackerVJP()),g,nothing)
    res2e = adjoint_sensitivities(sol2,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autojacvec=ReverseDiffVJP()),g,nothing)
    res2f = adjoint_sensitivities(sol2,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autojacvec=ZygoteVJP()),g,nothing)
    res2g = adjoint_sensitivities(sol2,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autodiff=false,autojacvec=false),g,nothing)

    # without jac, without param_jac
    f3 = ODEFunction(f!)
    prob3 = SteadyStateProblem(f3,u0,p)
    sol3 = solve(prob3,DynamicSS(Rodas5()))
    res3a = adjoint_sensitivities(sol3,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(),g,dg!)
    res3b = adjoint_sensitivities(sol3,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(),g,nothing)
    res3c = adjoint_sensitivities(sol3,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autodiff=false),g,nothing)
    res3d = adjoint_sensitivities(sol3,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autojacvec=TrackerVJP()),g,nothing)
    res3e = adjoint_sensitivities(sol3,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autojacvec=ReverseDiffVJP()),g,nothing)
    res3f = adjoint_sensitivities(sol3,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autojacvec=ZygoteVJP()),g,nothing)
    res3g = adjoint_sensitivities(sol3,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autodiff=false,autojacvec=false),g,nothing)

    @test norm(res_analytical' .- res1a) < 1e-7
    @test norm(res_analytical' .- res1b) < 1e-7
    @test norm(res_analytical' .- res1c) < 1e-7
    @test norm(res_analytical' .- res1d) < 1e-7
    @test norm(res_analytical' .- res1e) < 1e-7
    @test_broken norm(res_analytical' .- res1f) < 1e-7
    @test norm(res_analytical' .- res1g) < 1e-7
    @test norm(res_analytical' .- res2a) < 1e-7
    @test norm(res_analytical' .- res2b) < 1e-7
    @test norm(res_analytical' .- res2c) < 1e-7
    @test norm(res_analytical' .- res2d) < 1e-7
    @test norm(res_analytical' .- res2e) < 1e-7
    @test_broken norm(res_analytical' .- res2f) < 1e-7
    @test norm(res_analytical' .- res2g) < 1e-7
    @test norm(res_analytical' .- res3a) < 1e-7
    @test norm(res_analytical' .- res3b) < 1e-7
    @test norm(res_analytical' .- res3c) < 1e-7
    @test norm(res_analytical' .- res3d) < 1e-7
    @test norm(res_analytical' .- res3e) < 1e-7
    @test_broken norm(res_analytical' .- res3f) < 1e-7
    @test norm(res_analytical' .- res3g) < 1e-7

    @info "oop checks"
    function foop(u,p,t)
      dx = p[1] + p[2]*u[1]
      dy = p[3]*u[1] + p[4]*u[2]
      [dx,dy]
    end
    proboop = SteadyStateProblem(foop,u0,p)
    soloop = solve(proboop,DynamicSS(Rodas5()))

    res4a = adjoint_sensitivities(soloop,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(),g,dg!)
    res4b = adjoint_sensitivities(soloop,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(),g,nothing)
    res4c = adjoint_sensitivities(soloop,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autodiff=false),g,nothing)
    res4d = adjoint_sensitivities(soloop,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autojacvec=TrackerVJP()),g,nothing)
    res4e = adjoint_sensitivities(soloop,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autojacvec=ReverseDiffVJP()),g,nothing)
    res4f = adjoint_sensitivities(soloop,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autojacvec=ZygoteVJP()),g,nothing)
    res4g = adjoint_sensitivities(soloop,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autodiff=false,autojacvec=false),g,nothing)
    res4h = adjoint_sensitivities(soloop,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(autodiff=true,autojacvec=false),g,nothing)

    @test norm(res_analytical' .- res4a) < 1e-7
    @test norm(res_analytical' .- res4b) < 1e-7
    @test norm(res_analytical' .- res4c) < 1e-7
    @test norm(res_analytical' .- res4d) < 1e-7
    @test norm(res_analytical' .- res4e) < 1e-7
    @test norm(res_analytical' .- res4f) < 1e-7
    @test norm(res_analytical' .- res4g) < 1e-7
    @test norm(res_analytical' .- res4h) < 1e-7
  end

  @testset "for u0: (should be zero, steady state does not depend on initial condition)" begin
    res5 = ForwardDiff.gradient(prob.u0) do u0
      tmp_prob = remake(prob,u0=u0)
      sol = solve(tmp_prob,
        SSRootfind(nlsolve = (f!,u0,abstol) -> (res=NLsolve.nlsolve(f!,u0,autodiff=:forward,method=:newton,iterations=Int(1e6),ftol=1e-14);res.zero) )
        )
      A = convert(Array,sol)
      g(A,p,nothing)
    end
    @test abs(dot(res5,res5)) < 1e-7
  end
end

using Zygote
@testset "concrete_solve derivatives steady state solver" begin

  function g1(u,p,t)
    sum(u)
  end

  function g2(u,p,t)
    sum((2.0.-u).^2)/2
  end

  u0 = zeros(2)
  p = [2.0,-2.0,1.0,-4.0]

  @testset "iip" begin
    function f!(du,u,p,t)
      du[1] = p[1] + p[2]*u[1]
      du[2] = p[3]*u[1] + p[4]*u[2]
    end
    prob = SteadyStateProblem(f!,u0,p)

    sol = solve(prob,DynamicSS(Rodas5()))
    res1 = adjoint_sensitivities(sol,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(),g1,nothing)
    res2 = adjoint_sensitivities(sol,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(),g2,nothing)

    dp1 = Zygote.gradient(p->sum(solve(prob,DynamicSS(Rodas5()),u0=u0,p=p,sensealg=SteadyStateAdjoint())),p)
    dp2 = Zygote.gradient(p->sum((2.0.-solve(prob,DynamicSS(Rodas5()),u0=u0,p=p,sensealg=SteadyStateAdjoint())).^2)/2.0,p)

    @test res1 ≈ dp1[1] rtol=1e-12
    @test res2 ≈ dp2[1] rtol=1e-12

    res1 = Zygote.gradient(p->sum(Array(solve(prob,DynamicSS(Rodas5()),u0=u0,p=p,sensealg=SteadyStateAdjoint()))[1]),p)
    dp1 = Zygote.gradient(p->sum(solve(prob,DynamicSS(Rodas5()),u0=u0,p=p,save_idxs=1:1,sensealg=SteadyStateAdjoint())),p)
    dp2 = Zygote.gradient(p->solve(prob,DynamicSS(Rodas5()),u0=u0,p=p,save_idxs=1,sensealg=SteadyStateAdjoint())[1],p)
    @test res1[1] ≈ dp1[1] rtol=1e-10
    @test res1[1] ≈ dp2[1] rtol=1e-10
  end

  @testset "oop" begin
    function f(u,p,t)
      dx = p[1] + p[2]*u[1]
      dy = p[3]*u[1] + p[4]*u[2]
      [dx,dy]
    end
    proboop = SteadyStateProblem(f,u0,p)

    soloop = solve(proboop,DynamicSS(Rodas5()))
    res1oop = adjoint_sensitivities(soloop,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(),g1,nothing)
    res2oop = adjoint_sensitivities(soloop,DynamicSS(Rodas5()),sensealg=SteadyStateAdjoint(),g2,nothing)

    dp1oop = Zygote.gradient(p->sum(solve(proboop,DynamicSS(Rodas5()),u0=u0,p=p,sensealg=SteadyStateAdjoint())),p)
    dp2oop = Zygote.gradient(p->sum((2.0.-solve(proboop,DynamicSS(Rodas5()),u0=u0,p=p,sensealg=SteadyStateAdjoint())).^2)/2.0,p)

    @test res1oop ≈ dp1oop[1] rtol=1e-12
    @test res2oop ≈ dp2oop[1] rtol=1e-12

    res1oop = Zygote.gradient(p->sum(Array(solve(proboop,DynamicSS(Rodas5()),u0=u0,p=p,sensealg=SteadyStateAdjoint()))[1]),p)
    dp1oop = Zygote.gradient(p->sum(solve(proboop,DynamicSS(Rodas5()),u0=u0,p=p,save_idxs=1:1,sensealg=SteadyStateAdjoint())),p)
    dp2oop = Zygote.gradient(p->solve(proboop,DynamicSS(Rodas5()),u0=u0,p=p,save_idxs=1,sensealg=SteadyStateAdjoint())[1],p)
    @test res1oop[1] ≈ dp1oop[1] rtol=1e-10
    @test res1oop[1] ≈ dp2oop[1] rtol=1e-10
  end
end


######
# methods that converge quickly even when α is non-simple
# These usually take 1 additional function call compared to a similar ordered
# method
######

"""
    Roots.Order1B()
    Roots.King()

A superlinear (order `1.6...`) modification of the secant method for multiple roots.
Presented in A SECANT METHOD FOR MULTIPLE ROOTS, by RICHARD F. KING, BIT 17 (1977), 321-328

The basic idea is similar to Schroder's method: apply the secant method
to  `f/f'`. However, this uses `f' ~ fp = (fx - f(x-fx))/fx` (a Steffensen step). In
this implementation, `Order1B`, when `fx` is too big, a single secant step of `f`
is used.

The *asymptotic* error, `eᵢ = xᵢ - α`, is given by
`eᵢ₊₂ = 1/2⋅G''/G'⋅ eᵢ⋅eᵢ₊₁ + (1/6⋅G'''/G' - (1/2⋅G''/G'))^2⋅eᵢ⋅eᵢ₊₁⋅(eᵢ+eᵢ₊₁)`.

"""
struct King <: AbstractSecant end
struct Order1B <: AbstractSecant end

struct KingState{T,S} <: AbstractUnivariateZeroState{T,S}
    xn1::T
    xn0::T
    fxn1::S
    fxn0::S
    G0::S
end

function init_state(M::Union{King, Order1B}, F, x₀, x₁::T, fx₀, fx₁::S) where {T,S}
    fₛ₀ = F(x₀ - fx₀*oneunit(T)/oneunit(S))
    G₀ = -fx₀^2 / (fₛ₀ - fx₀)
    KingState(x₁, x₀, fx₁, fx₀, G₀)
end
initial_fncalls(::Union{King, Order1B}) = 3

function update_state(method::Order1B, F, o::KingState{T,S}, options, l=NullTracks())  where {T, S}
    if do_guarded_step(Order1B(), o)
        state, flag = update_state(Order1(), F, o, options, l)

        x0,fx0 = state.xn0, state.fxn0 # clunky! Need to update G₀ after Order1() step
        fₛ = F(x0 - fx0*oneunit(T)/oneunit(S))
        incfn(l)
        G₀ = -fx0^2 / (fₛ - fx0)
        @set! state.G0 = G₀

        return (state, flag)
    else
        update_state(King(), F, o, options, l)
    end
end

function update_state(method::King, F,
                      o::KingState{T,S}, options, l=NullTracks())  where {T, S}

    x0, x1 = o.xn0, o.xn1
    fx0, fx1 = o.fxn0, o.fxn1
    G₀ = o.G0
    fₛ₁ = F(x1 - fx1*oneunit(T)/oneunit(S))
    incfn(l)
    G₁ = -fx1^2 / (fₛ₁ - fx1)

    m = (x1-x0) / (G₁-G₀) # approximate value of `m`, the multiplicity

    if abs(m) <= 1e-2 * oneunit(m)
        log_message(l, "Estimate for multiplicity has issues. ")
        return (o, true)
    end

    Δ = G₁ * (x1 - x0) / (G₁ - G₀)

    if isissue(Δ)
        log_message(l, "Increment `Δx` has issues. ")
        return o, true
    end

    x0, x1 = x1, x1 - Δ
    fx0, fx1 = fx1, F(x1)
    incfn(l)

    @set! o.xn0 = x0
    @set! o.xn1 = x1
    @set! o.fxn0 = fx0
    @set! o.fxn1 = fx1
    @set! o.G0 = G₁

    return o, false

end

##################################################

### Order2B() Esser method
"""
    Roots.Order2B()
    Roots.Esser()

Esser's method. This is a quadratically convergent method that, like
Schroder's method, does not depend on the multiplicity of the
zero. Schroder's method has update step `x - r2/(r2-r1) * r1`, where `ri =
fⁱ⁻¹/fⁱ`. Esser approximates `f' ~ f[x-h, x+h], f'' ~
f[x-h,x,x+h]`, where `h = fx`, as with Steffensen's method, Requiring 3
function calls per step. The implementation `Order2B` uses a secant
step when `|fx|` is considered too large.

Esser, H. Computing (1975) 14: 367. DOI: [10.1007/BF02253547](https://doi.org/10.1007/BF02253547)
Eine stets quadratisch konvergente Modifikation des Steffensen-Verfahrens

Example
```
f(x) = cos(x) - x
g(x) = f(x)^2
x0 = pi/4
find_zero(f, x0, Order2(), verbose=true)        #  3 steps / 7 function calls
find_zero(f, x0, Roots.Order2B(), verbose=true) #  4 / 9
find_zero(g, x0, Order2(), verbose=true)        #  22 / 45
find_zero(g, x0, Roots.Order2B(), verbose=true) #  4 / 10
```
"""
struct Esser <: AbstractSecant end
struct Order2B <: AbstractSecant end

function update_state(method::Order2B, fs, o::AbstractUnivariateZeroState{T,S}, options, l=NullTracks())  where {T, S}
     update_state_guarded(method, Secant(), Esser(), fs, o, options, l)
end

function update_state(method::Esser, F,
                      o::AbstractUnivariateZeroState{T,S}, options, l=NullTracks())  where {T, S}

    x1, fx1 = o.xn1, o.fxn1

    f0 = fx1

    f1  = F(x1 + f0 * oneunit(T) / oneunit(S))
    f_1 = F(x1 - f0 * oneunit(T) / oneunit(S))
    incfn(l, 2)

    # h = f0
    # r1 = f/f' ~ f/f[x+h,x-h]
    # r2 = f'/f'' ~ f[x+h, x-h]/f[x-h,x,x+h]
    r1 = f0 * 2*f0 / (f1 - f_1) * oneunit(T) / oneunit(S)
    r2 = (f1 - f_1)/(f1 - 2*f0 + f_1) * f0/2 * oneunit(T) / oneunit(S)

    k = r2/(r2-r1)  # ~ m

    if abs(k) <= 1e-2 * oneunit(k)
        log_message(l, "Estimate for multiplicity had issues. ")
        return o, true
    end

    delta = k * r1

    if isissue(delta)
        log_message(l, "Increment `Δx` has issues. ")
        return o, true
    end

    x0, x1 = x1, x1 - delta
    fx0, fx1 = fx1, F(x1)
    incfn(l)

    @set! o.xn0 = x0
    @set! o.xn1 = x1
    @set! o.fxn0 = fx0
    @set! o.fxn1 = fx1

    return o, false

end

"""
    Roots.Schroder()

Schröder's method, like Halley's method, utilizes `f`, `f'`, and
`f''`. Unlike Halley it is quadratically converging, but this is
independent of the multiplicity of the zero (cf. Schröder, E. "Über
unendlich viele Algorithmen zur Auflösung der Gleichungen."
Math. Ann. 2, 317-365, 1870;
[mathworld](http://mathworld.wolfram.com/SchroedersMethod.html)). Schröder's
method applies Newton's method to `f/f'`, a function with all
simple zeros.

Example
```
m = 2
f(x) = (cos(x)-x)^m
fp(x) = (-x + cos(x))*(-2*sin(x) - 2)
fpp(x) = 2*((x - cos(x))*cos(x) + (sin(x) + 1)^2)
find_zero((f, fp, fpp), pi/4, Roots.Halley())    # 14 steps
find_zero((f, fp, fpp), 1.0, Roots.Schroder())    # 3 steps
```

(Whereas, when `m=1`, Halley is 2 steps to Schröder's 3.)

If function evaluations are expensive one can pass in a function which
returns `(f, f/f',f'/f'')` as follows

```
find_zero(x -> (sin(x), sin(x)/cos(x), -cos(x)/sin(x)), 3.0, Roots.Schroder())
```

This can be advantageous if the derivatives are easily computed from
the value of `f`, but otherwise would be expensive to compute.

The error, `eᵢ = xᵢ - α`, is the same as `Newton` with `f` replaced by `f/f'`.

"""
struct Schroder <: AbstractHalleyLikeMethod
end
const Schroeder = Schroder # either spelling
const Schröder = Schroder

## Shroder-like methods
function init_state(M::Schroder, F::Callable_Function, x)
    x₀ = float(first(x))
    fx₀, (Δ, ΔΔ) = F(x₀)
    x₁ = x₀ - ΔΔ/(ΔΔ - Δ) * Δ # m*r1
    state = init_state(M, F, x₀, x₁, fx₀, fx₀)
end

function update_state(method::Schroder, F, o::AbstractUnivariateZeroState{T,S},
                      options::UnivariateZeroOptions, l=NullTracks()) where {T,S}
    xn = o.xn1
    fxn = o.fxn1
    r1, r2 = o.Δ, o.ΔΔ

    Δ =  r2 / (r2 - r1) * r1  # m * r1

    if isissue(Δ)
        log_message(l, "Issue with increment")
        return o, true
    end

    xn1::T = xn - Δ

    fxn1::S, (r1::T, r2::T) = F(xn1)
    incfn(l,3)

    @set! o.xn0 = xn
    @set! o.xn1 = xn1
    @set! o.fxn0 = fxn
    @set! o.fxn1 = fxn1
    @set! o.Δ = r1
    @set! o.ΔΔ = r2

    return o, false
end

## Thukral 3,4,5 (2) is Schroder
"""
    AbstractThukralBMethod

Abstract type for `ThukralXB` methods for `X` being `2`,`3`,`4`, or `5`.

These are a family of methods which are
* efficient (order `X`) for non-simple roots (e.g. `Thukral2B` is the `Schroder` method)
* taking `X+1` function calls per step
* require `X` derivatives. These can be passed as a tuple of functions, `(f, f', f'', …)`, *or* as
a function returning the ratios: `x -> (f(x), f(x)/f'(x), f'(x)/f''(x), …)`.

## Example

```julia
using ForwardDiff
Base.adjoint(f::Function)  = x  -> ForwardDiff.derivative(f, float(x))
f(x) = (exp(x) + x - 2)^6
x0 = 1/4
find_zero((f, f', f''), x0, Roots.Halley())               # 14 iterations; 45 function evaluations
find_zero((f, f', f''), big(x0), Roots.Thukral2B())       #  4 iterations; 15 function evaluations
find_zero((f, f', f'', f'''), big(x0), Roots.Thukral3B()) #  3 iterations; 16 function evaluations
```

## Refrence

*Introduction to a family of Thukral ``k``-order method for finding multiple zeros of nonlinear equations*,
R. Thukral, JOURNAL OF ADVANCES IN MATHEMATICS 13(3):7230-7237, DOI: [10.24297/jam.v13i3.6146](https://doi.org/10.24297/jam.v13i3.6146).
"""
abstract type AbstractThukralBMethod <: AbstractHalleyLikeMethod end
initial_fncalls(M::AbstractThukralBMethod) = fn_argout(M)

struct Thukral2B <: AbstractThukralBMethod end
fn_argout(::Thukral2B) = 3

struct Thukral3B <: AbstractThukralBMethod end
fn_argout(::Thukral3B) = 4

struct Thukral4B <: AbstractThukralBMethod end
fn_argout(::Thukral4B) = 5

struct Thukral5B <: AbstractThukralBMethod end
fn_argout(::Thukral5B) = 6

struct ThukralBState{N,T,S} <: AbstractUnivariateZeroState{T,S}
    xn1::T
    xn0::T
    Δs::NTuple{N, T}
    fxn1::S
    fxn0::S
end

function init_state(M::AbstractThukralBMethod, F::Callable_Function, x)
    x₁ = float(first(x))
    fx₁, Δs = F(x₁)
    state = init_state(M, F, nan(x₁), x₁, nan(fx₁), fx₁; Δs=Δs)
end

function init_state(M::AbstractThukralBMethod, F, x₀, x₁::T, fx₀, fx₁; Δs = nothing) where {T}
    ThukralBState(x₁, x₀, NTuple{fn_argout(M)-1,T}(Δs), fx₁, fx₀)
end

function update_state(M::AbstractThukralBMethod, F, o::AbstractUnivariateZeroState{T,S},
                      options::UnivariateZeroOptions, l=NullTracks()) where {T,S}

    x₀ = o.xn1

    Δ = compute_thukral_Δ(M, o)

    isissue(Δ) && return (o, true)
    x₁ = x₀ - Δ
    fx₁, Δs = F(x₁)
    incfn(l, fn_argout(M))

    @set! o.xn0 = x₀
    @set! o.fxn0 = o.fxn1
    @set! o.Δs =  NTuple{fn_argout(M)-1,T}(Δs)
    @set! o.xn1 = x₁
    @set! o.fxn1 = fx₁

    return (o, false)
end

function compute_thukral_Δ(M::Thukral2B, o)
    r₁, r₂ = o.Δs
    t₁, t₂ = 1/r₁, 1/r₂
    Δ = one(o.xn1)
    Δ /= (t₁ - t₂)
    Δ
end

function compute_thukral_Δ(M::Thukral3B, o)
    r₁, r₂, r₃ = o.Δs
    t₁, t₂, t₃ = 1/r₁, 1/r₂, 1/r₃
    Δ = (2t₁ - 2t₂)
    Δ /= (2t₁^2 - 3t₁*t₂ + t₂*t₃)
    Δ
end

function compute_thukral_Δ(M::Thukral4B, o)
    r₁, r₂, r₃, r₄ = o.Δs
    t₁, t₂, t₃, t₄ = 1/r₁, 1/r₂, 1/r₃, 1/r₄
    Δ = 6t₁^2 - 9t₁*t₂ +3t₂*t₃
    Δ /= 6t₁^3 - 12 * t₁^2*t₂ + 4t₁*t₂*t₃ - t₂*t₃*t₄ + 3*t₁*t₂^2
    Δ
end

function compute_thukral_Δ(M::Thukral5B, o)
    r₁, r₂, r₃, r₄, r₅ = o.Δs
    t₁, t₂, t₃, t₄, t₅ = 1/r₁, 1/r₂, 1/r₃, 1/r₄, 1/r₅
    Δ = 24*t₁^3 - 48t₁^2*t₂  + 16*t₁*t₂*t₃ - 4*t₂*t₃*t₄ + 12t₁*t₂^2
    Δ /= 24t₁^4  - 60t₁^3*t₂ + 20*t₁^2*t₂*t₃ - 5*t₁*t₂*t₃*t₄ + 30t₁^2*t₂^2 - 10*t₁*t₂^2*t₃ + t₂*t₃*t₄*t₅
    Δ
end


function icrgen_qc(winrate, n1, n2;uncommon_icr_probs = [0, 0.85, 0.10, 0.05], rare_icr_probs = [0,0,2/3,1/3])
	first_icr = Dict( (w,uncommon_icr_probs) for w in 0:7)
	second_icr = deepcopy(first_icr)
	second_icr[6] = rare_icr_probs
	second_icr[7] = rare_icr_probs
	third_icr = deepcopy(second_icr)
	for i in 4:5
		third_icr[i] = rare_icr_probs
	end

	avg_first_icr = sum( pdfnegbin_truncated(k; p = winrate, r = 3, maxwins = 7) .* first_icr[k] for k in 0:7 )
	avg_second_icr = sum( pdfnegbin_truncated(k; p = winrate, r = 3, maxwins = 7) .* second_icr[k] for k in 0:7 )
	avg_third_icr = sum( pdfnegbin_truncated(k; p = winrate, r = 3, maxwins = 7) .* third_icr[k] for k in 0:7 )

	if n1 > 0
		icrvec = [[ Categorical(avg_first_icr), Categorical(avg_second_icr), Categorical(avg_third_icr) ] for i in 1:n1]
		for i in 1:n2
			push!(icrvec, Categorical{Float64}[])
		end
		icrcycle = Cycle(  icrvec )
	else
		icrcycle = Cycle( [Categorical{Float64}[]] )
	end

	icrcycle
end

function icrgen_singleton(winrate, n1, n2;uncommon_icr_probs = [0, 0.85, 0.10, 0.05], rare_icr_probs = [0,0,2/3,1/3])
	first_icr = Dict( (w,rare_icr_probs) for w in 0:5)
	second_icr = Dict( (w, uncommon_icr_probs) for w in 0:5)

	avg_first_icr = sum( pdfnegbin_truncated(k; p = winrate, r = 2, maxwins = 5) .* first_icr[k] for k in 0:5 )
	avg_second_icr = sum( pdfnegbin_truncated(k; p = winrate, r = 2, maxwins = 5) .* second_icr[k] for k in 0:5 )
	if n1 > 0
		icrvec = [[ Categorical(avg_first_icr), Categorical(avg_second_icr)] for i in 1:n1]
		for i in 1:n2
			push!(icrvec, Categorical{Float64}[])
		end
		icrcycle = Cycle(  icrvec )
	else
		icrcycle = Cycle( [Categorical{Float64}[]] )
	end

	icrcycle
end


@testset "CUFFT" begin

using CuArrays.CUFFT

using CUDAdrv

using FFTW

# notes:
#   plan_bfft does not need separate testing since it is used by plan_ifft

N1 = 8
N2 = 32
N3 = 64
N4 = 8

MYRTOL = 1e-5
MYATOL = 1e-8

## complex

function out_of_place(X::AbstractArray{T,N}) where {T <: Complex,N}
    fftw_X = fft(X)
    d_X = CuArray(X)
    p = plan_fft(d_X)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    pinv = plan_ifft(d_Y)
    d_Z = pinv * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)

    pinv2 = inv(p)
    d_Z = pinv2 * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)
end

function in_place(X::AbstractArray{T,N}) where {T <: Complex,N}
    fftw_X = fft(X)
    d_X = CuArray(X)
    p = plan_fft!(d_X)
    p * d_X
    Y = collect(d_X)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    pinv = plan_ifft!(d_X)
    pinv * d_X
    Z = collect(d_X)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)
end

function batched(X::AbstractArray{T,N},region) where {T <: Complex,N}
    fftw_X = fft(X,region)
    d_X = CuArray(X)
    p = plan_fft(d_X,region)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test_maybe_broken isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    pinv = plan_ifft(d_Y,region)
    d_Z = pinv * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)
end

@testset for T in [ComplexF32, ComplexF64]

@testset "1D" begin
    dims = (N1,)
    X = rand(T, dims)
    out_of_place(X)
end
@testset "1D inplace" begin
    dims = (N1,)
    X = rand(T, dims)
    in_place(X)
end

@testset "2D" begin
    dims = (N1,N2)
    X = rand(T, dims)
    out_of_place(X)
end
@testset "2D inplace" begin
    dims = (N1,N2)
    X = rand(T, dims)
    in_place(X)
end

@testset "Batch 1D" begin
    dims = (N1,N2)
    X = rand(T, dims)
    batched(X,1)

    dims = (N1,N2)
    X = rand(T, dims)
    batched(X,2)

    dims = (N1,N2)
    X = rand(T, dims)
    batched(X,(1,2))
end

@testset "3D" begin
    dims = (N1,N2,N3)
    X = rand(T, dims)
    out_of_place(X)
end

@testset "3D inplace" begin
    dims = (N1,N2,N3)
    X = rand(T, dims)
    in_place(X)
end

@testset "Batch 2D (in 3D)" begin
    dims = (N1,N2,N3)
    for region in [(1,2),(2,3),(1,3)]
        X = rand(T, dims)
        batched(X,region)
    end

    X = rand(T, dims)
    @test_throws ArgumentError batched(X,(3,1))
end

@testset "Batch 2D (in 4D)" begin
    dims = (N1,N2,N3,N4)
    for region in [(1,2),(1,4),(3,4)]
        X = rand(T, dims)
        batched(X,region)
    end
    for region in [(1,3),(2,3),(2,4)]
        X = rand(T, dims)
        @test_throws ArgumentError batched(X,region)
    end

end

end

## real

function out_of_place(X::AbstractArray{T,N}) where {T <: Real,N}
    fftw_X = rfft(X)
    d_X = CuArray(X)
    p = plan_rfft(d_X)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    pinv = plan_irfft(d_Y,size(X,1))
    d_Z = pinv * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)

    pinv2 = inv(p)
    d_Z = pinv2 * d_Y
    Z = collect(d_Z)
    @test_maybe_broken isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)

    pinv3 = inv(pinv)
    d_W = pinv3 * d_X
    W = collect(d_W)
    @test isapprox(W, Y, rtol = MYRTOL, atol = MYATOL)
end

function batched(X::AbstractArray{T,N},region) where {T <: Real,N}
    fftw_X = rfft(X,region)
    d_X = CuArray(X)
    p = plan_rfft(d_X,region)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    pinv = plan_irfft(d_Y,size(X,region[1]),region)
    d_Z = pinv * d_Y
    Z = collect(d_Z)
    @test isapprox(Z, X, rtol = MYRTOL, atol = MYATOL)
end

@testset for T in [Float32, Float64]

@testset "1D" begin
    X = rand(T, N1)
    out_of_place(X)
end

@testset "Batch 1D" begin
    dims = (N1,N2)
    X = rand(T, dims)
    batched(X,1)

    dims = (N1,N2)
    X = rand(T, dims)
    batched(X,2)

    dims = (N1,N2)
    X = rand(T, dims)
    batched(X,(1,2))
end

@testset "2D" begin
    X = rand(T, N1,N2)
    out_of_place(X)
end

@testset "Batch 2D (in 3D)" begin
    dims = (N1,N2,N3)
    for region in [(1,2),(2,3),(1,3)]
        X = rand(T, dims)
        batched(X,region)
    end

    X = rand(T, dims)
    @test_throws ArgumentError batched(X,(3,1))
end

@testset "Batch 2D (in 4D)" begin
    dims = (N1,N2,N3,N4)
    for region in [(1,2),(1,4),(3,4)]
        X = rand(T, dims)
        batched(X,region)
    end
    for region in [(1,3),(2,3),(2,4)]
        X = rand(T, dims)
        @test_throws ArgumentError batched(X,region)
    end
end

@testset "3D" begin
    X = rand(T, N1, N2, N3)
    out_of_place(X)
end

end

## complex integer

function out_of_place(X::AbstractArray{T,N}) where {T <: Complex{<:Integer},N}
    fftw_X = fft(X)
    d_X = CuArray(X)
    p = plan_fft(d_X)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    d_Y = fft(d_X)
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)
end

@testset for T in [Complex{Int32}, Complex{Int64}]

@testset "1D" begin
    dims = (N1,)
    X = rand(T, dims)
    out_of_place(X)
end

end

## real integer

function out_of_place(X::AbstractArray{T,N}) where {T <: Integer,N}
    fftw_X = rfft(X)
    d_X = CuArray(X)
    p = plan_rfft(d_X)
    d_Y = p * d_X
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)

    d_Y = rfft(d_X)
    Y = collect(d_Y)
    @test isapprox(Y, fftw_X, rtol = MYRTOL, atol = MYATOL)
end

@testset for T in [Int32, Int64]

@testset "1D" begin
    X = rand(T, N1)
    out_of_place(X)
end

end

@testset "streams" begin
    X = rand(N1)
    d_X = CuArray(X)
    p = plan_fft(d_X)
    CUFFT.set_stream(p, CUDAdrv.CuDefaultStream())
end

end


## Preloads
using Statistics
using FFTW
using Plots
using BenchmarkTools
using Profile
using LinearAlgebra
using Measures
using ProgressMeter
push!(LOAD_PATH, pwd()*"/main")
using DHC_2DUtils

fhash = fink_filter_hash(1, 8, wd=2, nx=256, Omega=false)

randA = rand(128,128)

no_shift = DHC_compute(randA,fhash)

shift = DHC_compute(circshift(randA,(1,1)),fhash)

maximum(shift - no_shift)

coeff_vec = zeros(256*256,2452)
@showprogress for i=1:256
    @showprogress for j=1:256
        coeff_vec[i+(j-1)*256,:] = DHC_compute(circshift(randA,(i,j)),fhash)
    end
end

using Distributed
addprocs(7);

@everywhere begin
    using Statistics
    using FFTW
    using Plots
    using BenchmarkTools
    using Profile
    using LinearAlgebra
    using Measures
    using ProgressMeter
    push!(LOAD_PATH, pwd()*"/main")
    using DHC_2DUtils

    fhash = fink_filter_hash(1, 8, wd=2, nx=128, Omega=false)

    function calc_shift(x)
        i,j,img = x
        return DHC_compute(circshift(img,(i,j)),fhash)
    end
end

mnist_DHC_out = @showprogress pmap(calc_shift, Iterators.product(1:128,1:128,[randA]))

mnist_DHC_out = hcat(mnist_DHC_out...)

out_test = mnist_DHC_out

maximum(out_test .- out_test[:,1])


# fractals.jl

###############################################################################
## The Mandelbrot set
###############################################################################
export Mandelbrot,mandelbrot
# TODO: fractals should be subsets of the complex plane
struct Mandelbrot{T} <: EuclideanDomain{2,T}
    maxiter     ::  Int
    threshold   ::  T
    maskcache   ::  Dict
    box         ::  Domain
end

function mandelbrot(maxiter = 1000, threshold = 1000.0)
    box = rectangle(-1.0, 0.35, -0.65, 0.65)
    M1 = 136
    M2 = 200
    mask1 = computemandelbrotgrid(EquispacedGrid(M1,-1.0, 0.3)× EquispacedGrid(M2,-0.65, 0.65), maxiter, threshold)
    mask2 = computemandelbrotgrid(EquispacedGrid(M1,-1.0, 0.3)× EquispacedGrid(M2,-0.65, 0.65), maxiter, threshold)
    cache = Dict{Int,BitArray{2}}()
    cache[M1] = mask1
    cache[M2] = mask2
    Mandelbrot(maxiter, threshold, cache, box)
end

function mandelbrotiteration(x, maxiter, threshold)
    c = 2*(x[1]+1im*x[2])
    T = typeof(c)
    z = zero(T)
    iter = 0
    while abs(z) < threshold && iter < maxiter
        z = z^2 + c
        iter += 1
    end
    abs(z) < threshold
end

function computemandelbrotgrid(grid, maxiter, threshold)
    mask = BitArray(undef, size(grid))
    fill!(mask,0)
    for i_2 = 1:size(grid, 2)
        for i_1 = 1:size(grid, 1)
            mask[i_1,i_2] = mandelbrotiteration(grid[i_1,i_2], maxiter, threshold)
        end
    end
    mask
end

function indomain_broadcast(grid::ProductGrid, m::Mandelbrot)
    if (infimum(support(grid)) ≈ minimum(boundingbox(m))) && (supremum(support(grid)) ≈ maximum(boundingbox(m)))
        if haskey(m.maskcache, size(grid,1))
            mask = m.maskcache[size(grid,1)]
        else # compute mask and cache it
            mask = computemandelbrotgrid(grid, m.maxiter, m.threshold)
            m.maskcache[size(grid,1)] = mask
        end
    else # Don't cache if the grid doesn't match the bounding box
        mask = computemandelbrotgrid(grid, m.maxiter, m.threshold)
    end
    mask
end

function isapprox(t::Tuple{T,T}, v::SVector{2,T}) where {T}
    return t[1]≈v[1] && t[2]≈v[2]
end
function isapprox(v::SVector{2,T}, t::Tuple{T,T}) where {T}
    return t[1]≈v[1] && t[2]≈v[2]
end
boundingbox(m::Mandelbrot) = m.box

indomain(x::SVector{2}, m::Mandelbrot) = mandelbrotiteration(x, m.maxiter, m.threshold)

show(io::IO, m::Mandelbrot) = print(io, "The Mandelbrot set")

################
## Julia sets
################
export JuliaSet, juliaset
struct JuliaSet{T} <: EuclideanDomain{2,T}
    c           ::  Complex{T}
    maxiter     ::  Int
    maskcache   ::  Dict
    box         ::  Domain
end

function juliaset(c = -0.122565+0.744866im, maxiter = 1000)
    box = rectangle(-0.2, 1.2, -0.4, 0.4)

    mask1 = computejuliasetgrid(EquispacedGrid(200,-0.2, 1.2)× EquispacedGrid(200,-0.4, 0.4), c, maxiter)
    mask2 = computejuliasetgrid(EquispacedGrid(200,-0.2, 1.2)× EquispacedGrid(200,-0.4, 0.4), c, maxiter)
    cache = Dict{Int,BitArray{2}}()
    cache[100] = mask1
    cache[200] = mask2
    JuliaSet(c, maxiter, cache, box)
end

function juliasetiteration(x, c, maxiter)
    gamma = 1 + sqrt(1-4*c)
    z = x[1] + 1im*x[2]
    for i = 1:maxiter
        z = gamma*z*(1-z)
    end
    abs(z) < 1000
end

function computejuliasetgrid(grid, c, maxiter)
    m = size(grid)
    mask = BitArray(undef, m)
    fill!(mask, 0)
    for i_2 = 1:m[2]
        for i_1 = 1:m[1]
            a,b = grid[i_1,i_2]
            mask[i_1,i_2] = juliasetiteration(grid[i_1,i_2], c, maxiter)
        end
    end
    mask
end

function indomain_broadcast(grid, js::JuliaSet)
    if (map(x->infimum(support(x)), elements(grid)) ≈ infimum(boundingbox(js))) && (map(x->supremum(support(x)),elements(grid)) ≈ supremum(boundingbox(js)))
        if haskey(js.maskcache, size(grid,1))
            mask = js.maskcache[size(grid,1)]
        else # compute mask and cache it
            mask = computejuliasetgrid(grid, js.c, js.maxiter)
            js.maskcache[size(grid,1)] = mask
        end
    else # Don't cache if the grid doesn't match the bounding box
        mask = computejuliasetgrid(grid, js.c, js.maxiter)
    end
    mask
end

indomain(x::SVector{2}, js::JuliaSet) = juliasetiteration(x, js.c, js.maxiter)

show(io::IO, js::JuliaSet) = print(io, "A particular Julia Set also known as the Douady rabbit")

boundingbox(js::JuliaSet) = js.box


using Base.Iterators
using DataStructures
using ArgParse

function readInput(filename::String)
    regex = r"Player (\d+) starting position: (\d+)"
    player1, value1 = -1, -1
    player2, value2 = -1, -1

    open(filename) do file
        for line in readlines(file)
            m = match(regex, line)
            if m[1] == "1"
                player1, value1 = parse.(Int, (m[1], m[2]))
            else
                player2, value2 = parse.(Int, (m[1], m[2]))
            end
        end
    end
    return (player1, value1), (player2, value2)
end

function playGame(p1, p2; winning_score::Int = 1000)
    roll_dice(offset::Int) = sum(((offset:offset+2) .- 1) .% 100 .+ 1)

    p1, p2 = p1 - 1, p2 - 1
    s1, s2, dice_i = 0, 0, 0
    for i in countfrom(1)
        if i % 2 == 1
            p1 = (p1 + roll_dice(dice_i + 1)) % 10
            s1 += p1 + 1

            s1 >= winning_score && return i * s2 *3
        else
            p2 = (p2 + roll_dice(dice_i + 1)) % 10
            s2 += p2 + 1

            s2 >= winning_score && return i * s1 * 3
        end

        dice_i = (dice_i + 3) % 100
    end
end

function dirac_game(p1, p2; winning_score::Int=21)
    s1, p1 = 0, p1 - 1
    s2, p2 = 0, p2 - 1
    wins = [0, 0]

    three_dice = vec(sum.(collect(product(1:3, 1:3, 1:3))))

    prev_step = Dict{Tuple{Int, Int, Int, Int}, Int}((p1, s1, p2, s2) => 1)
    turn = true
    while length(prev_step) != 0
        next_step = DefaultDict{Tuple{Int, Int, Int, Int}, Int}(0)
        for k in keys(prev_step)
            v = prev_step[k]
            p1, s1, p2, s2 = k

            for dice in three_dice
                if turn
                    s1_ = s1 + (p1 + dice) % 10 + 1
                    if s1_ >= winning_score
                        wins[1] += v
                    else
                        p1_ = (p1 + dice) % 10
                        next_step[(p1_, s1_, p2, s2)] += v
                    end
                else
                    s2_ = s2 + (p2 + dice) % 10 + 1
                    if s2_ >= winning_score
                        wins[2] += v
                    else
                        p2_ = (p2 + dice) % 10
                        next_step[(p1, s1, p2_, s2_)] += v
                    end
                end
            end
        end
        turn = !turn
        prev_step = next_step
    end
    return maximum(wins)
end

function parse_cli()
    s = ArgParseSettings()
    @add_arg_table s begin
        "--example"
            help = "Calculate example"
            action = :store_true
    end
    return parse_args(s)
end

args = parse_cli();
if args["example"]
    (_, p1), (_, p2) = readInput("example.txt")

    score = playGame(p1, p2)
    @assert score == 739785

    w = dirac_game(p1, p2)
    @assert w == 444356092776315
end

(_, p1), (_, p2) = readInput("input.txt")
score = playGame(p1, p2)
w = dirac_game(p1, p2)

println(score)
println(w)


# Tests on Multivariate LogNormal distributions

using Distributions,  PDMats
using LinearAlgebra, Random, Test

####### Core testing procedure

function test_mvlognormal(g::MvLogNormal, n_tsamples::Int=10^6,
                          rng::Union{AbstractRNG, Missing} = missing)
    d = length(g)
    mn = mean(g)
    md = median(g)
    mo = mode(g)
    S = cov(g)
    s = var(g)
    e = entropy(g)
    @test partype(g) == Float64
    @test isa(mn, Vector{Float64})
    @test isa(md, Vector{Float64})
    @test isa(mo, Vector{Float64})
    @test isa(s, Vector{Float64})
    @test isa(S, Matrix{Float64})
    @test length(mn) == d
    @test length(md) == d
    @test length(mo) == d
    @test length(s) == d
    @test size(S) == (d, d)
    @test s          ≈ diag(S)
    @test md         ≈ exp.(mean(g.normal))
    @test mn         ≈ exp.(mean(g.normal) + var(g.normal)/2)
    @test mo         ≈ exp.(mean(g.normal) - var(g.normal))
    @test entropy(g) ≈ d*(1 + Distributions.log2π)/2 + logdetcov(g.normal)/2 + sum(mean(g.normal))
    gg = typeof(g)(MvNormal(params(g)...))
    @test Vector(g.normal.μ) == Vector(gg.normal.μ)
    @test Matrix(g.normal.Σ) == Matrix(gg.normal.Σ)
    @test insupport(g,ones(d))
    @test !insupport(g,zeros(d))
    @test !insupport(g,-ones(d))

    # sampling
    if ismissing(rng)
        X = rand(g, n_tsamples)
    else
        X = rand(rng, g, n_tsamples)
    end
    emp_mn = vec(mean(X, dims=2))
    emp_md = vec(median(X, dims=2))
    Z = X .- emp_mn
    emp_cov = (Z * Z') ./ n_tsamples
    for i = 1:d
        @test isapprox(emp_mn[i]   , mn[i] , atol=(sqrt(s[i] / n_tsamples) * 8.0))
    end
    for i = 1:d
        @test isapprox(emp_md[i]   , md[i] , atol=(sqrt(s[i] / n_tsamples) * 8.0))
    end
    for i = 1:d, j = 1:d
        @test isapprox(emp_cov[i,j], S[i,j], atol=(sqrt(s[i] * s[j]) * 20.0) / sqrt(n_tsamples))
    end

    # evaluation of logpdf and pdf
    for i = 1:min(100, n_tsamples)
        @test logpdf(g, X[:,i]) ≈ log(pdf(g, X[:,i]))
    end
    @test logpdf(g, X) ≈ log.(pdf(g, X))
    @test isequal(logpdf(g, zeros(d)),-Inf)
    @test isequal(logpdf(g, -mn),-Inf)
    @test isequal(pdf(g, zeros(d)),0.0)
    @test isequal(pdf(g, -mn),0.0)

    # test the location and scale functions
    @test isapprox(location(g), location(MvLogNormal,:meancov,mean(g),cov(g))   , atol=1e-8)
    @test isapprox(location(g), location(MvLogNormal,:mean,mean(g),scale(g))    , atol=1e-8)
    @test isapprox(location(g), location(MvLogNormal,:median,median(g),scale(g)), atol=1e-8)
    @test isapprox(location(g), location(MvLogNormal,:mode,mode(g),scale(g))    , atol=1e-8)
    @test isapprox(scale(g)   , scale(MvLogNormal,:meancov,mean(g),cov(g))      , atol=1e-8)

    @test isapprox(location(g), location!(MvLogNormal,:meancov,mean(g),cov(g),zero(mn))   , atol=1e-8)
    @test isapprox(location(g), location!(MvLogNormal,:mean,mean(g),scale(g),zero(mn))    , atol=1e-8)
    @test isapprox(location(g), location!(MvLogNormal,:median,median(g),scale(g),zero(mn)), atol=1e-8)
    @test isapprox(location(g), location!(MvLogNormal,:mode,mode(g),scale(g),zero(mn))    , atol=1e-8)
    @test isapprox(scale(g)   , Distributions.scale!(MvLogNormal,:meancov,mean(g),cov(g),zero(S)), atol=1e-8)

    lc1,sc1 = params(MvLogNormal,mean(g),cov(g))
    lc2,sc2 = params!(MvLogNormal,mean(g),cov(g),similar(mn),similar(S))
    @test isapprox(location(g), lc1, atol=1e-8)
    @test isapprox(location(g), lc2, atol=1e-8)
    @test isapprox(scale(g)   , sc1, atol=1e-8)
    @test isapprox(scale(g)   , sc2, atol=1e-8)
end

####### Validate results for a single-dimension MvLogNormal by comparing with univariate LogNormal
@testset "Comparing results from MvLogNormal with univariate LogNormal" begin
l1 = LogNormal(0.1,0.4)
l2 = MvLogNormal(0.1*ones(1),0.4)
@test [mean(l1)]     ≈ mean(l2)
@test [median(l1)]   ≈ median(l2)
@test [mode(l1)]     ≈ mode(l2)
@test [var(l1)]      ≈ var(l2)
@test  entropy(l1)   ≈ entropy(l2)
@test logpdf(l1,5.0) ≈ logpdf(l2,[5.0])
@test pdf(l1,5.0)    ≈ pdf(l2,[5.0])
@test (Random.seed!(78393) ; [rand(l1)]) == (Random.seed!(78393) ; rand(l2))
@test [rand(MersenneTwister(78393), l1)] == rand(MersenneTwister(78393), l2)
end

###### General Testing

mu = [0.1, 0.2, 0.3]
va = [0.16, 0.25, 0.36]
C = [0.4 -0.2 -0.1; -0.2 0.5 -0.1; -0.1 -0.1 0.6]

@testset "Testing MvLogNormal with $key" for (key, rng) in
    Dict("rand(...)" => missing,
         "rand(rng, ...)" => MersenneTwister(123))

    @testset "Testing MvLogNormal with $key for $(typeof(g)) with $(Distributions.distrname(g.normal))" for (g, μ, Σ) in [
        (MvLogNormal(mu,PDMats.PDMat(C)), mu, C),
        (MvLogNormal(PDMats.PDiagMat(sqrt.(va))), zeros(3), Matrix(Diagonal(va))),
        (MvLogNormal(mu, sqrt(0.2)), mu, Matrix(0.2I, 3, 3)),
        (MvLogNormal(3, sqrt(0.2)), zeros(3), Matrix(0.2I, 3, 3)),
        (MvLogNormal(mu, sqrt.(va)), mu, Matrix(Diagonal(va))),
        (MvLogNormal(sqrt.(va)), zeros(3), Matrix(Diagonal(va))),
        (MvLogNormal(mu, C), mu, C),
        (MvLogNormal(C), zeros(3), C) ]

        m,s = params(g)
        @test Vector(m) ≈ μ
        test_mvlognormal(g, 10^4, rng)
    end
end

##### Constructors and conversions
d = MvLogNormal(Array{Float32}(mu), PDMats.PDMat(Array{Float32}(C)))
@test typeof(convert(MvLogNormal{Float64}, d)) == typeof(MvLogNormal(mu, PDMats.PDMat(C)))
@test typeof(convert(MvLogNormal{Float64}, d.normal.μ, d.normal.Σ)) == typeof(MvLogNormal(mu, PDMats.PDMat(C)))


# multigrid transfer operators
function intergrid_1d_interp(N, Minv)

    m = size(Minv)[1]
    ic = Int(ceil(m/2.0))
    sten = transpose(Minv) * Float64[(0.5)^(j-1)/factorial(j-1) for j=1:m]

    Ih = zeros(Float64, (N,N))
    Is = I + Ih
    for i=1:m
        k = i - ic
        vec = sten[i] * ones(N - abs(k))
        Ih += diagm(k => vec)
    end

    Iout = zeros(Float64, (2*N,N))
    Iout[1:2:end,:] = Is
    Iout[2:2:end,:] = Ih
    return Iout
end

function transfer_matrices(Nx0, Ny0, Mxinv, Myinv, P0, P1)

    Ix = intergrid_1d_interp(Nx0, Mxinv)
    Iy = intergrid_1d_interp(Ny0, Myinv)
    Ic = kron(sparse(Iy), sparse(Ix))
    Rc = transpose(Ic) / 4.0

    Restr = P0 * Rc * transpose(P1)
    Interp = P1 * Ic * transpose(P0)
    return Restr, Interp
end

# the idea behind this struct is that instances represent a matrix at multiple resolutions. Not only
# can the same matrix at different resolutions be represented by the same variable, but the algebra can be
# simplified in a more elegant way too.
struct MultiResMatrix

    mats :: Array{SparseMatrixCSC{Float64, Int32}, 1}
    lvl_by_row_size :: Dict{Int64, Int64}
    lvl_by_column_size :: Dict{Int64, Int64}

end

# no need for Julia to see this struct as a custom array, but we can overload the multiplication
# function so that multiplying from the left or the right automatically works. Both M*X and X*M will multiply
# X by the correct version of M.
function Base.:*(A::MultiResMatrix, X)
    lvl = A.lvl_by_row_size[size(X)[1]]
    return A.mats[lvl] * X
end
function Base.:*(X, A::MultiResMatrix)
    lvl = A.lvl_by_column_size[size(X)[2]]
    return X * A.mats[lvl]
end

function MultiResMatrix(A::SparseMatrixCSC, grd::Grid)

    As = SparseMatrixCSC[A]
    lvl_by_row_size = Dict(size(A)[2] => 1)
    lvl_by_column_size = Dict(size(A)[1] => 1)

    for l=1:length(grd.restrictions)
        Restr = grd.restrictions[l]
        Interp = grd.interpolations[end-(l-1)]
        Ap = Restr * As[end] * Interp
        append!(As, [Ap])

        lvl_by_row_size[size(Ap)[2]] = l+1
        lvl_by_column_size[size(Ap)[1]] = l+1
    end

    return MultiResMatrix(As, lvl_by_row_size, lvl_by_column_size)
end

function setup_jacobi_smoothing(A::SparseMatrixCSC, w::Float64, grd::Grid)

    A_mr = MultiResMatrix(A, grd)
    Ms = SparseMatrixCSC[]
    Ds = SparseMatrixCSC[]

    for a in A_mr.mats
        D = w*inv(Diagonal(a))
        M = I - D*a
        append!(Ms, [M])
        append!(Ds, [D])
    end

    M_mr = MultiResMatrix(Ms, A_mr.lvl_by_row_size, A_mr.lvl_by_column_size)
    D_mr = MultiResMatrix(Ds, A_mr.lvl_by_row_size, A_mr.lvl_by_column_size)

    return A_mr, M_mr, D_mr
end

function jacobi_smooth(n::Int64, M::MultiResMatrix, D::MultiResMatrix, x0::Vector{Float64}, r::Vector{Float64})

    x = x0[:]
    f = D * r
    for dummy=1:n
        x[:] = M * x + f
    end

    return x
end

function residual(A::MultiResMatrix, x::Vector{Float64}, b::Vector{Float64})
    return b - A * x
end

function vcycle(A::MultiResMatrix, x0::Vector{Float64}, b::Vector{Float64}, M::MultiResMatrix, D::MultiResMatrix, n::Int64, grd::Grid)

    # compute initial residual
    r = b - A*x0

    # relax on coarser grids
    for Restr in grd.restrictions
        x = jacobi_smooth(n, M, D, zeros(Float64, length(r)), r)
        r = Restr*(r - A*x)
    end

    # solve exactly on course grid, then interpolate
    x_crs = grd.interpolations[1] * (A.mats[end] \ r)
    r = grd.interpolations[1]*r - A * (x + x_crs)

    # relax on finer grids
    for Interp in grd.interpolations[2:end]
        x = jacobi_smooth(n, M, D, zeros(Float64, length(r)), r)
        r = Interp*(r - A*x)
    end

    # final smoothing
    x = jacobi_smooth(n, M, D, zeros(Float64, length(r)), r)
    return x0 + x
end

function pure_jacobi(A0::SparseMatrixCSC, x0::Vector{Float64}, b::Vector{Float64}, w::Float64, grd::Grid)

    # setup multiresolution matrices
    A, M, D = setup_jacobi_smoothing(A0, w, grd)

    # compute norm of b, initialize error
    bnorm = norm(b)
    err = 1.0
    x = x0[:]
    println()

    # interate until norm(residual) / bnorm < 10^-8
    n_iters = 0
    max_iters = 1000
    while (err > 1e-8) && (n_iters < max_iters)
        x[:] = jacobi_smooth(100, M, D, x, b) #XXX: This actually doesn't make sense, should just make jacobi_smooth depend on an initial guess.
        r = residual(A, x, b)
        err = norm(r) / bnorm
        n_iters += 1
        println(err)
    end

    return x
end

function multigrid_solve(A0::SparseMatrixCSC, x::Vector{Float64}, b::Vector{Float64}, w::Float64, grd::Grid)

    # setup multiresolution matrices
    A, M, D = setup_jacobi_smoothing(A0, w, grd)

    # compute norm of b, initialize error
    bnorm = norm(b)
    err = 1.0
    println()

    # interate until norm(residual) / bnorm < 10^-8
    n_iters = 0
    max_iters = 1000
    while (err > 1e-8) && (n_iters < max_iters)
        x = vcycle(A, x, b, M, D, 100, grd)
        err = norm(residual(A, x, b)) / bnorm
        n_iters += 1
        println(err)
    end

    return x
end


### A Pluto.jl notebook ###
# v0.17.2

using Markdown
using InteractiveUtils

# ╔═╡ c24044bc-2fc3-4e34-b99c-7dc4ec6a7515
import Pkg

# ╔═╡ 42da96ce-8122-4aa0-99d6-635f68db6e66
Pkg.activate(".")

# ╔═╡ 88e46392-1fbc-4079-8207-dd5b3ac01d7d
using DeltaRCWA

# ╔═╡ 16316725-c608-49a5-b90c-8caff973683a
using Plots

# ╔═╡ 83d9578c-3960-486c-ad3d-b4449f7bae6f
md"
# Fresnel scattering at a dielectric interface

Note that our GSTCs simplify to continuity of the tangential fields when
the surface conductivities vanish.
Therefore, our scattering matrix should recover the Fresnel coefficients
when we allow the medium on each side of the interface to differ.

The analytical solution for this problem is known and I will just cite it [here](https://homerreid.github.io/scuff-em-documentation/tests/FresnelScattering/FresnelScattering/)
``\newcommand{\bm}{\boldsymbol}``
"

# ╔═╡ 6cbc5f34-fe14-45e8-9fee-9787d76ad0fe
struct EmptySheet <: Sheet end

# ╔═╡ fb8df44c-a354-42de-9e99-df64dd4c5727
DeltaRCWA.ElectricResponseStyle(::Type{EmptySheet}) = Admittance()

# ╔═╡ f3d36c35-c09b-483f-bb41-046b6ecd0773
const Water = UniformMedium{1.7, 1.0}()

# ╔═╡ e899dbc0-03bb-4a13-9c09-bfad33e5c1d3
const stack = SheetStack(EmptySheet(), (Vacuum, Water))

# ╔═╡ 43d29086-510d-4870-94b7-2ff5ff34324b
n = 100 # number of grid points

# ╔═╡ 5377ae36-32fc-41c5-bd03-c737422b17f4
prob = DeltaRCWAProblem(stack, ((n, 1.0), ), 10.0, [i==2 ? 1 : 0 for i in 1:n], zeros(n))

# ╔═╡ 65c2e966-8cb0-4d76-96db-18f9045695a3
sol = solve(prob)

# ╔═╡ bf1213cd-2ec1-4d1e-8ebc-5be5af26621c
plot(sol; combine=true)

# ╔═╡ ba0c4060-f9d4-43b6-bc31-b48dcd81dd91
md"""
``k_z = \sqrt{\omega^2 \epsilon \mu - k_x^2 - k_y^2}``
"""

# ╔═╡ 687627f9-6938-4d7a-9395-9f1783b8a449
md"""
## Scattering matrix in 2d

We restrict ourselves to the decoupled TE/TM mode case (2d with no off-diagonal impedance matrix elements) for simplicity (we can make the 3d case work too).
Starting from the GSTCs for the TM modes (when ``k_y=0`` and the TE mode zero, ``E_y=H_x=0``)
```math
\begin{align}
-(H_y^{(2,+)} + H_y^{(2,-)} - H_y^{(1,+)} - H_y^{(1,-)})
	&= \sigma^e_{xx} (E_x^{(2,+)} + E_x^{(2,-)} + E_x^{(1,+)} + E_x^{(1,-)})/2
\\
(E_x^{(2,+)} + E_x^{(2,-)} - E_x^{(1,+)} - E_x^{(1,-)})
	&= -\sigma^m_{yy} (H_y^{(2,+)} + H_y^{(2,-)} + H_y^{(1,+)} + H_y^{(1,-)})/2
\end{align}
```
and the corresponding planewave relations (indexed by medium ``i``)
```math
\begin{align}
% -\omega^2 \epsilon^i \mu^i H_x^i &= \omega \epsilon^i k_z^i E_y^i
% \\
(k_x^2 -\omega^2 \epsilon^i \mu^i) H_y^i =  -(k_z^i)^2 H_y^i &= -\omega \epsilon^i k_z^i E_x^i
\\
-\omega^2 \epsilon^i \mu^i E_x^i &= -\omega \mu^i k_z^i H_y^i
% \\
% (k_x^2 -\omega^2 \epsilon^i \mu^i) E_y^i &= \omega \mu^i k_z^i H_x^i
\end{align}
```
which both are the same and let us eliminate the electric field using
```math
\begin{align}
E_x^i &= \frac{k_z^i}{\omega \epsilon^i} H_y^i
\end{align}
```
Taking into account the direction of propagation means
```math
\begin{align}
E_x^{(i,\pm)} &= \pm\frac{k_z^i}{\omega \epsilon^i} H_y^{(i,\pm)}
\end{align}
```
Note that this is the point where the 2d case is advantageous because the decoupled modes make the inversion in this case just the inverse of a diagonal matrix.
Hence the GSTCs are
```math
\begin{align}
&- \left( H_y^{(2,+)} + H_y^{(2,-)} - H_y^{(1,+)} - H_y^{(1,-)} \right)
\\ &\qquad = \sigma^e_{xx} \left( \frac{k_z^2}{2\omega\epsilon^2} \left( H_y^{(2,+)} - H_y^{(2,-)} \right) + \frac{k_z^1}{2\omega\epsilon^1} \left( H_y^{(1,+)} - H_y^{(1,-)} \right) \right)
\\
&\frac{2k_z^2}{\omega\epsilon^2} \left( H_y^{(2,+)} - H_y^{(2,-)} \right) - \frac{2k_z^1}{\omega\epsilon^1} \left( H_y^{(1,+)} - H_y^{(1,-)} \right)
\\ &\qquad = -\sigma^m_{yy} \left( H_y^{(2,+)} + H_y^{(2,-)} + H_y^{(1,+)} + H_y^{(1,-)} \right)
\end{align}
```
which gives the following scattering matrix
```math
\begin{align}
&\begin{pmatrix}
I + \sigma^e_{xx} \frac{k_z^1}{2\omega\epsilon^1}
&
-I - \sigma^e_{xx} \frac{k_z^2}{2\omega\epsilon^2}
\\
\sigma^m_{yy} + \frac{2k_z^1}{\omega\epsilon^1}
&
\sigma^m_{yy} + \frac{2k_z^2}{\omega\epsilon^2}
\end{pmatrix}
\begin{pmatrix}
H_y^{(1,-)}
\\
H_y^{(2,+)}
\end{pmatrix}
\\ =
&\begin{pmatrix}
-I + \sigma^e_{xx} \frac{k_z^1}{2\omega\epsilon^1}
&
I - \sigma^e_{xx} \frac{k_z^2}{2\omega\epsilon^2}
\\
-\sigma^m_{yy} + \frac{2k_z^1}{\omega\epsilon^1}
&
-\sigma^m_{yy} + \frac{2k_z^2}{\omega\epsilon^2}
\end{pmatrix}
\begin{pmatrix}
H_y^{(1,+)}
\\
H_y^{(2,-)}
\end{pmatrix}
\end{align}
```
"""

# ╔═╡ 860dc600-c86d-43a2-a88a-eb876eecbd6e
md"""
## Recovering Fresnel scattering

We can recover Fresnel scattering by setting the metasurface conductivities to zero, inverting the left hand side and reading off the scattered amplitudes.
To illustrate this, let's calculate this in this special case, which is just requiring the continuity of the tangential fields
```math
\begin{pmatrix}
I & -I
\\
\epsilon^2 k_z^1 & \epsilon^1 k_z^2
\end{pmatrix}
\begin{pmatrix}
H_y^{(1,-)}
\\
H_y^{(2,+)}
\end{pmatrix}
=
\begin{pmatrix}
-I & I
\\
\epsilon^2 k_z^1
& \epsilon^1 k_z^2
\end{pmatrix}
\begin{pmatrix}
H_y^{(1,+)}
\\
H_y^{(2,-)}
\end{pmatrix}
```
Let's invert the left hand side
```math
\begin{align}
\begin{pmatrix}
H_y^{(1,-)}
\\
H_y^{(2,+)}
\end{pmatrix}
&=
\begin{pmatrix}
(1 + \frac{\epsilon^2 k_z^1}{\epsilon^1 k_z^2})^{-1}
&
(\epsilon^1 k_z^2 + \epsilon^1 k_z^2)^{-1}
\\
-(1 + \frac{\epsilon^1 k_z^2}{\epsilon^2 k_z^1})^{-1}
&
(\epsilon^1 k_z^2 + \epsilon^2 k_z^1)^{-1}
\end{pmatrix}
\begin{pmatrix}
-I & I
\\
\epsilon^2 k_z^1 & \epsilon^1 k_z^2
\end{pmatrix}
\begin{pmatrix}
H_y^{(1,+)}
\\
H_y^{(2,-)}
\end{pmatrix}
\\ &=
(\epsilon^1 k_z^2 + \epsilon^2 k_z^1)^{-1}
\begin{pmatrix}
\epsilon^2 k_z^1 - \epsilon^1 k_z^2 & 2 \epsilon^1 k_z^2
\\
2 \epsilon^2 k_z^1 & \epsilon^1 k_z^2 - \epsilon^2 k_z^1
\end{pmatrix}
\begin{pmatrix}
H_y^{(1,+)}
\\
H_y^{(2,-)}
\end{pmatrix}
\end{align}
```
Note how the scattering depends only on the angle of incidence from the normal axis (the ratio of ``\omega`` and ``k_z``) and the impedances of the materials.
An analogous calculation can be done for the TE mode
"""

# ╔═╡ 271422af-db75-4373-be1c-104745f941bf
md"""
## Scattering matrix in 3d

The planewave form of Ampère's law we will use within dielectric medium ``i`` with permittivity ``\epsilon^i`` and permeability ``\mu^i`` is
```math
\bm{R} \bm{\Omega}^i \bm{R}^{-1} \bm{E}_\parallel^i = \omega \bm{R} \mu^i k_z^i \bm{H}_\parallel^i
```
where
```math
\bm{\Omega}^i = \bm{K} - \omega^2 \epsilon^i \mu^i \bm{I}
, \qquad
\bm{K} =
\begin{pmatrix}
k_x k_x & k_x k_y
\\
k_y k_x & k_y k_y
\end{pmatrix}
, \qquad
\bm{R} =
\begin{pmatrix}
0 & -1
\\
1 & 0
\end{pmatrix}
.
```
The GSTCs (with an ``\bm{R}`` and a 2 moved about) are
```math
\begin{align}
2 \bm{\rho}^e \bm{R} (\bm{H}_\parallel^2 - \bm{H}_\parallel^1)
&= \bm{E}_\parallel^2 + \bm{E}_\parallel^1
\\
\bm{E}_\parallel^2 - \bm{E}_\parallel^1
&= \bm{R} \bm{\sigma}^m (\bm{H}_\parallel^2 + \bm{H}_\parallel^1)/2
\end{align}
```
It is basic algebra to show that ``\bm{\Omega}^i`` commutes with ``\bm{\Omega}^j`` so multiplying the GSTCs by ``\bm{\Omega}^1 \bm{\Omega}^2 \bm{R}^{-1} = \bm{\Omega}^2 \bm{\Omega}^1 \bm{R}^{-1}`` gives
```math
\begin{align}
2 \bm{\Omega}^1 \bm{\Omega}^2 \bm{R}^{-1}
\bm{\rho}^e \bm{R} (\bm{H}_\parallel^2 - \bm{H}_\parallel^1)
&= \bm{\Omega}^1 \bm{R}^{-1} \bm{R} \bm{\Omega}^2 \bm{R}^{-1} \bm{E}_\parallel^2
+ \bm{\Omega}^2 \bm{R}^{-1} \bm{R} \bm{\Omega}^1 \bm{R}^{-1} \bm{E}_\parallel^1
\\
\bm{\Omega}^1 \bm{R}^{-1} \bm{R} \bm{\Omega}^2 \bm{R}^{-1} \bm{E}_\parallel^2
- \bm{\Omega}^2 \bm{R}^{-1} \bm{R} \bm{\Omega}^1 \bm{R}^{-1} \bm{E}_\parallel^1
&= \bm{\Omega}^1 \bm{\Omega}^2 \bm{\sigma}^m (\bm{H}_\parallel^2 + \bm{H}_\parallel^1)/2
\end{align}
```
Next we will do 3 operations in the same step:
1. Break each mode into forward and backward propagating components: ``\bm{H}_\parallel^i \to \bm{H}_\parallel^{(i,+)} + \bm{H}_\parallel^{(i,-)},\ \bm{E}_\parallel^i \to \bm{E}_\parallel^{(i,+)} + \bm{E}_\parallel^{(i,-)}``
2. Eliminate the electric field with Ampère's law: ``\bm{R} \bm{\Omega}^i \bm{R}^{-1} \bm{E}_\parallel^{(i,\pm)} = \pm \omega \bm{R} \mu^i k_z^i \bm{H}_\parallel^{(i,\pm)}``
3. Cancel any pairs of ``\bm{R}`` and ``\bm{R}^{-1}``
```math
\begin{align}
&2 \bm{\Omega}^1 \bm{\Omega}^2 \bm{R}^{-1} \bm{\rho}^e \bm{R} \left( \bm{H}_\parallel^{(2,+)} + \bm{H}_\parallel^{(2,-)} - \bm{H}_\parallel^{(1,+)} - \bm{H}_\parallel^{(1,-)} \right)
\\ &\qquad = \omega \left(
\bm{\Omega}^1 \mu^2 k_z^2 \left( \bm{H}_\parallel^{(2,+)} - \bm{H}_\parallel^{(2,-)} \right)
+ \bm{\Omega}^2 \mu^1 k_z^1 \left( \bm{H}_\parallel^{(1,+)} - \bm{H}_\parallel^{(1,-)} \right) \right)
\\
&\omega \left(
\bm{\Omega}^1 \mu^2 k_z^2 \left( \bm{H}_\parallel^{(2,+)} - \bm{H}_\parallel^{(2,-)} \right)
- \bm{\Omega}^2 \mu^1 k_z^1 \left( \bm{H}_\parallel^{(1,+)} - \bm{H}_\parallel^{(1,-)} \right) \right)
\\ &\qquad = \bm{\Omega}^1 \bm{\Omega}^2 \bm{\sigma}^m \left( \bm{H}_\parallel^{(2,+)} + \bm{H}_\parallel^{(2,-)} + \bm{H}_\parallel^{(1,+)} + \bm{H}_\parallel^{(1,-)} \right)/2
\end{align}
```
Next if we move the scattered components, (1,-) and (2,+), to the left and the incident components to the right, (1,+) and (2,-), we can express the scattering matrix as
```math
\begin{align}
&\begin{pmatrix}
-\bm{\Omega}^1 \bm{\Omega}^2 \bm{R}^{-1} \bm{\rho}^e \bm{R} + \frac{1}{2} \omega \bm{\Omega}^2 \mu^1 k_z^1
&
\bm{\Omega}^1 \bm{\Omega}^2 \bm{R}^{-1} \bm{\rho}^e \bm{R} - \frac{1}{2} \omega \bm{\Omega}^1 \mu^2 k_z^2
\\
-\bm{\Omega}^1 \bm{\Omega}^2 \bm{\sigma}^m + 2 \omega \bm{\Omega}^2 \mu^1 k_z^1
&
-\bm{\Omega}^1 \bm{\Omega}^2 \bm{\sigma}^m + 2 \omega \bm{\Omega}^1 \mu^2 k_z^2
\end{pmatrix}
\begin{pmatrix}
\bm{H}_\parallel^{(1,-)}
\\
\bm{H}_\parallel^{(2,+)}
\end{pmatrix}
\\ =
&\begin{pmatrix}
\bm{\Omega}^1 \bm{\Omega}^2 \bm{R}^{-1} \bm{\rho}^e \bm{R} + \frac{1}{2} \omega \bm{\Omega}^2 \mu^1 k_z^1
&
-\bm{\Omega}^1 \bm{\Omega}^2 \bm{R}^{-1} \bm{\rho}^e \bm{R} - \frac{1}{2} \omega \bm{\Omega}^1 \mu^2 k_z^2
\\
\bm{\Omega}^1 \bm{\Omega}^2 \bm{\sigma}^m + 2 \omega \bm{\Omega}^2 \mu^1 k_z^1
&
\bm{\Omega}^1 \bm{\Omega}^2 \bm{\sigma}^m + 2 \omega \bm{\Omega}^1 \mu^2 k_z^2
\end{pmatrix}
\begin{pmatrix}
\bm{H}_\parallel^{(1,+)}
\\
\bm{H}_\parallel^{(2,-)}
\end{pmatrix}
\end{align}
```
Note that if the two media are the same (i.e. ``\epsilon^1=\epsilon^2, \mu^1=\mu^2``) then ``\bm{\Omega}^1 = \bm{\Omega}^2`` and the scattering matrix has an overall factor of ``\Omega`` that can be removed, which recovers the result for the scattering matrix in a homogenous medium.
In fact we can manipulate this expression further using the specific form of the inverse
```math
(\bm{\Omega}^i)^{-1} = \bm{R}^{-1} \bm{\Omega}^i \bm{R} \frac{1}{\omega^2 \epsilon^i \mu^i (k_z^i)^2}
```
due to the fact ``\bm{\Omega}^i`` is a square block matrix whose blocks are diagonal matrices.
"""

# ╔═╡ 325b8359-ebac-434b-bba1-31aba05c2c39
md"
## Higher order corrections for a metasurface at a dielectric interface

The GSTCs we use are published by Epstein (DOI: 10.1109/TAP.2014.2354419)

These are based on ones published by Kuester (DOI: 10.1109/TAP.2003.817560), which
make the following assumptions.
'In this paper, we will derive what we believe to be the most flexible form of a
generalized sheet transition condition (GSTC) applicable to metafilms located in a
homogeneous medium. The case of a metafilm on a dielectric interface will be
considered in a separate paper.'

Kuester published a follow-up article (DOI: 10.1109/APS.2010.5562250)
with a first order correction to the previous one which should generalize to the
case where the metafilm has different media on either side.
The form of the GSTC is the same as before, except the impedance matrices are modified functions of the actual material parameters
"

# ╔═╡ Cell order:
# ╟─83d9578c-3960-486c-ad3d-b4449f7bae6f
# ╠═c24044bc-2fc3-4e34-b99c-7dc4ec6a7515
# ╠═42da96ce-8122-4aa0-99d6-635f68db6e66
# ╠═88e46392-1fbc-4079-8207-dd5b3ac01d7d
# ╠═6cbc5f34-fe14-45e8-9fee-9787d76ad0fe
# ╠═fb8df44c-a354-42de-9e99-df64dd4c5727
# ╠═f3d36c35-c09b-483f-bb41-046b6ecd0773
# ╠═e899dbc0-03bb-4a13-9c09-bfad33e5c1d3
# ╠═43d29086-510d-4870-94b7-2ff5ff34324b
# ╠═5377ae36-32fc-41c5-bd03-c737422b17f4
# ╠═65c2e966-8cb0-4d76-96db-18f9045695a3
# ╠═16316725-c608-49a5-b90c-8caff973683a
# ╠═bf1213cd-2ec1-4d1e-8ebc-5be5af26621c
# ╟─ba0c4060-f9d4-43b6-bc31-b48dcd81dd91
# ╟─687627f9-6938-4d7a-9395-9f1783b8a449
# ╟─860dc600-c86d-43a2-a88a-eb876eecbd6e
# ╟─271422af-db75-4373-be1c-104745f941bf
# ╟─325b8359-ebac-434b-bba1-31aba05c2c39


module DoubleExponentialFormulas

using LinearAlgebra: norm

export
    QuadTS,
    QuadES,
    QuadSS,
    QuadDE,
    quadde,
    quaddeo

"""
    startindex(f::Function, weights, istart::Integer)

Returns the first index i which `f(weights[i][1])` is not NaN,
where `i >= istart`.

This function is employed to avoid sampling `f(x)` with too large `abs(x)` in
trapezoidal rule. For example, the limit of `x*exp(-x^2)` may result in NaN
when `x` approaches infinity since the expression reduced to be `Inf*0` in
float point number computations.
"""
function startindex(f::Function, weights, istart::Integer)
    iend = length(weights)
    for i in istart:iend
        x, _ = @inbounds weights[i]
        y = try
            f(x)
        catch
            continue
        end
        all(.!isnan.(y)) && return i
    end
    return iend
end

include("mapsum.jl")

"""
    estimate_error(T::Type{<:AbstractFloat}, prevI, I)

Estimate an error from the true integral(, though it may not be very accurate).

With double exponential formula, the error from the true value `I` decays
exponentially.
    ΔI(h) = I - I(h) ≈ exp(-C/h)
`I(h)` is the numerical integral with a step size `h`, `C` is a
constant depending on the integrand. If the step size is halved,
    ΔI(h/2) = I - I(h/2) ≈ exp(-2C/h) ≈ ( ΔI(h) )²,
the significant digits gets almost twice. Therefore, `I(h/2)` should be
considerably closer to the true value than `I(h)`.
    ΔI(h/2) ≈ ( ΔI(h) )² ≈ ( I - I(h) )² ≈ ( I(h/2) - I(h) )²
Introduce a magnification coefficient `M` for safety.
    ΔI(h/2) ≈ M*( I(h/2) - I(h) )²
Furthermore, the floating point numbers generally have a finite
significant digits; `ΔI(h/2)` would asymptotically get closer to `I*ε`
at minimum, rather than 0.
    ΔI(h/2) ≈ M*( I(h/2) - I(h) )² + I*ε
            ≈ M*( I(h/2) - I(h) )² + I(h/2)*ε

FIXME: The logical ground of the safety factor is very weak. Any better way?
"""
function estimate_error(T::Type{<:AbstractFloat}, prevI, I)
    ε = eps(T)
    M = 20
    return M*norm(I - prevI)^2 + norm(I)*ε
end

include("quadts.jl")
include("quades.jl")
include("quadss.jl")
include("quadde.jl")

include("quaddeo.jl")

end # module


"""
    idx_for(i::Int)

Simple helper function for converting 0-based indices to Julia's 1-based indices.
"""
function idx_for(i::Int)
    return i + 1
end

"""
    heap_sort!(arr::Vector{T}, gt = >, N::Int = length(arr)) where {T}

Sort the given vector (in-place) using the Heapsort algorithm.

Heapsort consists of two stages:
  1. Building a (max) heap of the array
  2. Repeatedly extracting the largest element and inserting it at the front of the sorted part of the array

After the largest element has been extracted, the tree is updated to maintain the heap property via a "sifting" operation.

  Storing a heap in an array is pretty straightforward - for every node with index n, its children are stored at indices
  2n + 1 and 2n + 2 (for 0-based indices). Index 0 contains the root node.
  Since Julia's indices are 1-based, we need to change this a little bit. We're using a trivial helper function
  idx_for to convert from 0-based to 1-based.

  See https://en.wikipedia.org/wiki/Heapsort for a complete explanation of Heapsort.

  Contributed By:- [Frank Schmitt](https://github.com/frankschmitt)
"""
function heap_sort!(arr::Vector{T}, gt = >, N::Int = length(arr)) where {T}
    n = N
    i = div(n, 2)
    t = -1

    while true
        if i > 0    # First stage - Sorting the heap
            i -= 1    # Save its index to i
            t = arr[idx_for(i)]    # Save parent value to t
        else          # Second stage - Extracting elements in-place
            n -= 1    # Make the new heap smaller
            if (n == 0)
                return # When the heap is empty, we are done
            end
            t = arr[idx_for(n)]               # Save last value (it will be overwritten)
            arr[idx_for(n)] = arr[idx_for(0)] # Save largest value at the end of arr
        end

        parent = i # We will start pushing down t from parent
        child = i * 2 + 1 # parent's left child

        # Sift operation - pushing the value of t down the heap
        while child < n
            if (
                child + 1 < n &&
                gt(arr[idx_for(child + 1)], arr[idx_for(child)])
            )
                child += 1 # Choose the largest child
            end
            if (gt(arr[idx_for(child)], t))                   # If any child is bigger than the parent
                arr[idx_for(parent)] = arr[idx_for(child)]    # Move the largest child up
                parent = child                                # Move parent pointer to this child
                child = parent * 2 + 1                            #  find the next child
            else
                break # t's place is found
            end
        end
        arr[idx_for(parent)] = t
    end
end