# Affinity propagation
#
#   Reference:
#       Clustering by Passing Messages Between Data Points.
#       Brendan J. Frey and Delbert Dueck
#       Science, vol 315, pages 972-976, 2007.
#

#### Interface

mutable struct AffinityPropResult <: ClusteringResult
    exemplars::Vector{Int}      # indexes of exemplars (centers)
    assignments::Vector{Int}    # assignments for each point
    counts::Vector{Int}         # number of samples in each cluster
    iterations::Int             # number of iterations executed
    converged::Bool             # converged or not
end

const _afp_default_maxiter = 200
const _afp_default_damp = 0.5
const _afp_default_tol = 1.0e-6
const _afp_default_display = :none

function affinityprop(S::AbstractArray{T};
                      maxiter::Integer=_afp_default_maxiter,
                      tol::Real=_afp_default_tol,
                      damp::Real=_afp_default_damp,
                      display::Symbol=_afp_default_display) where T<:AbstractFloat

    # check arguments
    n = size(S, 1)
    size(S, 2) == n || error("S must be a square matrix.")
    n >= 2 || error("the number of samples must be at least 2.")
    tol > 0 || error("tol must be a positive value.")
    0 <= damp < 1 || error("damp must be a non-negative real value below 1.")

    # invoke core implementation
    _affinityprop(S, round(Int, maxiter), tol, convert(T, damp), display_level(display))
end

#### Implementation

function _affinityprop(S::AbstractArray{T},
                       maxiter::Int,
                       tol::Real,
                       damp::T,
                       displevel::Int) where T<:AbstractFloat
    n = size(S, 1)
    n2 = n * n

    # initialize messages
    R = zeros(T, n, n)  # responsibilities
    A = zeros(T, n, n)  # availabilities

    # prepare storages
    Rt = Matrix{T}(undef, n, n)
    At = Matrix{T}(undef, n, n)

    if displevel >= 2
        @printf "%7s %12s | %8s \n" "Iters" "objv-change" "exemplars"
        println("-----------------------------------------------------")
    end

    t = 0
    converged = false
    while !converged && t < maxiter
        t += 1

        # compute new messages
        _afp_compute_r!(Rt, S, A)
        _afp_dampen_update!(R, Rt, damp)

        _afp_compute_a!(At, R)
        _afp_dampen_update!(A, At, damp)

        # determine convergence
        ch = max(Linfdist(A, At), Linfdist(R, Rt)) / (one(T) - damp)
        converged = (ch < tol)

        if displevel >= 2
            # count the number of exemplars
            ne = _afp_count_exemplars(A, R)
            @printf("%7d %12.4e | %8d\n", t, ch, ne)
        end
    end

    # extract exemplars and assignments
    exemplars = _afp_extract_exemplars(A, R)
    assignments, counts = _afp_get_assignments(S, exemplars)

    if displevel >= 1
        if converged
            println("Affinity propagation converged with $t iterations: $(length(exemplars)) exemplars.")
        else
            println("Affinity propagation terminated without convergence after $t iterations: $(length(exemplars)) exemplars.")
        end
    end

    # produce output struct
    return AffinityPropResult(exemplars, assignments, counts, t, converged)
end

# compute responsibilities
function _afp_compute_r!(R::Matrix{T}, S::AbstractArray{T}, A::Matrix{T}) where T
    n = size(S, 1)

    I1 = Vector{Int}(undef, n)  # I1[i] is the column index of the maximum element in (A+S)[i,:]
    Y1 = Vector{T}(undef, n)    # Y1[i] is the maximum element in (A+S)[i,:]
    Y2 = Vector{T}(undef, n)    # Y2[i] is the second maximum element in (A+S)[i,:]

    # Find the first and second maximum elements along each row
    @inbounds for i = 1:n
        v1 = A[i,1] + S[i,1]
        v2 = A[i,2] + S[i,2]
        if v1 > v2
            I1[i] = 1
            Y1[i] = v1
            Y2[i] = v2
        else
            I1[i] = 2
            Y1[i] = v2
            Y2[i] = v1
        end
    end
    @inbounds for j = 3:n, i = 1:n
        v = A[i,j] + S[i,j]
        if v > Y2[i]
            if v > Y1[i]
                Y2[i] = Y1[i]
                I1[i] = j
                Y1[i] = v
            else
                Y2[i] = v
            end
        end
    end

    # compute R values
    @inbounds for j = 1:n, i = 1:n
        mv = (j == I1[i] ? Y2[i] : Y1[i])
        R[i,j] = S[i,j] - mv
    end

    return R
end

# compute availabilities
function _afp_compute_a!(A::Matrix{T}, R::Matrix{T}) where T
    n = size(R, 1)
    z = zero(T)
    for j = 1:n
        @inbounds rjj = R[j,j]

        # compute s <- sum_{i \ne j} max(0, R(i,j))
        s = z
        for i = 1:n
            if i != j
                @inbounds r = R[i,j]
                if r > z
                    s += r
                end
            end
        end

        for i = 1:n
            if i == j
                @inbounds A[i,j] = s
            else
                @inbounds r = R[i,j]
                u = rjj + s
                if r > z
                    u -= r
                end
                A[i,j] = ifelse(u < z, u, z)
            end
        end
    end
    return A
end

# dampen update
function _afp_dampen_update!(x::Array{T}, xt::Array{T}, damp::T) where T
    ct = one(T) - damp
    for i = 1:length(x)
        @inbounds x[i] = ct * xt[i] + damp * x[i]
    end
    return x
end

# count the number of exemplars
function _afp_count_exemplars(A::Matrix, R::Matrix)
    n = size(A,1)
    c = 0
    for i = 1:n
        @inbounds if A[i,i] + R[i,i] > 0
            c += 1
        end
    end
    return c
end

# extract all exemplars
function _afp_extract_exemplars(A::Matrix, R::Matrix)
    n = size(A,1)
    r = Int[]
    for i = 1:n
        @inbounds if A[i,i] + R[i,i] > 0
            push!(r, i)
        end
    end
    return r
end

# get assignments
function _afp_get_assignments(S::AbstractArray, exemplars::Vector{Int})
    n = size(S, 1)
    k = length(exemplars)
    Se = S[:, exemplars]
    a = Vector{Int}(undef, n)
    cnts = zeros(Int, k)
    for i = 1:n
        p = 1
        v = Se[i,1]
        for j = 2:k
            s = Se[i,j]
            if s > v
                v = s
                p = j
            end
        end
        a[i] = p
    end
    for i = 1:k
        a[exemplars[i]] = i
    end
    for i = 1:n
        @inbounds cnts[a[i]] += 1
    end
    return (a, cnts)
end


"""
    Wishart(ν, S)
```julia
ν::Real           degrees of freedom (whole number or a real number greater than p - 1)
S::AbstractPDMat  p x p scale matrix
```
The [Wishart distribution](http://en.wikipedia.org/wiki/Wishart_distribution)
generalizes the gamma distribution to ``p\\times p`` real, positive semidefinite
matrices ``\\mathbf{H}``.

If ``\\nu>p-1``, then ``\\mathbf{H}\\sim \\textrm{W}_p(\\nu, \\mathbf{S})``
has rank ``p`` and its probability density function is

```math
f(\\mathbf{H};\\nu,\\mathbf{S}) = \\frac{1}{2^{\\nu p/2} \\left|\\mathbf{S}\\right|^{\\nu/2} \\Gamma_p\\left(\\frac {\\nu}{2}\\right ) }{\\left|\\mathbf{H}\\right|}^{(\\nu-p-1)/2} e^{-(1/2)\\operatorname{tr}(\\mathbf{S}^{-1}\\mathbf{H})}.
```

If ``\\nu\\leq p-1``, then ``\\mathbf{H}`` is rank ``\\nu`` and it has
a density with respect to a suitably chosen volume element on the space of
positive semidefinite matrices. See [here](https://doi.org/10.1214/aos/1176325375).

For integer ``\\nu``, a random matrix given by

```math
\\mathbf{H} = \\mathbf{X}\\mathbf{X}^{\\rm{T}},
\\quad\\mathbf{X} \\sim \\textrm{MN}_{p,\\nu}(\\mathbf{0}, \\mathbf{S}, \\mathbf{I}_{\\nu})
```

has ``\\mathbf{H}\\sim \\textrm{W}_p(\\nu, \\mathbf{S})``.
For non-integer ``\\nu``, Wishart matrices can be generated via the
[Bartlett decomposition](https://en.wikipedia.org/wiki/Wishart_distribution#Bartlett_decomposition).
"""
struct Wishart{T<:Real, ST<:AbstractPDMat, R<:Integer} <: ContinuousMatrixDistribution
    df::T          # degree of freedom
    S::ST          # the scale matrix
    logc0::T       # the logarithm of normalizing constant in pdf
    rank::R        # rank of a sample
    singular::Bool # singular of nonsingular wishart?
end

#  -----------------------------------------------------------------------------
#  Constructors
#  -----------------------------------------------------------------------------

function Wishart(df::T, S::AbstractPDMat{T}) where T<:Real
    df > 0 || throw(ArgumentError("df must be positive. got $(df)."))
    p = dim(S)
    singular = df <= p - 1
    if singular
        isinteger(df) || throw(
            ArgumentError("df of a singular Wishart distribution must be an integer (got $df)")
        )
    end
    rnk::Integer = ifelse(singular, df, p)
    logc0 = wishart_logc0(df, S, rnk)
    _df, _logc0 = promote(df, logc0)
    Wishart{typeof(_df), typeof(S), typeof(rnk)}(_df, S, _logc0, rnk, singular)
end

function Wishart(df::Real, S::AbstractPDMat)
    T = Base.promote_eltype(df, S)
    Wishart(T(df), convert(AbstractArray{T}, S))
end

Wishart(df::Real, S::Matrix) = Wishart(df, PDMat(S))
Wishart(df::Real, S::Cholesky) = Wishart(df, PDMat(S))

#  -----------------------------------------------------------------------------
#  REPL display
#  -----------------------------------------------------------------------------

show(io::IO, d::Wishart) = show_multline(io, d, [(:df, d.df), (:S, d.S)])

#  -----------------------------------------------------------------------------
#  Conversion
#  -----------------------------------------------------------------------------

function convert(::Type{Wishart{T}}, d::Wishart) where T<:Real
    P = convert(AbstractArray{T}, d.S)
    Wishart{T, typeof(P), typeof(d.rank)}(T(d.df), P, T(d.logc0), d.rank, d.singular)
end
function convert(::Type{Wishart{T}}, df, S::AbstractPDMat, logc0, rnk, singular) where T<:Real
    P = convert(AbstractArray{T}, S)
    Wishart{T, typeof(P), typeof(rnk)}(T(df), P, T(logc0), rnk, singular)
end

#  -----------------------------------------------------------------------------
#  Properties
#  -----------------------------------------------------------------------------

insupport(::Type{Wishart}, X::AbstractMatrix) = ispossemdef(X)
function insupport(d::Wishart, X::AbstractMatrix)
    size(X) == size(d) || return false
    if d.singular
        return ispossemdef(X, rank(d))
    else
        return isposdef(X)
    end
end

dim(d::Wishart) = dim(d.S)
size(d::Wishart) = (p = dim(d); (p, p))
rank(d::Wishart) = d.rank
params(d::Wishart) = (d.df, d.S)
@inline partype(d::Wishart{T}) where {T<:Real} = T

mean(d::Wishart) = d.df * Matrix(d.S)

function mode(d::Wishart)
    r = d.df - dim(d) - 1
    r > 0 || throw(ArgumentError("mode is only defined when df > p + 1"))
    return Matrix(d.S) * r
end

function meanlogdet(d::Wishart)
    logdet_S = logdet(d.S)
    p = dim(d)
    v = logdet_S + p * oftype(logdet_S, logtwo)
    df = oftype(logdet_S, d.df)
    for i in 0:(p - 1)
        v += digamma((df - i) / 2)
    end
    return d.singular ? oftype(v, -Inf) : v
end

function entropy(d::Wishart)
    d.singular && throw(ArgumentError("entropy not defined for singular Wishart."))
    p = dim(d)
    df = d.df
    return -d.logc0 - ((df - p - 1) * meanlogdet(d) - df * p) / 2
end

#  Gupta/Nagar (1999) Theorem 3.3.15.i
function cov(d::Wishart, i::Integer, j::Integer, k::Integer, l::Integer)
    S = d.S
    return d.df * (S[i, k] * S[j, l] + S[i, l] * S[j, k])
end

function var(d::Wishart, i::Integer, j::Integer)
    S = d.S
    return d.df * (S[i, i] * S[j, j] + S[i, j] ^ 2)
end

#  -----------------------------------------------------------------------------
#  Evaluation
#  -----------------------------------------------------------------------------

function wishart_logc0(df::T, S::AbstractPDMat{T}, rnk::Integer) where {T<:Real}
    p = dim(S)
    if df <= p - 1
        return singular_wishart_logc0(p, df, S, rnk)
    else
        return nonsingular_wishart_logc0(p, df, S)
    end
end

function logkernel(d::Wishart, X::AbstractMatrix)
    if d.singular
        return singular_wishart_logkernel(d, X)
    else
        return nonsingular_wishart_logkernel(d, X)
    end
end

#  Singular Wishart pdf: Theorem 6 in Uhlig (1994 AoS)
function singular_wishart_logc0(p::Integer, df::T, S::AbstractPDMat{T}, rnk::Integer) where {T<:Real}
    logdet_S = logdet(S)
    h_df = oftype(logdet_S, df) / 2
    return -h_df * (logdet_S + p * oftype(logdet_S, logtwo)) - logmvgamma(rnk, h_df) + ((rnk * (rnk - p)) // 2) * oftype(logdet_S, logπ)
end

function singular_wishart_logkernel(d::Wishart, X::AbstractMatrix)
    p = dim(d)
    r = rank(d)
    L = eigvals(Hermitian(X), (p - r + 1):p)
    return ((d.df - (p + 1)) * sum(log, L) - tr(d.S \ X)) / 2
end

#  Nonsingular Wishart pdf
function nonsingular_wishart_logc0(p::Integer, df::T, S::AbstractPDMat{T}) where {T<:Real}
    logdet_S = logdet(S)
    h_df = oftype(logdet_S, df) / 2
    return -h_df * (logdet_S + p * oftype(logdet_S, logtwo)) - logmvgamma(p, h_df)
end

function nonsingular_wishart_logkernel(d::Wishart, X::AbstractMatrix)
    return ((d.df - (dim(d) + 1)) * logdet(cholesky(X)) - tr(d.S \ X)) / 2
end

#  -----------------------------------------------------------------------------
#  Sampling
#  -----------------------------------------------------------------------------

function _rand!(rng::AbstractRNG, d::Wishart, A::AbstractMatrix)
    if d.singular
        axes2 = axes(A, 2)
        r = rank(d)
        randn!(rng, view(A, :, axes2[1:r]))
        fill!(view(A, :, axes2[(r + 1):end]), zero(eltype(A)))
    else
        _wishart_genA!(rng, A, d.df)
    end
    unwhiten!(d.S, A)
    A .= A * A'
end

function _wishart_genA!(rng::AbstractRNG, A::AbstractMatrix, df::Real)
    # Generate the matrix A in the Bartlett decomposition
    #
    #   A is a lower triangular matrix, with
    #
    #       A(i, j) ~ sqrt of Chisq(df - i + 1) when i == j
    #               ~ Normal()                  when i > j
    #
    T = eltype(A)
    z = zero(T)
    axes1 = axes(A, 1)
    @inbounds for (j, jdx) in enumerate(axes(A, 2)), (i, idx) in enumerate(axes1)
        A[idx, jdx] = if i < j
            z
        elseif i > j
            randn(rng, T)
        else
            rand(rng, Chi(df - i + 1))
        end
    end
    return A
end

#  -----------------------------------------------------------------------------
#  Test utils
#  -----------------------------------------------------------------------------

function _univariate(d::Wishart)
    check_univariate(d)
    df, S = params(d)
    α = df / 2
    β = 2 * first(S)
    return Gamma(α, β)
end

function _rand_params(::Type{Wishart}, elty, n::Int, p::Int)
    n == p || throw(ArgumentError("dims must be equal for Wishart"))
    ν = elty(n - 1 + abs(10 * randn()))
    X = rand(elty, n, n)
    X .= 2 .* X .- 1
    S = X * X'
    return ν, S
end


#! basic implementation

"""
basic cobb douglas
"""
function cobbdouglas(x, w)
    z = one(eltype(x))
    for i in eachindex(x)
        z *= x[i]^w[i]
    end
    return z
end

"""
basic ces
"""
function _ces(x, w, r)
    z = zero(eltype(x))
    for i in eachindex(x)
        z += w[i] * x[i]^r
    end
    return z^inv(r)
end

"""
constant elasticity of substitution function
"""
function ces(x, w, r)
    n = length(x)
    (n == length(w)) || throw(DimensionMismatch())
    r == zero(eltype(x)) ? cobbdouglas(x, w) : _ces(x, w, r)
end


########################################################################################
# Integrated Brownian Motion
########################################################################################
"""Generate the discrete dynamics for a q-IBM model. INCLUDES AUTOMATIC PRECONDITIONING!

Careful: Dimensions are ordered differently than in `probnum`!"""
function ibm(d::Integer, q::Integer)
    F̃ = diagm(1 => ones(q))
    I_d = diagm(0 => ones(d))
    F = kron(F̃, I_d)  # In probnum the order is inverted

    @fastmath function A!(A::AbstractMatrix, h::Real)
        # Assumes that A comes from a previous computation => zeros and one-diag
        val = one(h)
        for i in 1:q
            val = val / i
            for j in 1:d*(q+1-i)
                @inbounds A[j,j+(d*i)] = val
            end
        end
    end

    @fastmath function _transdiff_ibm_element(row::Int, col::Int, h::Real)
        idx = 2 * q + 1 - row - col
        fact_rw = factorial(q - row)
        fact_cl = factorial(q - col)
        return h / (idx * fact_rw * fact_cl)
    end
    @fastmath function Q!(Q::AbstractMatrix, h::Real, σ²::Real=1.0)
        val = one(h)
        @simd for col in 0:q
            @simd for row in col:q
                val = _transdiff_ibm_element(row, col, h) * σ²
                @simd for i in 0:d-1
                    @inbounds Q[1 + col*d + i,1 + row*d + i] = val
                    @inbounds Q[1 + row*d + i,1 + col*d + i] = val
                end
            end
        end
    end

    return A!, Q!
end

"""Same as above, but without the automatic preconditioning"""
function vanilla_ibm(d::Integer, q::Integer)
    F̃ = diagm(1 => ones(q))
    I_d = diagm(0 => ones(d))
    F = kron(F̃, I_d)  # In probnum the order is inverted

    @fastmath function A!(A::AbstractMatrix, h::Real)
        # Assumes that A comes from a previous computation => zeros and one-diag
        val = one(h)
        for i in 1:q
            val = val * h / i
            for j in 1:d*(q+1-i)
                @inbounds A[j,j+(d*i)] = val
            end
        end
    end

    @fastmath function _transdiff_ibm_element(row::Int, col::Int, h::Real)
        idx = 2 * q + 1 - row - col
        fact_rw = factorial(q - row)
        fact_cl = factorial(q - col)
        return h^idx / (idx * fact_rw * fact_cl)
    end
    @fastmath function Q!(Q::AbstractMatrix, h::Real, σ²::Real=1.0)
        val = one(h)
        @simd for col in 0:q
            @simd for row in col:q
                val = _transdiff_ibm_element(row, col, h) * σ²
                @simd for i in 0:d-1
                    @inbounds Q[1 + col*d + i,1 + row*d + i] = val
                    @inbounds Q[1 + row*d + i,1 + col*d + i] = val
                end
            end
        end
    end

    return A!, Q!
end


immutable Dual <: Real
    val::Float64
    adj::Float64
end

Base.show(io::IO, x::Dual) = print(io, "$(x.val) + $(x.adj)ϵ")

Base.convert(::Type{Dual}, x::Real) = Dual(x, 0)  # first argument is not named, but is if type Type{Dual}
Base.convert(::Type{Dual}, x::Dual) = x  # corner case, since Dual <: Real

Base.promote_rule(::Type{Float64}, ::Type{Dual}) = Dual
Base.promote_rule(::Type{Int64}, ::Type{Dual}) = Dual
Base.promote_rule(::Type{Float32}, ::Type{Dual}) = Dual

import Base: +, -, /, *, abs, <, ==, <=
+(x::Dual, y::Dual) = Dual(x.val + y.val, x.adj + y.adj)
-(x::Dual) = Dual(-x.val, -x.adj)
-(x::Dual, y::Dual) = x + -y
*(x::Dual, y::Dual) = Dual(x.val * y.val, x.val * y.adj + x.adj + y.val)
/(x::Dual, y::Dual) = Dual(x.val/y.val, x.adj/y.val - (x.val * y.adj)/y.val^2)
abs(x::Dual) = x.val > 0 ? x : -x
==(x::Dual, y::Dual) = x.val == y.val && x.adj == y.adj
<(x::Dual, y::Dual) = x.val == y.val ? x.adj < y.adj : x.val < y.val
<=(x::Dual, y::Dual) = x < y || x == y

import Base.log
log(x::Dual) = Dual(log(x.val), x.adj/x.val)


immutable Polynomial <: AbstractInterpolation
    Polynomial() = Polynomial
end

function default_settings{D,T}(::Type{Polynomial},axes::NTuple{D,T})
    N = length(axes[1])
    return Irregular{1,N}, None
end

type PolynomialInterpolation{T,N}
  x::Vector{T}
  y::Vector{T}
  coeff::Vector{T}
end

function preprocess(::Type{Polynomial}, axes, f)
    x = convert(Array, axes[1])
    y = convert(Array, f)
    nx = length(x)
    coeff = ones(nx)

    for i=1:nx
        for j = 1:nx
            i == j && continue
            coeff[i] = coeff[i]*(x[i] - x[j])
        end
        coeff[i] = (y[i]/coeff[i])
    end
    x,y,coeff = promote(x,y,coeff)
    return PolynomialInterpolation{eltype(x),nx}(x,y,coeff)
end

function apply_bc(::Type{Polynomial},BC::Type{Nearest}, I::Symbol, x::Symbol)
    apply_nearest(I,1,x,1)
end

function apply_bc(::Type{Polynomial},BC::Type{Reflect}, I::Symbol, x::Symbol)
    apply_reflect(I,1,x,1)
end

function apply_bc{V}(::Type{Polynomial},BC::Type{Value{V}}, I::Symbol, x::Symbol)
    apply_value(BC, I, 1, x, 1)
end

function apply_bc(::Type{Polynomial},BC::Type{None}, I::Symbol, x::Symbol)
    apply_none()
end

function apply_bc(::Type{Polynomial},BC::Type{Error}, I::Symbol, x::Symbol)
    apply_error(I, 1, x, 1)
end

@inline function get_indices{D,N}(::Type{Polynomial}, G::Type{Irregular{D,N}}, I, x)
    return ()
end

@inline function get_value{T,N}(::Type{Polynomial}, I::PolynomialInterpolation{T,N}, x, inds)
    coeff = I.coeff
    xx = I.x
    y = 0.0
    @inbounds for i=1:N
        tmp = 1.0
        for j = 1:N
            j == i && continue
            tmp = tmp * (x[1] - xx[j])
        end
        y = y + tmp*coeff[i]
    end
    return y
end


using Images, ImageFeatures, ImageDraw, ImageMorphology

#### Example 1

img = Gray.(load("sample-images/cat-3417184_640.jpg"))
img_f = Float16.(restrict(img))

new_img = Gray.(hcat(
    Float32.(img_f) .* (~fastcorners(img_f, 12, 0.15)),
    Float32.(img_f) .* (~fastcorners(img_f, 12, 0.25)),
    Float32.(img_f) .* (~fastcorners(img_f, 12, 0.35))
))

new_img = Gray.(hcat(
    Float32.(img_f) .* (~fastcorners(img_f, 12, 0.15)),
    Float32.(img_f) .* (~imcorner(img_f, method=harris)),
    Float32.(img_f) .* (~imcorner(img_f, method=shi_tomasi))
))

imshow(new_img)

### Example 2

img = Gray.(load("sample-images/newspaper-37782_640.png"))
img_f = Float16.(img)

new_img = Gray.(hcat(
    img_f,
    Float16.(img_f) .* (~fastcorners(img_f, 12, 0.15))
))

imshow(new_img)

### Example 3
img = Gray.(restrict(load("sample-images/cat-3417184_640.jpg")))

img = restrict(load("sample-images/board-157165_640.png"))
img_f = Float16.(Gray.(img))

img_harris = copy(img)
img_harris[dilate(imcorner(img_f, method=harris)) .> 0.01] = colorant"yellow"

img_shi = copy(img)
img_shi[dilate(imcorner(img_f, method=shi_tomasi)) .> 0.01] = colorant"yellow"

img_rosenfield = copy(img)
img_rosenfield[dilate(imcorner(img_f, method=kitchen_rosenfeld)) .> 0.01] = colorant"yellow"

img_fast = copy(img)
img_fast[dilate(fastcorners(img_f, 12, 0.05)) .> 0.01] = colorant"yellow"

new_img = vcat(
    hcat(
        img_harris,
        img_shi),
    hcat(
        img_rosenfield,
        img_fast)
)

new_img[Int(size(new_img, 1) / 2), :] = colorant"yellow"
new_img[:, Int(size(new_img, 2) / 2)] = colorant"yellow"

imshow(new_img)

### Example 4
img = Gray.(restrict(load("sample-images/cat-3417184_640.jpg")))
img_f = Float16.(Gray.(img))

img_harris = copy(img)
img_harris[dilate(imcorner(img_f, method=harris)) .> 0.01] = colorant"yellow"

img_shi = copy(img)
img_shi[dilate(imcorner(img_f, method=shi_tomasi)) .> 0.01] = colorant"yellow"

img_rosenfield = copy(img)
img_rosenfield[dilate(imcorner(img_f, method=kitchen_rosenfeld)) .> 0.01] = colorant"yellow"

img_fast = copy(img)
img_fast[dilate(fastcorners(img_f, 12, 0.05)) .> 0.01] = colorant"yellow"

new_img = vcat(
    hcat(
        img_harris,
        img_shi),
    hcat(
        img_rosenfield,
        img_fast)
)

imshow(new_img)

### Example 5

img = Gray.(restrict(load("sample-images/cat-3417184_640.jpg")))
img_f = Float16.(Gray.(img))

img_harris = copy(img)
img_harris[dilate(imcorner(img_f, Percentile(95), method=harris)) .> 0.01] = colorant"yellow"

img_shi = copy(img)
img_shi[dilate(imcorner(img_f, Percentile(95), method=shi_tomasi)) .> 0.01] = colorant"yellow"

img_rosenfield = copy(img)
img_rosenfield[dilate(imcorner(img_f, Percentile(95), method=kitchen_rosenfeld)) .> 0.01] = colorant"yellow"

img_fast = copy(img)
img_fast[dilate(fastcorners(img_f, 12, 0.05)) .> 0.01] = colorant"yellow"

new_img = vcat(
    hcat(
        img_harris,
        img_shi),
    hcat(
        img_rosenfield,
        img_fast)
)

imshow(new_img)

### Performance comparison

img = restrict(load("sample-images/cat-3417184_640.jpg"))
img_f = Float16.(Gray.(img))

@btime fastcorners(img_f, 12, 0.15);
@btime fastcorners(img_f, 12, 0.05);
@btime imcorner(img_f, method=harris);
@btime imcorner(img_f, Percentile(95), method=harris);
@btime imcorner(img_f, method=shi_tomasi);
@btime imcorner(img_f, method=kitchen_rosenfeld);


using Test
using Random

function michaelis_menten(u, p, t)
    [0.6 - 1.5u[1]/(0.3+u[1])] # Should be 0.6*0.3+0.6u[1] - 1.5u[1] = u[2]*u[1]-0.3*u[2]
end

u0 = [0.5]
tspan = (0.0, 4.0)

problem_1 = ODEProblem(michaelis_menten, u0, tspan)
solution_1 = solve(problem_1, Tsit5(), saveat = 0.1)
problem_2 = ODEProblem(michaelis_menten, 2*u0, tspan)
solution_2 = solve(problem_2, Tsit5(), saveat = 0.1)
X = [solution_1[:,:] solution_2[:,:]]
ts = [solution_1.t; solution_2.t]

DX = similar(X)
for (i, xi) in enumerate(eachcol(X))
    DX[:, i] = michaelis_menten(xi, [], ts[i])
end

@parameters t
@variables u[1:2]
u = collect(u)
h = [monomial_basis(u[1:1], 4)...]
basis = Basis([h; h .* u[2]], u)

@testset "Ideal data" begin

    prob = ContinuousDataDrivenProblem(X, ts, DX)

    opts = [ImplicitOptimizer(5e-1);ImplicitOptimizer(0.4:0.1:0.7)]
    for opt in opts
        res = solve(prob, basis, opt,u[2:2] ,normalize = false, denoise = false, maxiter = 10000)
        m = metrics(res)
        @test all(m[:L₂] .< 1e-1)
        @test all(m[:AIC] .> 1000.0)
        @test all(m[:R²] .> 0.6)
    end

end

Random.seed!(2345)
X = X .+ 1e-3*randn(size(X))

@testset "Noisy data" begin

    prob = ContinuousDataDrivenProblem(X, ts, GaussianKernel())

    for opt in [ImplicitOptimizer(3e-1);ImplicitOptimizer(3e-1:0.1:7e-1)]
        res = solve(prob, basis, opt, u[2:2], normalize = false, denoise = true)
        m = metrics(res)
        @test all(m[:L₂] .< 1e-1)
        @test all(m[:AIC] .> 1000.0)
        @test all(m[:R²] .> 0.9)
    end

end


"""
    struct PSpace1D{TR,TI<:Integer,TA<:AbstractArray} <: AbstractPhysicalSpace
        x0::TR
        x1::TR
        nx::TI
        x::TA
        dx::TA
    end

1D physical space with structured mesh

"""
struct PSpace1D{TR,TI<:Integer,TA<:AbstractArray} <: AbstractPhysicalSpace1D
    x0::TR
    x1::TR
    nx::TI
    x::TA
    dx::TA
end

function PSpace1D(X0::TR, X1::TR, NX::TI, NG = 0::Integer) where {TR,TI}
    TX = ifelse(TR == Float32, Float32, Float64)

    δ = (X1 - X0) / NX
    x = OffsetArray{TX}(undef, 1-NG:NX+NG)
    dx = similar(x)

    # uniform mesh
    for i in eachindex(x)
        x[i] = X0 + (i - 0.5) * δ
        dx[i] = δ
    end

    return PSpace1D{TR,TI,typeof(x)}(X0, X1, NX, x, dx)
end

PSpace1D() = PSpace1D(0, 1, 100)
PSpace1D(X0::T, X1::T) where {T} = PSpace1D(X0, X1, 100)

"""
    struct PSpace2D{TR<:Real,TI<:Integer,TA<:AbstractMatrix{<:Real},TB<:AbstractArray{<:Real,4}} <: AbstractPhysicalSpace2D
        x0::TR
        x1::TR
        nx::TI
        y0::TR
        y1::TR
        ny::TI
        x::TA
        y::TA
        dx::TA
        dy::TA
        vertices::TB
    end

2D Physical space with structured mesh

"""
struct PSpace2D{
    TR<:Real,
    TI<:Integer,
    TA<:AbstractMatrix{<:Real},
    TB<:AbstractArray{<:Real,4},
} <: AbstractPhysicalSpace2D
    x0::TR
    x1::TR
    nx::TI
    y0::TR
    y1::TR
    ny::TI
    x::TA
    y::TA
    dx::TA
    dy::TA
    vertices::TB
end

function PSpace2D(
    X0::TR,
    X1::TR,
    NX::TI,
    Y0::TR,
    Y1::TR,
    NY::TI,
    NGX = 0::Integer,
    NGY = 0::Integer,
) where {TR,TI}
    TX = ifelse(TR == Float32, Float32, Float64)

    δx = (X1 - X0) / NX
    δy = (Y1 - Y0) / NY
    x = OffsetArray{TX}(undef, 1-NGX:NX+NGX, 1-NGY:NY+NGY)
    y = OffsetArray{TX}(undef, 1-NGX:NX+NGX, 1-NGY:NY+NGY)
    dx = OffsetArray{TX}(undef, 1-NGX:NX+NGX, 1-NGY:NY+NGY)
    dy = OffsetArray{TX}(undef, 1-NGX:NX+NGX, 1-NGY:NY+NGY)

    for j in axes(x, 2)
        for i in axes(x, 1)
            x[i, j] = X0 + (i - 0.5) * δx
            y[i, j] = Y0 + (j - 0.5) * δy
            dx[i, j] = δx
            dy[i, j] = δy
        end
    end

    vertices = OffsetArray{TX}(undef, 1-NGX:NX+NGX, 1-NGY:NY+NGY, 4, 2)
    for j in axes(vertices, 2), i in axes(vertices, 1)
        vertices[i, j, 1, 1] = x[i, j] - 0.5 * dx[i, j]
        vertices[i, j, 2, 1] = x[i, j] + 0.5 * dx[i, j]
        vertices[i, j, 3, 1] = x[i, j] + 0.5 * dx[i, j]
        vertices[i, j, 4, 1] = x[i, j] - 0.5 * dx[i, j]

        vertices[i, j, 1, 2] = y[i, j] - 0.5 * dy[i, j]
        vertices[i, j, 2, 2] = y[i, j] - 0.5 * dy[i, j]
        vertices[i, j, 3, 2] = y[i, j] + 0.5 * dy[i, j]
        vertices[i, j, 4, 2] = y[i, j] + 0.5 * dy[i, j]
    end

    return PSpace2D{TR,TI,typeof(x),typeof(vertices)}(
        X0,
        X1,
        NX,
        Y0,
        Y1,
        NY,
        x,
        y,
        dx,
        dy,
        vertices,
    )
end

PSpace2D() = PSpace2D(0, 1, 45, 0, 1, 45)
PSpace2D(X0::T, X1::T, Y0::T, Y1::T) where {T} = PSpace2D(X0, X1, 45, Y0, Y1, 45)

"""

2D Circular space in polar coordinates

"""
struct CSpace2D{
    TR<:Real,
    TI<:Integer,
    TA<:AbstractMatrix{<:Real},
    TB<:AbstractArray{<:Real,4},
} <: AbstractPhysicalSpace2D
    r0::TR
    r1::TR
    nr::TI
    θ0::TR
    θ1::TR
    nθ::TI
    r::TA
    θ::TA
    x::TA
    y::TA
    dr::TA
    dθ::TA
    darc::TA
    vertices::TB
end

function CSpace2D(
    R0::TR,
    R1::TR,
    NR::TI,
    θ0,
    θ1,
    Nθ::TI,
    NGR = 0::TI,
    NGθ = 0::TI,
) where {TR,TI<:Integer}
    TX = ifelse(TR == Float32, Float32, Float64)

    δr = (R1 - R0) / NR
    δθ = (θ1 - θ0) / Nθ
    r = OffsetArray{TX}(undef, 1-NGR:NR+NGR, 1-NGθ:Nθ+NGθ)
    θ = OffsetArray{TX}(undef, 1-NGR:NR+NGR, 1-NGθ:Nθ+NGθ)
    x = OffsetArray{TX}(undef, 1-NGR:NR+NGR, 1-NGθ:Nθ+NGθ)
    y = OffsetArray{TX}(undef, 1-NGR:NR+NGR, 1-NGθ:Nθ+NGθ)
    dr = OffsetArray{TX}(undef, 1-NGR:NR+NGR, 1-NGθ:Nθ+NGθ)
    dθ = OffsetArray{TX}(undef, 1-NGR:NR+NGR, 1-NGθ:Nθ+NGθ)
    darc = OffsetArray{TX}(undef, 1-NGR:NR+NGR, 1-NGθ:Nθ+NGθ)

    for j in axes(r, 2)
        for i in axes(r, 1)
            r[i, j] = R0 + (i - 0.5) * δr
            θ[i, j] = θ0 + (j - 0.5) * δθ
            x[i, j] = r[i, j] * cos(θ[i, j])
            y[i, j] = r[i, j] * sin(θ[i, j])
            dr[i, j] = δr
            dθ[i, j] = δθ
            darc[i, j] = r[i, j] * dθ[i, j]
        end
    end

    vertices = OffsetArray{TX}(undef, 1-NGR:NR+NGR, 1-NGθ:Nθ+NGθ, 4, 2)
    for j in axes(vertices, 2), i in axes(vertices, 1)
        vertices[i, j, 1, 1] = (r[i, j] - 0.5 * dr[i, j]) * cos(θ[i, j] - 0.5 * dθ[i, j])
        vertices[i, j, 2, 1] = (r[i, j] + 0.5 * dr[i, j]) * cos(θ[i, j] - 0.5 * dθ[i, j])
        vertices[i, j, 3, 1] = (r[i, j] + 0.5 * dr[i, j]) * cos(θ[i, j] + 0.5 * dθ[i, j])
        vertices[i, j, 4, 1] = (r[i, j] - 0.5 * dr[i, j]) * cos(θ[i, j] + 0.5 * dθ[i, j])

        vertices[i, j, 1, 2] = (r[i, j] - 0.5 * dr[i, j]) * sin(θ[i, j] - 0.5 * dθ[i, j])
        vertices[i, j, 2, 2] = (r[i, j] + 0.5 * dr[i, j]) * sin(θ[i, j] - 0.5 * dθ[i, j])
        vertices[i, j, 3, 2] = (r[i, j] + 0.5 * dr[i, j]) * sin(θ[i, j] + 0.5 * dθ[i, j])
        vertices[i, j, 4, 2] = (r[i, j] - 0.5 * dr[i, j]) * sin(θ[i, j] + 0.5 * dθ[i, j])
    end

    return CSpace2D{TR,TI,typeof(x),typeof(vertices)}(
        R0,
        R1,
        NR,
        TX(θ0),
        TX(θ1),
        Nθ,
        r,
        θ,
        x,
        y,
        dr,
        dθ,
        darc,
        vertices,
    )
end

"""
    uniform_mesh(x0::Real, xnum::Int, dx::Real)

Generate uniform mesh
"""
function uniform_mesh(x0, xnum::T, dx) where {T<:Int}
    points = zeros(xnum)
    for i = 1:xnum
        points[i] = x0 + (i - 0.5) * dx
    end

    return points
end

"""
    ndgrid(v::AbstractVector)
    ndgrid(v1::AbstractVector{T}, v2::AbstractVector{T}) where T
    ndgrid(vs::AbstractVector{T}...) where T

Equivalent N-dimensional mesh generator as matlab
"""
ndgrid(v::AbstractVector) = copy(v)

function ndgrid(v1::AbstractVector{T}, v2::AbstractVector{T}) where {T}
    m, n = length(v1), length(v2)
    v1 = reshape(v1, m, 1)
    v2 = reshape(v2, 1, n)

    return (repeat(v1, 1, n), repeat(v2, m, 1))
end

function ndgrid(vs::AbstractVector{T}...) where {T}
    ndgrid_fill(a, v, s, snext) = begin
        for j = 1:length(a)
            a[j] = v[div(rem(j - 1, snext), s)+1]
        end
    end

    n = length(vs)
    sz = map(length, vs)
    out = ntuple(i -> Array{T}(undef, sz), n)
    s = 1
    for i = 1:n
        a = out[i]::Array
        v = vs[i]
        snext = s * size(a, i)
        ndgrid_fill(a, v, s, snext)
        s = snext
    end

    return out
end

"""
    2D: meshgrid(x::AbstractArray{<:Real,1}, y::AbstractArray{<:Real,1})
    3D: meshgrid(x::AbstractArray{<:Real,1}, y::AbstractArray{<:Real,1}, z::AbstractArray{<:Real,1})

Equivalent structured mesh generator as matlab
"""
meshgrid(v::AbstractVector{T}) where {T} = meshgrid(v, v)

function meshgrid(x::T, y::T) where {T<:AbstractVector{<:Real}}
    X = [i for j in y, i in x]
    Y = [j for j in y, i in x]

    return X, Y
end

function meshgrid(x::T, y::T, z::T) where {T<:AbstractVector{<:Real}}
    X = [i for k in z, j in y, i in x]
    Y = [j for k in z, j in y, i in x]
    Z = [k for k in z, j in y, i in x]

    return X, Y, Z
end

"""
    find_idx(x::AbstractArray{<:Real,1}, p::Real; mode = :nonuniform::Symbol)

Find the location index of a point in mesh

- @arg x: center locations of 1D mesh
- @arg p: point location
- @arg mode: choose uniform / nonuniform formulations

"""
function find_idx(
    x::T,
    p::Real;
    mode = :nonuniform::Symbol,
) where {T<:AbstractArray{<:Real,1}}

    if mode == :uniform
        dx = x[2] - x[1]
        return Int(ceil((p - x[1] + 0.5 * dx) / dx)) # point location
    else
        return argmin(abs.(x .- p)) # center location
    end

end

# ------------------------------------------------------------
# Extended Base.show()
# ------------------------------------------------------------

function Base.show(io::IO, ps::PSpace1D{TR,TI,TA}) where {TR,TI,TA}
    print(
        io,
        "PhysicalSpace1D{$TR,$TI,$TA}\n",
        "domain: ($(ps.x0),$(ps.x1))\n",
        "resolution: $(ps.nx)\n",
        "ghost: $(1-firstindex(ps.x))\n",
    )
end

function Base.show(io::IO, ps::PSpace2D{TR,TI,TA}) where {TR,TI,TA}
    print(
        io,
        "PhysicalSpace2D{$TR,$TI,$TA}\n",
        "domain: ($(ps.x0),$(ps.x1)) × ($(ps.y0),$(ps.y1))\n",
        "resolution: $(ps.nx) × $(ps.ny)\n",
        "ghost in x: $(1-firstindex(ps.x[:, 1]))\n",
        "ghost in y: $(1-firstindex(ps.y[1, :]))\n",
    )
end

function Base.show(io::IO, ps::CSpace2D{TR,TI,TA}) where {TR,TI,TA}
    print(
        io,
        "CircularSpace2D{$TR,$TI,$TA}\n",
        "domain: ($(ps.r0),$(ps.r1)) × ($(ps.θ0),$(ps.θ1))\n",
        "resolution: $(ps.nr) × $(ps.nθ)\n",
        "ghost in r: $(1-firstindex(ps.x[:, 1]))\n",
        "ghost in θ: $(1-firstindex(ps.y[1, :]))\n",
    )
end


using StatsBase, Distributions
using MDP

# Based on the description in Sec 8 of the following Technical Report
# > Shalabh Bhatnagar, Richard S. Sutton, Mohammad Ghavamzadeh, and Mark Lee
# > "Natural Actor–Critic Algorithms",
# > University of Alberta Department of Computing Science Technical Report TR09-10
# > June 2009.

# In the description of Garnet, there is a parameter τ that determines
# the non-stationarity of the mode. Here, we only implement the case when
# τ = 0. We will also assume that b << n, so we generate a sparse P matrix.

# The description of the garnet model says that the rewards should have normal
# distribution. We let the user to explicitly specify the reward distribution.

function garnet(n, m, b; rewards=Uniform(0,1), objective=:Max)
   P = SparseMatrixCSC[ spzeros(n,n) for a in 1:m ]
   S = 1:n
   for a in 1:m
       for i in 1:n
           nextState = sample(S, b; replace=false)
           nextProb  = rand(b)
           nextProb /= sum(nextProb)
           P[a][i, nextState] = nextProb
       end
   end
   c = rand(rewards, n,m)
   model = ProbModel(c, P; objective=objective)
end

function randomPolicy(model :: ProbModel; stochastic = true)
  (n, m) = (model.stateSize, model.actionSize)

  Q = Matrix[zero(model.P[1])]
  c = zeros(n,1)

  if !stochastic
    a = sample(1:m, n; replace=true)
    for i in 1:n
      Q[1][i,:] = model.P[a[i]][i,:]
      c[i,1] = model.cost[i, a[i]]
    end
  else
    for i in 1:n
      p  = rand(m)
      p /= sum(p)

      for a in 1:m
        Q[1][i,:] += model.P[a][i,:]*p[a]
        c[i,1]    += model.cost[i,a]*p[a]
      end
    end
  end

  return ProbModel(c, Q)
end


abstract MFOperation <: Membershipfunction

#############################
# Interection of 2 Membership Functions
immutable T1IntersectOP <: MFOperation
    t1::Membershipfunction
    t2::Membershipfunction
end
# export T1IntersectOP

import Base: getindex
getindex(op::T1IntersectOP, value::Real) = tnorm(op.t1[value], op.t2[value])
export getindex

# Union of 2 Membership Functions
immutable T1UnionOP <: MFOperation
    t1::Membershipfunction
    t2::Membershipfunction
end
# export T1UnionOP

import Base: getindex
getindex(op::T1UnionOP, value::Real) = tconorm(op.t1[value], op.t2[value])
export getindex

# Complement of
immutable T1ComplementOP <: MFOperation
    t1::Membershipfunction
end
import Base: getindex
getindex(op::T1ComplementOP, value::Real) = 1.0 - op.t1[value]
export getindex
# export T1ComplementOP

Base.complement(mf::Membershipfunction) = T1ComplementOP(mf)
export complement

Base.intersect(mf::Membershipfunction, mfs::Membershipfunction...) = intersect(T1IntersectOP(mf,mfs[1]),mfs[2:end]...)
Base.intersect(mf::Membershipfunction, mf1::Membershipfunction) = T1IntersectOP(mf,mf1)
export intersect

Base.union(mf::Membershipfunction, mfs::Membershipfunction...) = union(T1UnionOP(mf,mfs[1]),mfs[2:end]...)
Base.union(mf::Membershipfunction, mf1::Membershipfunction) = T1UnionOP(mf,mf1)
export union


# Estimates the Klein consumption equation by GMM
using DelimitedFiles, Statistics
function KleinMoments(theta, y, x, z)
	e = y - x*theta
	m = e.*z
end

function main()
data = readdlm("klein.data")
# construct missing lags, and drop first row that has missing data
profits = data[:,3]
output = data[:,7]
data = [data lag(profits,1) lag(output,1)]
data = data[2:end,:]
n = size(data,1)
# define instruments
exogs = [1, 6, 8, 9, 10, 11, 12]
exogs = data[:,exogs]
exogs = [ones(n,1) exogs]
# CONSUMPTION
println("CONSUMPTION EQUATION")
# define variables in consumption equation
y = data[:,2]
profits = data[:,3]
lagprofits = data[:,11]
wp = data[:,4]
wg = data[:,8]
wages = wp + wg
# regressors in consumption equation
x = [profits lagprofits wages]
x = [ones(n,1) x]
# GMM estimation
theta = x\y
weight = 1.0
names = ["Constant", "Profits", "Profits-1", "Wages"]
moments = θ -> KleinMoments(θ, y, x, exogs)
# initial consistent estimate: only used to get moment covariance (needed for t-stats) no screen output
theta, obj_value, D, ms, convergence = gmm(moments, vec(theta), weight)
# moment covariance assuming no autocorrelation
momentcov = cov(ms)
weight = inv(momentcov)
# estimation results using efficient weight (no autocorrelation)
gmmtitle = "Klein model 1 GMM example, plain covariance"
gmmresults(moments, theta, weight, gmmtitle, names)
# moment covariance assuming autocorrelation
# note: if there really is autocorrelation,
# then lagged endogs need to be dropped as
# instruments. This is not done here, as this
# is just meant as an example of use of NW
# covariance estimator
momentcov = NeweyWest(ms)
weight = inv(momentcov)
# estimation results using efficient weight (NW)
gmmtitle = "Klein model 1 GMM example, NW covariance"
gmmresults(moments, theta, weight, gmmtitle, names)
return
end
main()


f = open("input02.txt", "r")
    l = readlines(f)
close(f)

horizontal = 0
deep = 0
for a in l
    s = split(a, " ")
    if s[1] == "forward"
        global horizontal += parse(Int, s[2])
    elseif s[1] == "down"
        global deep += parse(Int, s[2])
    elseif s[1] == "up"
        global deep += -parse(Int, s[2])
    else
        println("Unknown entry in input")
    end
end
product = horizontal * deep
println("Part 1")
println("Horizontal : $horizontal")
println("Depth : $deep")
println("Product : $product")

horizontal = 0
deep = 0
aim = 0
for a in l
    s = split(a, " ")
    if s[1] == "down"
        global aim += parse(Int, s[2])
    elseif s[1] == "up"
        global aim += -parse(Int, s[2])
    elseif s[1] == "forward"
        global horizontal += parse(Int, s[2])
        global deep += (parse(Int, s[2]) * aim)
    else
        println("Unknown entry in input")
    end
end
product = horizontal * deep
println("Part 2")
println("Horizontal : $horizontal")
println("Depth : $deep")
println("Product : $product")


using Spectra
using Base.Test

# Dummy data
x = collect(0.:1.:100.)
y = collect(0.:1.:100.)
x2 = collect(-100.:0.1:100.)
y_poly4 = 1.0+ x2 + 0.25*x2.^2 + 0.0001.*x2.^3 + 1.e-6.*x2.^4
roi = [0.0 100.0]

# testing baselines
y_fit_gcvspl,bas_fit_gcvspl = baseline(x[:],y[:],roi,"gcvspline",p=1.0)
@test_approx_eq_eps(y_fit_gcvspl,(y[:]-y[:]),1e-10)
@test_approx_eq_eps(bas_fit_gcvspl,y[:],1e-10)

y_fit_dierckx,bas_fit_dierckx = baseline(x[:],y[:],roi,"Dspline",p=1.0)
@test_approx_eq_eps(y_fit_dierckx,(y[:]-y[:]),1e-10)
@test_approx_eq_eps(bas_fit_dierckx,y[:],1e-10)

y_fit_linear,bas_fit_linear = baseline(x[:],y[:],roi,"poly",p=2.0)
@test_approx_eq_eps(y_fit_linear,(y[:]-y[:]),1e-10)
@test_approx_eq_eps(bas_fit_linear,y[:],1e-10)

y_fit_poly,bas_fit_poly = baseline(x2[:],y_poly4[:],roi,"poly",p=4.0)
@test_approx_eq_eps(y_fit_poly,(y_poly4[:]-y_poly4[:]),1e-3)
@test_approx_eq_eps(bas_fit_poly,y_poly4[:],1e-3)

#Tolerances for ML algo are high, it is just to test that they are running well...
y_fit_kr,bas_fit_kr = baseline(x[:],y[:],roi,"KRregression")
@test_approx_eq_eps(y_fit_kr,(y[:]-y[:]),1e-1)
@test_approx_eq_eps(bas_fit_kr,y[:],1e-1)

y_fit_svm,bas_fit_svm = baseline(x[:],y[:],roi,"SVMregression")
@test_approx_eq_eps(y_fit_svm,(y[:]-y[:]),10.)
@test_approx_eq_eps(bas_fit_svm,y[:],10.)

# testing long correction function
x_for_long = [20.,21.,22.,23.,24.,25.]
y_for_long = [1.0,1.0,1.0,1.0,1.0,1.0]
h_plank = 6.626070040e-34 #Plank constant
k_bolt = 1.38064852e-23  #Boltzman constant
c_light = 299792458. #speed of light in vaccum
nu0 = 1.0./(514.532).*1e9 #laser wavenumber at 514.532
nu = 100.0.*x_for_long # cm-1 to m-1
T = 23.0+273.15 # the temperature in K

x_long,long_res,eselong = tlcorrection([x_for_long y_for_long],23.0,514.532) # using the function
t0 = nu0.^3.*nu./((nu0-nu).^4)
t1= 1 - exp(-h_plank.*c_light.*nu./(k_bolt.*T)) # c in m/s  : t1 dimensionless
long_calc= y_for_long.*t0.*t1 # pour les y
long_calc = long_calc./trapz(x_for_long,long_calc) # area normalisation

@test_approx_eq(long_res,long_calc)
@test_approx_eq(x_for_long,x_long)

#figure()
#plot(x2,y_poly4,"o")
#plot(x2,bas_fit_poly,"-")


"""
    struct JohnsonState{T, U}

An [`AbstractPathState`](@ref) designed for Johnson shortest-paths calculations.
"""
struct JohnsonState{T <: Real,U <: Integer} <: AbstractPathState
    dists::Matrix{T}
    parents::Matrix{U}
end

"""
    johnson_shortest_paths(g, distmx=weights(g))

Use the [Johnson algorithm](https://en.wikipedia.org/wiki/Johnson%27s_algorithm)
to compute the shortest paths between all pairs of vertices in graph `g` using an
optional distance matrix `distmx`.

### Implementation Notes
If the parameter parallel is set true, dijkstra_shortest_paths will run in parallel.
Parallel bellman_ford_shortest_paths is currently unavailable
Return a [`LightGraphs.JohnsonState`](@ref) with relevant
traversal information.
Behaviour in case of negative cycle depends on bellman_ford_shortest_paths.
Throws NegativeCycleError() if a negative cycle is present.

### Performance
Complexity: O(|V|*|E|)
If distmx is not mutable or of type, DefaultDistance than a sparse matrix will be produced using distmx.
In the case that distmx is immutable, to reduce memory overhead,
if edge (a, b) does not exist in g then distmx[a, b] should be set to 0.
"""
function johnson_shortest_paths(g::AbstractGraph{U},
    distmx::AbstractMatrix{T}=weights(g)) where T <: Real where U <: Integer

    nvg = nv(g)
    type_distmx = typeof(distmx)
    #Change when parallel implementation of Bellman Ford available
    wt_transform = bellman_ford_shortest_paths(g, vertices(g), distmx).dists

    if !type_distmx.mutable && type_distmx !=  LightGraphs.DefaultDistance
        distmx = sparse(distmx) #Change reference, not value
    end

    #Weight transform not needed if all weights are positive.
    if type_distmx !=  LightGraphs.DefaultDistance
        for e in edges(g)
            distmx[src(e), dst(e)] += wt_transform[src(e)] - wt_transform[dst(e)]
        end
    end

    dists = Matrix{T}(undef, nvg, nvg)
    parents = Matrix{U}(undef, nvg, nvg)
    for v in vertices(g)
        dijk_state = dijkstra_shortest_paths(g, v, distmx)
        dists[v, :] = dijk_state.dists
        parents[v, :] = dijk_state.parents
    end

    broadcast!(-, dists, dists, wt_transform)
    for v in vertices(g)
        dists[:, v] .+= wt_transform[v] #Vertical traversal prefered
    end

    if type_distmx.mutable
        for e in edges(g)
            distmx[src(e), dst(e)] += wt_transform[dst(e)] - wt_transform[src(e)]
        end
    end

    return JohnsonState(dists, parents)
end

function enumerate_paths(s::JohnsonState{T,U}, v::Integer) where T <: Real where U <: Integer
    pathinfo = s.parents[v, :]
    paths = Vector{Vector{U}}()
    for i in 1:length(pathinfo)
        if (i == v) || (s.dists[v, i] == typemax(T))
            push!(paths, Vector{U}())
        else
            path = Vector{U}()
            currpathindex = i
            while currpathindex != 0
                push!(path, currpathindex)
                currpathindex = pathinfo[currpathindex]
            end
            push!(paths, reverse(path))
        end
    end
    return paths
end

enumerate_paths(s::JohnsonState) = [enumerate_paths(s, v) for v in 1:size(s.parents, 1)]
enumerate_paths(st::JohnsonState, s::Integer, d::Integer) = enumerate_paths(st, s)[d]


using LsqFit
using Optim
using MicroLogging
using PyPlot
using Roots

#Fitting a biased Lorentzian to amplitude data.
"""
Returns a five parameter, biased Lorentizian function

Args:
  x: domain of the returned function
  p: a vector of length 5 with function params --
      (amplitude, frequency, line width, linear skew, offset)
Returns:
    The computed function.

See appendix E.4 of Gao\'s thesis (2008)
"""
function bias_lorentzian(x, p)
  return p[1] ./ ((x .- p[2]).^2 .+ (p[3]/2)^2) .+ p[4]*(x .- p[2]) .+ p[5];
end

"""
Fit frequency-amplitude data to a skewed Lorentzian using nonlinear least-squares.
Model is y = a / ((x - b)^2 + (c/2)^2) + d*(x - b) + e

Args:
  freq: Frequency data.
  data: Data to be fitted.
  vars: Variance of data. Defaults to empty.
Returns:
  fit: LsqFitResult object.
  fit_function(x): Biased Lorentzian at fitted parameters.
"""
function fit_biased_lorentzian(freq, data, vars=[])

  BL_fit_dict(p) = Dict("a" => p[1], "b" => p[2], "c" => p[3],
                                "d" => p[4], "e" => p[5])
  p0 = initial_guess_blorentz(freq, data)
  return generic_fit(freq, data, bias_lorentzian, p0, BL_fit_dict,
    "y = a / ((x - b)^2 + (c/2)^2) + d*(x - b) + e", yvars=vars)
end

function initial_guess_blorentz(xpts, ypts)
  """
  Returns a length five vector with an initial guess of
  Lorentizian function parameters.
  """
  # find offset
  e = median(ypts)
  if abs(maximum(ypts) - median(ypts)) <= abs(minimum(ypts) - median(ypts))
        idx = findmin(ypts)[2]
        direction = -1
    else
        idx = argmax(ypts)
        direction = 1
  end
  # center frequency
  b = xpts[idx]
  half = abs(median(ypts) + ypts[idx]) / 2.
  if direction == 1
    idx_l = findfirst(x -> x > half, ypts)
    idx_r = findlast(x -> x > half, ypts)
  else
    idx_l = findfirst(x -> x < half, ypts)
    idx_r = findlast(x -> x < half, ypts)
  end
  c = abs(xpts[idx_l] - xpts[idx_r])
  a =  direction * c^2 * abs(ypts[idx] - e) / 4
  # slope of linear skew
  d = (ypts[end] - ypts[1])/(xpts[end] - xpts[1])
  return [a, b, c, d, e];
end

struct CircleFitParams
  """Container for the result of a circle fit.
    f0: Resonator center frequency.
    Qi: Resonator internal quality factor.
    Qc: Resonator coupling quality factor.
    ϕ: Impedance mismatch angle.
    τ: System phase delay.
    α: System gain angle.
    A: System gain amplitude.
  """
  f0
  Qi
  Qc
  ϕ
  τ
  α
  A
end

struct CircleFitResult
  fit_params::CircleFitParams
  sq_error::Float64
  SNR::Float64
  errors::CircleFitParams
  fit_curve::Function
end

#Fitting to the resonance circle of a quarter-wave resonator
function fit_resonance_circle(freq::Vector{T}, data::Vector{Complex{T}}; kwargs...) where T <: AbstractFloat
  """
  Fits complex-valued data

  Args:
    freq: Frequency data.
    data: Complex S21 data.
    kwargs: Use these if you already know something about the system and would like to override the
      following properties
      τ: Cable delay
      α: System phase shifted
      A: Overall system gain
  Returns:
    Result: The result of the fit.
  """
  @assert length(freq) == length(data) "Frequency and Data vectors must have same length!"
  @assert length(data) > 20 "Too few points."

  kwdict = Dict(kwargs)

  function χ2(data, R, xc, yc)
    x = real.(data)
    y = imag.(data)
    d = sqrt.( (x-xc).^2 + (y-yc).^2 ) - R
    return sum(d.^2)
  end

  #Fit to cable delay
  haskey(kwdict, :τ) ? τ = kwdict[:τ] : τ = fit_delay(freq, data)

  Sp = exp.(1im * 2. * π * freq * τ) .* data
  #Get best-fit circle and translate to origin.
  R, xc, yc = fit_circle(real.(Sp), imag.(Sp))

  #@assert χ2 < 10 "Could not fit circle to delay-corrected data. χ² = $(χ2)."
  @debug "Fit circle to delay corrected data with χ² = $(χ2(Sp, R, xc, yc))."

  #Translate circle to origin and fit overall phase delay and scaling
  St = Sp - (xc + 1im * yc)
  _, _, θ = fit_phase(freq, St)
  #find f → ∞ point -- this part seems fragile...
  β = mod(θ + π, π)
  P = xc + R*cos(β) + 1im*(yc + R*sin(β))

  @debug "Found scale parameters A = $(abs(P)), α = $(angle(P))."

  haskey(kwdict, :A) ? A = kwdict[:A] : A = abs(P)
  haskey(kwdict, :α) ? α = kwdict[:α] : α = angle(P)

  #Calibrate out α, A and get impedance mismatch angle
  Sc = Sp .* exp.(-1im * α) / A
  Rc, xcc, ycc = fit_circle(real.(Sc), imag.(Sc))

  #@assert χ2 < 10 "Could not fit circle to calibrated data. χ² = $(χ2)."
  @debug "Fit circle to calibrated data with χ² = $(χ2(Sc, Rc, xcc, ycc))."

  N = length(data)
  r = sqrt.((real(Sc) - xcc).^2 + (imag(Sc) - ycc).^2)
  SNR = Rc / sqrt(sum((r - Rc).^2)/(N-1))

  @debug "Data has an SNR of $(SNR)."

  ϕ = -asin(ycc/Rc)
  @debug "Found impedance mismatch angle: ϕ = $(ϕ)"
  #Final fit to phase to extract resonant frequency and Q
  Sct = Sc - (xcc + 1im * ycc)
  f0, Qr, _ = fit_phase(freq, Sct)

  Qc_cplx = Qr/(2 * Rc * exp(-1im * ϕ))
  Qi = 1/(1/Qr - real(1/Qc_cplx))
  Qc = 1/abs(1/Qc_cplx)

  @debug "Found quality factors: Qᵢ = $(Qi), Qc = $(Qc)."

  fit_params = CircleFitParams(f0, Qi, Qc, ϕ, τ, α, A)

  refined_params, errors = refine_circle_fit(freq, data, fit_params)

  χ0 = sum(abs.(data - lorentzian_resonance(fit_params, freq)).^2)
  χ1 = sum(abs.(data - lorentzian_resonance(refined_params, freq)).^2)
  @debug "Circle fit found χ² = $(χ0), refined to χ² = $(χ1)"

  if χ0 > χ1
    params = refined_params
    χ = χ1
  else
    params = fit_params
    χ = χ0
  end

  return CircleFitResult(params,
                          χ,
                          SNR,
                          errors,
                          x->lorentzian_resonance(fit_params, x))
end

function refine_circle_fit(freq, data, fit_params)
  """Refine circle fit and get errors using LsqFit.jl"""

  model(x, p) = abs.(lorentzian_resonance(p, x))

  initial_guess = [fit_params.f0, fit_params.Qi, fit_params.Qc,
                  fit_params.ϕ, fit_params.τ, fit_params.α, fit_params.A]

  fit = curve_fit(model, freq, abs.(data), initial_guess)
  errors = margin_error(fit)

  refined_params = CircleFitParams(fit.param[1], fit.param[2], fit.param[3],
                fit.param[4], fit.param[5], fit.param[6], fit.param[7])

  errs = CircleFitParams(errors[1], errors[2], errors[3],
                errors[4], errors[5], errors[6], errors[7])

  return (refined_params, errs)
end

function fit_phase(freq, data)
  """Fit phase of resonance.
  Args:
    freq: Frequency data.
    data: Complex S21 data.
  Returns:
    f0: Center frequency.
    Q: Quality factor.
    Θ0: Offset phase angle.
  """
  model(x, p) = p[1] + 2. * slope * atan.(2 * p[2] *(1. - x / p[3]))
  ϕ = unwrap(angle.(data))
  #guesses for initial parameters
  idx = findmin(abs.(ϕ - mean(ϕ)))[2]
  if mean(ϕ[1:9]) > mean(ϕ[end-9:end])
    j = findfirst(x -> x - ϕ[idx] < π/2., ϕ)
    k = findfirst(x -> x - ϕ[idx] < -π/2., ϕ)
    slope = 1
  else
    j = findfirst(x -> x - ϕ[idx] > π/2., ϕ)
    k = findfirst(x -> x - ϕ[idx] > -π/2., ϕ)
    slope = -1
  end
  Qguess = freq[idx]/abs.(freq[j] - freq[k])
  fit = curve_fit(model, freq, ϕ, [ϕ[idx], Qguess, freq[idx]])
  χ = sum(fit.resid.^2)

  @debug "Frequency vs. Phase fit found: f₀ = $(fit.param[3]), Q = $(fit.param[2]), θ₀ = $(fit.param[3])"
  @debug "Phase fit χ² = $(χ)"
  #@assert χ < 10 "Could not fit phase: χ² = $(χ)"

  return fit.param[3], fit.param[2], fit.param[1]
end

function fit_delay(freq, data)
  """Fit overall phase delay imposed by system transfer function.

  Args:
    freq: Frequency data.
    data: Complex S21 data.
  Returns:
    τ: Linear phase delay
  """
  function delay_model(x)
    ddata = data .* exp.(2. * π * 1im * freq * x)
    R, xc, yc = fit_circle(real.(ddata), imag.(ddata))
    d = sqrt.((real.(ddata) - xc).^2 + (imag.(ddata) - yc).^2) - R
    return sum(d.^2)
  end

  ϕ = unwrap(angle.(data))
  linfit(x,p) = p[1] + x * p[2]
  fit = curve_fit(linfit, freq, ϕ, [mean(ϕ), 0])
  ϕ0 = fit.param[2]
  result = optimize(delay_model, -1.5*abs(ϕ0), 1.5*abs(ϕ0)) #would prefer 1D gradient descent
  τ = Optim.minimizer(result)

  χ2 = delay_model(τ)
  @debug "Cable delay fit found: $(τ) with χ² = $(χ2)."
 #@assert χ2 < 10 "Could not calibrate out cable delay: χ² = $(χ2)."

  return τ
end

function lorentzian_resonance(p::CircleFitParams, f)
  """
  Return a resonance model in S21 amplitude over the range [x] and with
  parameters [p].

  See appendix E of Gao 2008 p. 155 Eq E.1
  """
  Qc_cplx = 1 / ((1/p.Qc)*exp(1im*p.ϕ))
  Q = 1/(1/p.Qi + real(1/Qc_cplx))
  return p.A*exp(1im*p.α).*exp.(-2π*1im*f*p.τ).*(1 - (Q/p.Qc)*exp(1im*p.ϕ)./(1 + 2*1im*Q*(f/p.f0 - 1)));
end

function lorentzian_resonance(p::Array, f)
  return lorentzian_resonance(CircleFitParams(p[1], p[2], p[3], p[4], p[5], p[6], p[7]), f)
end

function simulate_resonance(kwargs...)
  """
  Return synthetic data to test fit functions

  Returns: xpts, ypts, curve_params
  """
  f0 = 6.5 + randn()
  Qc = 120000 + randn()*2000
  Qi = 6e5 + randn()*1e4
  τ = 1.7
  ϕ = 0.17 * π
  Q = 1 ./ (1/Qi + real(1 ./ Qc*exp.(-1im*ϕ)))
  α = 0.35 * π
  A = 0.23
  df = f0/Q;
  freqs = range(f0 - 6*df, stop=f0 + 6*df, length=401)
  p = CircleFitParams(f0, Qi, Qc, ϕ, τ, α, A)
  data = lorentzian_resonance(p, freqs);
  # add some noise
  rand_phase = 2π*0.001*randn(length(freqs));
  rand_amp = 0.002*randn(length(freqs));
  data = data.*exp(-1im*rand_phase).*(1+rand_amp);

  return freqs, data, p
end

function fit_circle(x, y; refine=false)
  """Algebraic fit of (x,y) points to a circle. See:
      N. Chernov and C. J. Lesort, Least Squares Fitting of Circles,
        Journal of Mathematical Imaging and Vision, 23: 239-252, 2005.

      Args:
        x, y: X and Y points to be fitted to a circle.
      Returns:
        R: Circle radius.
        xc: x-coordinate of circle center.
        yc: y-coordinate of circle center.
  """
  @assert length(x) == length(y) "X and Y vector lengths must be equal for circle fit!"
  n = length(x)
  z = x.^2 + y.^2
  Mx = sum(x)
  My = sum(y)
  Mz = sum(z)
  Mxx = sum(x.^2)
  Myy = sum(y.^2)
  Mzz = sum(z.^2)
  Mxy = sum(x .* y)
  Mxz = sum(x .* z)
  Myz = sum(y .* z)
  M = [Mzz Mxz Myz Mz; Mxz Mxx Mxy Mx; Myz Mxy Myy My; Mz Mx My n];
  B = [0 0 0 -2; 0 1 0 0; 0 0 1 0; -2 0 0 0]

  F(η) = det(M - η*B)
  #We use Newton's method to guarantee convergence to the smallest positive eigenvalue
  ηs = newton(F, 0.0)
  ev = nullspace(M - ηs*B)

  xc = -ev[2]/2/ev[1]
  yc = -ev[3]/2/ev[1]
  R = sqrt(ev[2]^2 + ev[3]^2 - 4*ev[1]*ev[4]) / 2 ./ abs(ev[1])

  return R, xc, yc
end

function fit_circle_LM(x, y, initial_guess)
    """Iterative algorithm for refining a circle fit, using radial weighting."""

    function resid(p)
        C = sqrt.((x-p[2]).^2 + (y-p[3]).^2)
        W = 1 ./ sqrt.((p[2]-x).^2 + (p[3]-y).^2)
        return (p[1] - C) .* W
    end

    T = eltype(x)

    result = LsqFit.lmfit(resid, initial_guess, T[])

    return result.param[1], result.param[2], result.param[3]
end


function WNNM(Y::Array{Float32,3}, C::Float32, Sig::Array{Float32,1}, m::Array{Float32,3})

    F = svd(F2cplx(Y))

    lenSig = length(F.S)

    patNum = size(Y)[2]

    mean_Sig_arr = √(sum(Sig[:].^2) / length(Sig[:]))

    TempC  = Float32(C * √(patNum) * 2 * mean_Sig_arr^2)

    X_arr = ClosedQWNNM(F.S, TempC,eps())

    X = F.U * Diagonal(X_arr) * F.Vt

    return cplx2F(X) .+ m
end

function ClosedQWNNM(SigmaY::Array{Float32,1}, c::Float32, our_eps)

    temp = (SigmaY .- our_eps) .^ 2 - 4 * ((-1) * our_eps .* SigmaY .+ c)
    ind = Int32[]
    SigmaX = Float32[]
    for i in 1:length(temp)
        if temp[i] > 0
            push!(ind, i)
            push!(SigmaX, max(SigmaY[i] - our_eps + √(temp[i]) , 0) / 2)
        else
            push!(SigmaX, 0)
        end
    end

    return SigmaX
end

function F2cplx(M::Array{Float32,3})::Array{Complex{Float32},2}
    (h, w) = size(M)

    cM = adjoint(M)

    [cM[1:2*h, 1:w]; cM[1:2*h, w+1:2*w]]
end

function adjoint(A::Array{Float32,3})::Array{Complex{Float32},2}
    W = zeros(size(A)[1:2])
    X = A[:, :, 1]
    Y = A[:, :, 2]
    Z = A[:, :, 3]
    imm = Complex{Float32}(im)
    vcat(hcat(W + imm .* X, Y + imm .* Z),hcat(-Y + imm .* Z, W - imm .* X))
end

function cplx2F(M)
    (h, w) = size(M)

    cM = [M[1:h ÷ 2,1:w] M[1+h ÷ 2:h,1 : w]]

    unadjoint(cM)
end
function unadjoint(A)
    (r, c) = size(A)
    r2 = r ÷ 2
    c2 = c ÷ 2
    A1 = A[     1 : r2,      1 : c2]
    A2 = A[     1 : r2, c2 + 1 : c ]
    A3 = A[r2 + 1 : r,       1 : c2]
    A4 = A[r2 + 1 : r,  c2 + 1 : c ]
    C = zeros(Float32, r2, c2, 3)
    C[:,:,1] = Float32.(imag(A1 - A4))
    C[:,:,2] = Float32.(real(A2 - A3))
    C[:,:,3] = Float32.(imag(A2 + A3) )
    C./2
end


using DifferentialEquations
using Plots
barrier=4

condition(u,t,integrator)=abs(u[1])>barrier || abs(u[2])>barrier || abs(u[3])>barrier||abs(u[4])>barrier
affect!(integrator) = terminate!(integrator)
cb=DiscreteCallback(condition,affect!)
Energy=0.18;
H=(2*Energy)^2
h=sqrt(H)

PofY_find(YY)=sqrt((h+YY.^2-h*YY.^2 )./(1+h-h*YY.^2))
H1(p,q,t)= ( (q[2]-q[1])^2+(p[2]-p[1])^2 )*( (q[2]+q[1])^2+(p[2]+p[1])^2 )
H2(p,q,t)= (p[2]^4+2*p[2]^2*(q[1]^2-1)+(1+q[1]^2)^2 )*(q[2]^4+2*q[2]^2*(p[1]^2+1)+(p[1]^2-1)^2 )
HAM(p,q,t)=H1(p,q,t)/H2(p,q,t)
PofY_find(YY)=sqrt((h+YY.^2-h*YY.^2 )./(1+h-h*YY.^2))

n_iter=100
Y_start=1e-6
Y_end=1e-2
t_end=1e10

Y=range(Y_start,stop=Y_end,length=n_iter)
P=zeros(n_iter)
T=zeros(n_iter)

for k=1:n_iter
    global
    q0,p0=[zeros(2) for i in 1:2]
    P[k]=PofY_find(Y[k])
    q0[1]=0
    p0[1]=P[k]
    q0[2]=0
    p0[2]=Y[k]
    prob= HamiltonianProblem{true}(HAM, p0, q0, (0., t_end))
    t,A=solve(prob, Vern9(),reltol=1e-13,abstol=1e-20,maxiters=1e15, callback=cb)
    T[k]=t.t[end]
end
println("Max exit time is  ", maximum(T))


##
## Visualization of the Mandelbrot set in Julia
##
## Tomáš Kalvoda, FIT CTU, 2015
##

using Images, Colors, FixedPointNumbers
using ProgressMeter

function iter_to_color(i, maxiter)
    clr = convert(RGB, HSV(360.*i/maxiter,1,1))
    return [clr.r, clr.g, clr.b]
end

function getcolor(z::Complex{Float64}, maxiter, radius)
    i = 1; u = 0 + 0im;
    while i <= maxiter
        u = u^2 + z
        if abs(u) > radius
            return iter_to_color(i,maxiter)
        end
        i += 1
    end
    return [1., 1., 1.]
end

function draw_image(;remin=-3.0, remax=3.0, immin=-3.0, immax=3.0, width=400, height=400, maxiter=30, radius=20.)
    redelta = (remax-remin)/(width-1)
    imdelta = (immax-immin)/(height-1)

    img = colorim(Array(Float64, (width,height,3)))
    img["spatialorder"] = ["x", "y"]

    p = Progress(width, 1)

    for j = 0:width-1
        for k = 0:height-1
            img[j+1, k+1, :] = getcolor(remin + j*redelta + 1im*(immin + k*imdelta), maxiter, radius)
        end
        next!(p)
    end

    return fliplr(img)
end


@doc doc"""
    CholeskySpace{N} <: Manifold

The manifold of lower triangular matrices with positive diagonal and
a metric based on the cholesky decomposition. The formulae for this manifold
are for example summarized in Table 1 of [^Lin2019].

# Constructor

    CholeskySpace(n)

Generate the manifold of $n\times n$ lower triangular matrices with positive diagonal.

[^Lin2019]:
    > Lin, Zenhua: "Riemannian Geometry of Symmetric Positive Definite Matrices via
    > Cholesky Decomposition", arXiv: [1908.09326](https://arxiv.org/abs/1908.09326).
"""
struct CholeskySpace{N} <: Manifold end

CholeskySpace(n::Int) = CholeskySpace{n}()

@doc doc"""
    check_manifold_point(M::CholeskySpace, x; kwargs...)

Check whether the matrix `x` lies on the [`CholeskySpace`](@ref) `M`, i.e.
it's size fits the manifold, it is a lower triangular matrix and has positive
entries on the diagonal.
The tolerance for the tests can be set using the `kwargs...`.
"""
function check_manifold_point(M::CholeskySpace, x; kwargs...)
    if size(x) != representation_size(M)
        return DomainError(
            size(x),
            "The point $(x) does not lie on $(M), since its size is not $(representation_size(M)).",
        )
    end
    if !isapprox(norm(strictlyUpperTriangular(x)), 0.0; kwargs...)
        return DomainError(
            norm(UpperTriangular(x) - Diagonal(x)),
            "The point $(x) does not lie on $(M), since it strictly upper triangular nonzero entries",
        )
    end
    if any(diag(x) .<= 0)
        return DomainError(
            min(diag(x)...),
            "The point $(x) does not lie on $(M), since it hast nonpositive entries on the diagonal",
        )
    end
    return nothing
end

"""
    check_tangent_vector(M::CholeskySpace, x, v; kwargs... )

Checks whether `v` is a tangent vector to `x` on the [`CholeskySpace`](@ref) `M`, i.e.
atfer [`check_manifold_point`](@ref)`(M,x)`, `v` has to be of same dimension as `x`
and a symmetric matrix.
The tolerance for the tests can be set using the `kwargs...`.
"""
function check_tangent_vector(M::CholeskySpace, x, v; kwargs...)
    mpe = check_manifold_point(M, x)
    mpe !== nothing && return mpe
    if size(v) != representation_size(M)
        return DomainError(
            size(v),
            "The vector $(v) is not a tangent to a point on $(M) since its size does not match $(representation_size(M)).",
        )
    end
    if !isapprox(norm(strictlyUpperTriangular(v)), 0.0; kwargs...)
        return DomainError(
            norm(UpperTriangular(v) - Diagonal(v)),
            "The matrix $(v) is not a tangent vector at $(x) (represented as an element of the Lie algebra) since it is not lower triangular.",
        )
    end
    return nothing
end

@doc doc"""
    distance(M::CholeskySpace, x, y)

Compute the Riemannian distance on the [`CholeskySpace`](@ref) `M` between two
matrices `x`, `y` that are lower triangular with positive diagonal. The formula
reads

````math
d_{\mathcal M}(x,y) = \sqrt{\sum_{i>j} (x_{ij}-y_{ij})^2 +
\sum_{j=1}^m (\log x_{jj} - \log y_{jj})^2
}
````
"""
function distance(::CholeskySpace, x, y)
    return sqrt(
        sum((strictlyLowerTriangular(x) - strictlyLowerTriangular(y)) .^ 2) +
        sum((log.(diag(x)) - log.(diag(y))) .^ 2),
    )
end

@doc doc"""
    exp(M::CholeskySpace, x, v)

Compute the exponential map on the [`CholeskySpace`](@ref) `M` eminating from the lower
triangular matrix with positive diagonal `x` towards the lower triangular matrix `v`
The formula reads

````math
\exp_x v = \lfloor x \rfloor + \lfloor v \rfloor + \operatorname{diag}(x)
\operatorname{diag}(x)\exp\bigl( \operatorname{diag}(v)\operatorname{diag}(x)^{-1}\bigr),
````

where $\lfloor x\rfloor$ denotes the strictly lower triangular matrix of $x$ and
$\operatorname{diag}(x)$ the diagonal matrix of $x$
"""
exp(::CholeskySpace, ::Any...)

function exp!(::CholeskySpace, y, x, v)
    y .= (
        strictlyLowerTriangular(x) +
        strictlyLowerTriangular(v) +
        Diagonal(diag(x)) * Diagonal(exp.(diag(v) ./ diag(x)))
    )
    return y
end

@doc doc"""
    inner(M::CholeskySpace, x, v, w)

Compute the inner product on the [`CholeskySpace`](@ref) `M` at the
lower triangular matric with positive diagonal `x` and the two tangent vectors
`v`,`w`, i.e they are both lower triangular matrices with arbitrary diagonal.
The formula reads

````math
    g_{x}(v,w) = \sum_{i>j} v_{ij}w_{ij} + \sum_{j=1}^m v_{ii}w_{ii}x_{ii}^{-2}
````
"""
function inner(::CholeskySpace, x, v, w)
    return (
        sum(strictlyLowerTriangular(v) .* strictlyLowerTriangular(w)) +
        sum(diag(v) .* diag(w) ./ (diag(x) .^ 2))
    )
end

@doc doc"""
    log(M::CholeskySpace, v, x, y)

Compute the logarithmic map on the [`CholeskySpace`](@ref) `M` for the geodesic eminating
from the lower triangular matrix with positive diagonal `x` towards `y`.
The formula reads

````math
\log_x v = \lfloor x \rfloor - \lfloor y \rfloor
+\operatorname{diag}(x)\log\bigl(\operatorname{diag}(y)\operatorname{diag}(x)^{-1}\bigr),
````

where $\lfloor x\rfloor$ denotes the strictly lower triangular matrix of $x$ and
$\operatorname{diag}(x)$ the diagonal matrix of $x$
"""
log(::Cholesky, ::Any...)

function log!(::CholeskySpace, v, x, y)
    return copyto!(
        v,
        strictlyLowerTriangular(y) - strictlyLowerTriangular(x) +
        Diagonal(diag(x) .* log.(diag(y) ./ diag(x))),
    )
end

@doc doc"""
    manifold_dimension(M::CholeskySpace)

Return the manifold dimension for the [`CholeskySpace`](@ref) `M`, i.e. $\frac{N(N+1)}{2}$.
"""
@generated manifold_dimension(::CholeskySpace{N}) where {N} = div(N * (N + 1), 2)

@doc doc"""
    reporesentation_size(M::CholeskySpace)

Return the representation size for the [`CholeskySpace`](@ref)`{N}` `M`, i.e. `(N,N)`.
"""
@generated representation_size(::CholeskySpace{N}) where {N} = (N, N)

# two small helpers for strictly lower and upper triangulars
strictlyLowerTriangular(x) = LowerTriangular(x) - Diagonal(diag(x))

strictlyUpperTriangular(x) = UpperTriangular(x) - Diagonal(diag(x))

@doc doc"""
    vector_transport_to(M::CholeskySpace, x, v, y, ::ParallelTransport)

Parallely transport the tangent vector `v` at `x` along the geodesic to `y`
on to the [`CholeskySpace`](@ref) manifold `M`. The formula reads

````math
\mathcal P_{y\gets x}(v) = \lfloor v \rfloor
+ \operatorname{diag}(y)\operatorname{diag}(x)^{-1}\operatorname{diag}(v),
````

where $\lfloor\cdot\rfloor$ denotes the strictly lower triangular matrix,
and $\operatorname{diag}$ extracts the diagonal matrix.
"""
vector_transport_to(::CholeskySpace, ::Any, ::Any, ::Any, ::ParallelTransport)

function vector_transport_to!(::CholeskySpace, vto, x, v, y, ::ParallelTransport)
    return copyto!(
        vto,
        strictlyLowerTriangular(x) + Diagonal(diag(y) .* diag(v) ./ diag(x)),
    )
end

@doc doc"""
    zero_tangent_vector(M::CholeskySpace, x)

Return the zero tangent vector on the [`CholeskySpace`](@ref) `M` at `x`.
"""
zero_tangent_vector(::CholeskySpace, ::Any...)

zero_tangent_vector!(M::CholeskySpace, v, x) = fill!(v, 0)


using Test

include("../src/simulator.jl")
using .simulator.fresnelltools

@testset "test_fresnelltools" begin

    @testset "FresnellBoundraryNoangle" begin
      # A randomized test verifying the angle-independent fresnell boundrary
      # is equivalent to the angle dependent at 90 degrees.
      for _ in 1:100
        n1 = rand(ComplexF64)
        n2 = rand(ComplexF64)

        independent = FresnellBoundrary(n1, n2)
        te_matrix, tm_matrix, θ = FresnellBoundrary(n1, n2, 0)

        @test independent == te_matrix == tm_matrix
        @test θ == 0
      end
    end

    @testset "FresnellBoundraryNonexisting" begin
      # Testing a variety of angles and refractive indices to check that nonexisting borders do no reflect.
      for _ in 1:100
        θ1 = rand() * π / 2
        n = rand(ComplexF64) * 2
        te, tm, θ2 = FresnellBoundrary(n, n, θ1)

        @test θ1 ≈ θ2
        @test round.(te, digits=5) ≈ round.(tm, digits=5) ≈ [1.0 0.0; 0.0 1.0]
      end
    end

    @testset "FresnellSlabTest" begin
      for _ in 1:100
        length = rand() * 10
        slab = FresnellSlab(1, 1, length, 0)

        @test slab ≈ [ℯ^(- length * 1im) 0; 0 ℯ^(- length * 1im)]
      end

      @test FresnellSlab(1, 2, 0, 0) == [1 0; 0 1]
    end

end


@testset "MortalityTables from File" begin

    @testset "TableMetaData" begin
        @testset "whitespace management" begin
            # "2017 Loaded CSO Smoker Distinct Nonsmoker Female ALB" comes with a trailing whitespace
            @test ~isnothing(MortalityTables.table_source_map["2017 Loaded CSO Smoker Distinct Nonsmoker Female ALB"])
        end
    end

    @testset "SOA tables" begin

        @testset "Ultimate Only SOA table" begin
            cso1980 = MortalityTables.table("1980 CSO Basic Table – Male, ANB")
            @test cso1980.ultimate[35] ≈ 0.00118
            @test cso1980[35] ≈ 0.00118
            @test cso1980[35:36] ≈ [0.00118, 0.00128]
            @test cso1980[35 + 60] ≈ 0.27302
            @test cso1980[95] ≈ 0.27302
            @test_throws BoundsError cso1980[125]
            @test_throws ArgumentError survival(cso1980,10,15)
            @test_throws ArgumentError decrement(cso1980,10,15)
            @test omega(cso1980.ultimate) == 100
            @test MortalityTables.ω(cso1980) == 100
        end

        @testset "2001 VBT" begin
            vbt2001 =
                MortalityTables.table("2001 VBT Residual Standard Select and Ultimate - Male Nonsmoker, ANB")

            @test vbt2001.select[35][35] ≈ 0.00036
            @test vbt2001.select[35][35 + 60] ≈ 0.24298
            @test vbt2001.select[95][95] ≈ 0.23815
            @test vbt2001.ultimate[95] ≈ 0.24298
            @test vbt2001.ultimate[120] ≈ 1.0
            @test vbt2001.select[85][120] ≈ 1.0
            @test_throws BoundsError vbt2001.select[150]
            @test omega(vbt2001.select[20]) == 120
            @test omega(vbt2001.ultimate) == 120

            # this tests when there are ultimate rates available
            # but no select rates for the issue age
            @test_throws BoundsError vbt2001.select[116][116] ≈ 0.79988

        end

        @testset "2001 CSO" begin
            cso2001 =
                MortalityTables.table("2001 CSO Super Preferred Select and Ultimate - Male Nonsmoker, ANB")

            @test cso2001.select[35][35] ≈ 0.00037
            @test cso2001.select[35][35 + 60] ≈ 0.26719
            @test cso2001.ultimate[16] ≈ 0.00041
            @test cso2001.ultimate[95] ≈ 0.26719
            @test cso2001.ultimate[120] ≈ 1.0
            @test cso2001.select[85][120] ≈ 1.0
            @test ismissing(cso2001.select[15][15]) # age before table defines rates
            @test_throws BoundsError cso2001.select[150]
            @test_throws BoundsError cso2001.select[70][150]
            @test omega(cso2001.select[16]) == 120
            @test omega(cso2001.select[98]) == 120 # issue #71
            @test omega(cso2001.ultimate) == 120

            @testset "duration ranges" begin
                # select only
                @test all(cso2001.select[29][29:31] .== [0.00029, 0.00035, 0.0004])
                @test all(cso2001.select[35][35:59] .== [
                    0.00037,
                    0.00043,
                    0.00049,
                    0.00057,
                    0.00063,
                    0.0007,
                    0.00077,
                    0.00084,
                    0.00092,
                    0.00101,
                    0.00114,
                    0.00127,
                    0.00143,
                    0.00159,
                    0.00174,
                    0.00188,
                    0.00208,
                    0.00231,
                    0.00251,
                    0.00279,
                    0.00315,
                    0.00357,
                    0.00406,
                    0.00461,
                    0.00508,
                ])

                # crosses into ultimate
                @test all(cso2001.select[35][55:64] .== [
                    0.00315,
                    0.00357,
                    0.00406,
                    0.00461,
                    0.00508,
                    0.00621,
                    0.0069,
                    0.00773,
                    0.00867,
                    0.00965,
                ])
                @test all(cso2001.select[35][35:64] .== [
                    0.00037,
                    0.00043,
                    0.00049,
                    0.00057,
                    0.00063,
                    0.0007,
                    0.00077,
                    0.00084,
                    0.00092,
                    0.00101,
                    0.00114,
                    0.00127,
                    0.00143,
                    0.00159,
                    0.00174,
                    0.00188,
                    0.00208,
                    0.00231,
                    0.00251,
                    0.00279,
                    0.00315,
                    0.00357,
                    0.00406,
                    0.00461,
                    0.00508,
                    0.00621,
                    0.0069,
                    0.00773,
                    0.00867,
                    0.00965,
                ])
            end

        end

    end
end


function white_box_jsr(A,d=2)
    s = discreteswitchedsystem(A)
    optimizer_constructor = optimizer_with_attributes(Mosek.Optimizer, MOI.Silent() => true)
    soslyapb(s, d, optimizer_constructor=optimizer_constructor, tol=1e-4, verbose=0)
    seq = sosbuildsequence(s, d, p_0=:Primal)
    psw = findsmp(seq)
    return psw.growthrate
end
#=
function white_box_stabilization_quad(A,B,tol=1e-4)
    normAmax = 0
    dim,dimIn = size(B)
    for Ai in A
        normAi = opnorm(Ai)
        if normAi>normAmax
            normAmax = normAi
        end
    end
    ub = normAmax
    lb = 0
    K = zeros(dimIn,dim)
    P = Matrix(I,dim,dim)
    while ub-lb > tol
        gamma = (ub + lb)/2
        solver = optimizer_with_attributes(Mosek.Optimizer, MOI.Silent() => true)
        model = Model(solver)
        @variable(model, Q[1:dim, 1:dim] in PSDCone())
        @variable(model, Y[1:dimIn,1:dim])
        @variable(model, s>=0)
        @SDconstraint(model, Q >= Matrix(I,dim,dim))
        @objective(model, Min, s)
        for Ai in A
            @SDconstraint(model, [gamma^2*Q+s*Matrix(I,dim,dim) Q*Ai'+(B*Y)';Ai*Q+B*Y Q+s*Matrix(I,dim,dim)] >= 0)
        end
        JuMP.optimize!(model)
        if value.(s) < 1e-10
            ub = gamma
            P = inv(value.(Q))
            K = value.(Y)*P
        else
            lb = gamma
        end
    end
    return ub, K, P
end
=#

function white_box_stabilization_quad(A,B,tol=1e-4)
    normAmax = 0
    dim,dimIn = size(B)
    for Ai in A
        normAi = opnorm(Ai)
        if normAi>normAmax
            normAmax = normAi
        end
    end
    ub = normAmax
    lb = 0
    K = zeros(dimIn,dim)
    P = Matrix(I,dim,dim)
    while ub-lb > tol
        gamma = (ub + lb)/2
        solver = optimizer_with_attributes(Mosek.Optimizer, MOI.Silent() => true)
        model = Model(solver)
        @variable(model, Q[1:dim, 1:dim] in PSDCone())
        @variable(model, Y[1:dimIn,1:dim])
        @SDconstraint(model, Q >= Matrix(I,dim,dim))
        @objective(model, Max, 0)
        for Ai in A
            @SDconstraint(model, [gamma^2*Q Q*Ai'+(B*Y)';Ai*Q+B*Y Q] >= 0)
        end
        JuMP.optimize!(model)
        if termination_status(model) == MOI.OPTIMAL
            ub = gamma
            P = inv(value.(Q))
            K = value.(Y)*P
        else
            lb = gamma
        end
    end
    return ub, K, P
end

function white_box_LQR(A,B,Q,R)
    (dim,dimIn) = size(B)
    lower_triangular(P) = [P[i, j] for i = 1:size(P, 1) for j = 1:i]
    solver = optimizer_with_attributes(Mosek.Optimizer, MOI.Silent() => true)
    model = Model(solver)
    @variable(model, S[1:dim, 1:dim] in PSDCone())
    @variable(model, Y[1:dimIn,1:dim])
    @variable(model, t)
    @SDconstraint(model, S >= 0)
    @constraint(model, [t; 1; lower_triangular(S)] in MOI.LogDetConeTriangle(dim))
    @objective(model, Max, t)
    for Ai in A
        @SDconstraint(model, [S S*Ai'+(B*Y)' S Y';Ai*S+B*Y S zeros(dim,dim) zeros(dim,dimIn);S zeros(dim,dim) inv(Q) zeros(dim,dimIn);Y zeros(dimIn,dim) zeros(dimIn,dim) inv(R)] >= 0)
    end
    JuMP.optimize!(model)
    if termination_status(model) == MOI.OPTIMAL
        P = inv(value.(S))
        K = value.(Y)*P
        return K, P
    else
        println("The LQR problem is infeasible!")
        return zeros(dimIn,dim), zeros(dim,dim)
    end
end


module Rank2BinaryForms

using Combinatorics:combinations
using DocStringExtensions:SIGNATURES

# Use at least Nemo 0.24.0
# Install with
# ] add https://github.com/Nemocas/Nemo.jl.git
using Nemo

export crossratio, affineGamma, Image_affineGamma, set_getting_rootsofunity, buildpolynomial

#################################
#################################
# Roots of unity
#################################
#################################

rootsofunity(d) = k -> Nemo.root_of_unity(QQBar, d, k)

const ROOTS_OF_UNITY = Ref{Function}(rootsofunity)

"""
$(SIGNATURES)

Sets the method used to compute the `d`-th roots of unity. The input is a function `f` computing the roots of unity accepting calls as `f(d)(k)` and returning the `k`-th root of the `d`-th roots of unity. The order of the `d`-th roots of unity can be clockwise or anti-clockwise, but they must be ordered.
Examples of valid inputs are the following functions `f` and `g`:

```julia
    g(d) = k -> exp(2*pi*im*k/d)
    using Nemo
    CC = ComplexField(64);
    f(d) = k -> Nemo.root_of_unity(CC, d)^k;
```
"""
function set_getting_rootsofunity(f)
    ROOTS_OF_UNITY[] = f
end

## Other ways to get the roots of unity
# CC = ComplexField(64)
# g(d) = k -> Nemo.root_of_unity(CC, d)^k
# set_getting_rootsofunity(g)
#
# f(d) = k -> exp(2*pi*im*k/d)
# set_getting_rootsofunity(f)

#################################
#################################
# Cross ratio
#################################
#################################
"""
$(SIGNATURES)

Returns the cross ratio of four numbers.
There is also the method:

    crossratio(v) = crossratio(v...)
"""
crossratio(A, B, C, D) = (A - B) * (C - D) * inv((A - D) * (C - B))
crossratio(v) = crossratio(v...)

#################################
#################################
# Triplets and 4-tuples
#################################
#################################

"""
$(SIGNATURES)

Returns a vector of triplets `[i,j,k]` of integers `1 <= i < j < k <= d`, up to the action of the `d`-th dihedral group.
"""
triplets(d) = map(pair -> pushfirst!(pair, 1), combinations(2:d, 2))

"""
$(SIGNATURES)

Returns a vector of vectors resulting to `pushfirst` each integer `1:d` in `triplet`.
"""
tuples(d) = triplet -> [[i, triplet...] for i in 1:d if !(i in triplet)]
# Add zero and one as roots
# tuples(d) = triplet -> [[i, triplet...] for i in 1:d if !(i == last(triplet))]
# Add all the roots, if type has an infinity value.
# tuples(d) = triplet -> [[i, triplet...] for i in 1:d]

#################################
#################################
# Define Gamma
#################################
#################################
# Computing the i-th root of the polynomial image of Gamma
# @noinlne not needed
ithroot(d) = tuple -> crossratio(ROOTS_OF_UNITY[](d).(tuple))

# function affineGamma(d, triplet)
#     out = -1 .* ithroot(d).(tuples(d)(triplet))
#     return  [zero(first(out)), one(first(out)), out...]
# end
# affineGamma(d) = triplet -> affineGamma(d, triplet)

"""
$(SIGNATURES)

Returns a function that computes the unknown roots of the polynomial `Γ(triplet)`.
"""
affineGamma(d) = triplet -> ithroot(d).(tuples(d)(triplet))

"""
$(SIGNATURES)

Returns the roots of all the polynomials in the image of `Γ`.
"""
Image_affineGamma(d) = affineGamma(d).(triplets(d))

#################################
#################################
# Building the polynomials
#################################
#################################
"""
$(SIGNATURES)

Generic function to build the polynomials in the image of `Γ` from the list of roots.
"""
function buildpolynomial(roots, vars)
    x = first(vars)
    y = last(vars)
    out = x * y * (x + y)
    for r in roots
        out *= x - r * y
    end
    return out
end

function buildpolynomial(roots, vars::Vector{Symbol})
    x = first(vars)
    y = last(vars)
    out = Expr(:call, :*, x, y, Expr(:call, :+, x, y))
    for r in roots
        out = Expr(:call, :*, out, Expr(:call, :-, x, Expr(:call, :*, r, y)))
    end
    return out
end

buildpolynomial(vars) = roots -> buildpolynomial(roots, vars)

end


# TODO: handle multidimensional case

"""
    product_integral(f, g, [, solver])

Compute the integral of `f .* g` over [0,1].
"""
function product_integral(f, g, solver=HCubatureJL())
    function fun(u,_)
        f(u) .* g(u)
    end
    prob = QuadratureProblem{false}(fun, 0, 1, ())
    solve(prob, solver).u
end

"""
    kernel_integral(k, f, t [, solver])

Compute the kernel integral `∫ k(t,s)f(s)ds`.
"""
function kernel_integral(k, f, t, solver=HCubatureJL())
    fun(u) = k(t,u)
    return product_integral(fun, f, solver)
end

"""
    kernel_transpose_integral(k, f, t [, solver])

Compute the kernel integral `∫ k(s,t)f(s)ds`.
"""
function kernel_transpose_integral(k, f, t, solver=HCubatureJL())
    fun(u) = k(u,t)
    return product_integral(fun, f, solver)
end

"""
    kernel_double_integral(k, f, g, [, solver])

Compute the kernel integral `∫∫ k(t,s)f(t)g(s) ds dt`.
"""
function kernel_double_integral(k, f, g, solver=HCubatureJL())
    function fun(u,_)
        f(u[1]) .* k(u[1], u[2]) .* g(u[2])
    end
    prob = QuadratureProblem{false}(fun, [0,0], [1,1], ())
    solve(prob, solver).u
end


"""
    CoamoebaFiber3D(f::AbstractPolynomial, θ=(0., 0., 0.))

Construct the fiber `CoamoebaFiber3D` of the trivariate polynomial `f` at `θ`.
"""
function CoamoebaFiber3D(p::MP.AbstractPolynomialLike{S}, θ=(0.0, 0.0, 0.0)) where S
    vars = MP.variables(p)
    @assert (length(vars) == 3) "Expected a trivariate polynomial, but got $(p)."

    T = realified_coefficient_type(S)
    f_re, f_im = SP.Polynomial.(one(T) .* realify(p))

    sincosθ = _sincosθ3(θ)
    U = zeros(2, 6)
    x = SVector(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
    CoamoebaFiber3D(f_re, f_im, sincosθ, U, x)
end

_sincosθ3(θ) = (sin(θ[1]), sin(θ[2]), sin(θ[3]), cos(θ[1]), cos(θ[2]), cos(θ[3]))

function update_fiber!(F::CoamoebaFiber3D, θ)
    F.sincosθ = _sincosθ3(θ)
    F
end

function phi!(fiber::CoamoebaFiber3D, w)
    s1, s2, s3, c1, c2, c3 = fiber.sincosθ
    v1, v2, v3 = exp.(w)
    x = SVector(v1 * c1, v2 * c2, v3 * c3, v1 * s1, v2 * s2, v3 * s3)
    fiber.x = x
    x
end

function evaluate(fiber::CoamoebaFiber3D, w)
    x = phi!(fiber, w)
    SVector(SP.evaluate(fiber.f_re, x), SP.evaluate(fiber.f_im, x))
end

function jacobian(fiber::CoamoebaFiber3D, w, precomputed=false)
    x = precomputed ? fiber.x : phi!(fiber, w)
    U = fiber.U
    @inbounds begin
        ∇f_re = SP.gradient(fiber.f_re, x)
        for i=1:6
            U[1, i] = ∇f_re[i]
        end
        U[2, 4] = U[1, 1]
        U[2, 5] = U[1, 2]
        U[2, 6] = U[1, 3]
        U[2, 1] = -U[1, 4]
        U[2, 2] = -U[1, 5]
        U[2, 3] = -U[1, 6]
    end

    @inbounds out = begin
        ∂re∂w1 = muladd(U[1, 1], x[1], U[1, 4] * x[4])
        ∂re∂w2 = muladd(U[1, 2], x[2], U[1, 5] * x[5])
        ∂re∂w3 = muladd(U[1, 3], x[3], U[1, 6] * x[6])
        ∂im∂w1 = muladd(U[2, 1], x[1], U[2, 4] * x[4])
        ∂im∂w2 = muladd(U[2, 2], x[2], U[2, 5] * x[5])
        ∂im∂w3 = muladd(U[2, 3], x[3], U[2, 6] * x[6])

        SMatrix{2, 3}(∂re∂w1, ∂im∂w1, ∂re∂w2, ∂im∂w2, ∂re∂w3, ∂im∂w3)
    end
    out
end


function SeisPWD(in;w1=10,w2=10,dz_in=1,dx_in=1,format="angle")
	# dip estimation by Plane Wave Destruction
	# see http://sepwww.stanford.edu/data/media/public/sep//prof/pvi.pdf Chapter 4.

	d=0
	d = copy(in)
	n1,n2 = size(in)
	#format = get(param,"format","angle") # output format: "angle" (in degrees wrt vertical) or "dip" (in y samples over x samples)
	pp = zeros(n1,n2); dx = wavekill(1.,pp,n1,n2,d)
	pp =  ones(n1,n2); dt = wavekill(0.,pp,n1,n2,d)
	dtdx = dt.*dx
	dxdx = dx.*dx
	dtdt = dt.*dt
	for i2 = 1 : n2
		dtdt[:,i2] = triangle(w1,n1,dtdt[:,i2])
		dxdx[:,i2] = triangle(w1,n1,dxdx[:,i2])
		dtdx[:,i2] = triangle(w1,n1,dtdx[:,i2])
	end
	for i1 = 1 : n1
		dtdt[i1,:] = triangle(w2,n2,dtdt[i1,:])
		dxdx[i1,:] = triangle(w2,n2,dxdx[i1,:])
		dtdx[i1,:] = triangle(w2,n2,dtdx[i1,:])
	end

	coh = sqrt.((dtdx.*dtdx)./(dtdt.*dxdx))
	pp = -dtdx./dtdt

	coh = zeros(n1,n2);
	for i1 = 1 : n1
		for i2 = 1 : n2
			if (abs(dtdt[i1,i2]) > 1e-6)
				coh[i1,i2] = sqrt((dtdx[i1,i2].*dtdx[i1,i2])./(dtdt[i1,i2].*dxdx[i1,i2]))
				pp[i1,i2] = -dtdx[i1,i2]./dtdt[i1,i2]
			else
				coh[i1,i2] = 0.
				pp[i1,i2]  = 0.
			end
		end
	end
	for i2 = 1 : n2
		pp[:,i2] = triangle(w1,n1,pp[:,i2])
	end
	for i1 = 1 : n1
		pp[i1,:] = triangle(w2,n2,pp[i1,:])
	end

	res = wavekill(1.,pp,n1,n2,d)

	if (format == "angle")
		pp = atan.(pp*dz_in/dx_in)*180/pi;
	end

	return coh,pp,res
end

function wavekill(aa,bb,n1,n2,uu)
	vv = zeros(n1,n2)
	s11 = -aa-bb; s12 = aa-bb;
	s21 = -aa+bb; s22 = aa+bb;
	for i2 = 1 : n2-1
		for i1 = 1 : n1-1
			vv[i1,i2] = uu[i1  ,i2  ]*s11[i1,i2] + uu[i1  ,i2+1]*s12[i1,i2] + uu[i1+1,i2  ]*s21[i1,i2] + uu[i1+1,i2+1]*s22[i1,i2];
		end
	end
	vv[n1,:] = vv[n1-1,:]
	vv[:,n2] = vv[:,n2-1]
	return vv
end

function triangle(nbox,nd,xx)
	yy = zeros(nd)
	pp = boxconv(nbox,nd,xx)
	np = nbox+nd-1
	qq = boxconv(nbox,np,pp)
	nq = nbox+np-1
	for i = 1 : nd
		yy[i] = qq[i+nbox-1]
	end
	for i = 1 : nbox - 1
		yy[i] = yy[i] + qq[nbox-i]
	end
	for i = 1 : nbox - 1
		yy[nd-i+1] = yy[nd-i+1] + qq[nd+(nbox-1)+i]
	end
	return yy
end

function boxconv(nbox,nx,xx)
	ny = nx+nbox-1
	yy = zeros(ny)
	bb = zeros(nx+nbox)
	bb[1] = xx[1]
	for i = 2 : nx
		bb[i] = bb[i-1] + xx[i]
	end
	for i = nx+1 : ny
		bb[i] = bb[i-1]
	end
	for i = 1 : nbox
		yy[i] = bb[i]
	end
	for i = nbox+1 : ny
		yy[i] = bb[i] - bb[i - nbox]
	end
	for i = 1 : ny
		yy[i] = yy[i]/nbox
	end
	return yy
end


push!(LOAD_PATH, string(pwd(),"/src"))

using ModifiedStokesSolver
import AnalyticDomains
using Base.Test
using PyPlot
import IterativeSolvers

@testset "MultiplyConnected" begin

    # Setup discretization
    R1 = 1.0
    R2 = 0.25
    alpha = 3.0
    panelorder = 16
    numpanels1 = 40
    numpanels2 = 8
    tol = 1e-8

    # Setup problem
    xsrc1 = [0.01, 0.02]
    xsrc2 = [1.0, 1.0]
    srand(1)
    qsrc = 1-2*rand(2)
    nsrc = 1-2*rand(2)
    fsrc = 1-2*rand(2)
    ubc(x) = stresslet(x-xsrc1, qsrc, nsrc, alpha) +
        stokeslet(x-xsrc1, fsrc, alpha) +
        stresslet(x-xsrc2, qsrc, nsrc, alpha) +
        stokeslet(x-xsrc2, fsrc, alpha)

    # Discretize
    curve1 = AnalyticDomains.starfish(;amplitude = 0.1, radius = R1)
    curve2 = AnalyticDomains.starfish(;amplitude = 0.0, radius = R2, interior=true)
    dcurve = discretize([curve1, curve2], [numpanels1, numpanels2], panelorder)
    N = dcurve.numpoints

    # Right hand side
    rhs = zeros(2*N)
    for i=1:dcurve.numpoints
        rhs[2(i-1)+(1:2)] = ubc(dcurve.points[:,i])
    end

    # Dense matrix
    println("* Matrix assembly and solve")
    @time LHS = system_matrix(dcurve, alpha)
    @time sol = LHS\rhs
    density = reshape(sol, 2, N)

    # Interior test point
    zt = [(R1+R2)/2, 0.01]
    ref = ubc(zt)
    dlp = doublelayer(dcurve, density, zt, alpha)
    relerr = maximum(abs.(ref.-dlp)) / maximum(abs.(ref))
    println("* Interior eval")
    @show relerr
    @test relerr < 1e-15

    # Field test
    Ngrid = 50
    xmax = maximum(dcurve.points[1,:])
    xmin = minimum(dcurve.points[1,:])
    ymax = maximum(dcurve.points[2,:])
    ymin = minimum(dcurve.points[2,:])
    x = linspace(xmin, xmax, Ngrid)
    y = linspace(ymin, ymax, Ngrid)
    X, Y = ndgrid(x, y)
    zt = [vec(X) vec(Y)]'
    interior = interior_points(dcurve, zt)
    numeval = sum(interior)
    zt_interior = zt[:, interior]
    @time u = doublelayer_fast(dcurve, density, zt_interior, alpha, specquad=true)
    uref = Array{Float64}(2, numeval)
    for i=1:numeval
        uref[:,i] = ubc(zt_interior[:, i])
    end
    unorm = norm(vec(uref), Inf)
    # Error
    err = u .- uref
    maxerr = Array{Float64}(numeval)
    for i=1:numeval
        maxerr[i] = norm(err[:, i], Inf)
    end

    max_relerr_grid = norm(vec(maxerr), Inf) / unorm
    @show max_relerr_grid
    @test max_relerr_grid < 1e-13

    # Plot
    E = zeros(X)
    E[interior] = maxerr
    clf()
    pcolor(X, Y, log10.(E))
    colorbar()
end


@testset "TriMesh Metrics" begin

    _mesh = load_trimesh(
        joinpath(@__DIR__, "assets/teapot.obj"),
        joinpath(@__DIR__, "assets/sphere.obj"),
    )

    @testset "Laplacian loss" begin

        verts1 = [
            0.1 0.3 0.5
            0.5 0.2 0.1
            0.6 0.8 0.7
        ]
        verts2 = [
            0.1 0.3 0.3
            0.6 0.7 0.8
            0.2 0.3 0.4
            0.1 0.5 0.3
        ]
        verts3 = [
            0.7 0.3 0.6
            0.2 0.4 0.8
            0.9 0.5 0.2
            0.2 0.3 0.4
            0.9 0.3 0.8
        ]

        faces1 = [1 2 3]
        faces2 = [
            1 2 3
            2 3 4
        ]
        faces3 = [
            2 3 1
            1 2 4
            3 4 2
            5 4 3
            5 1 2
            5 4 2
            5 3 2
        ]

        T, R = Float32, UInt32
        verts_list = [T.(verts1'), T.(verts2'), T.(verts3')]
        faces_list = [R.(faces1'), R.(faces2'), R.(faces3')]
        m = TriMesh(verts_list, faces_list)

        _edges = get_edges_packed(m)
        V = size(get_verts_packed(m), 2)
        L = zeros(V, V)
        deg = zeros(V, V)
        for i = 1:size(_edges, 1)
            L[_edges[i, 1], _edges[i, 2]] = 1
            L[_edges[i, 2], _edges[i, 1]] = 1
        end
        deg = sum(L; dims = 2)
        deg = map(x -> (x > 0 ? 1 / x : x), deg)
        for i = 1:V, j = 1:V
            if i == j
                L[i, j] = -1
            elseif L[i, j] == 1
                L[i, j] = deg[i]
            end
        end

        verts = get_verts_packed(m)
        L = L * transpose(verts)
        L = Flux3D._norm(L; dims = 2)
        @test isapprox(mean(L), laplacian_loss(m))
        @test gradient(x -> laplacian_loss(x), m) isa Tuple
    end

    @testset "Edge Loss" begin
        m = deepcopy(_mesh)
        verts = get_verts_packed(m)
        edges = get_edges_packed(m)
        v1 = verts[:, edges[:, 1]]
        v2 = verts[:, edges[:, 2]]
        loss = mean((Flux3D._norm(v1 - v2; dims = 1)) .^ 2)
        @test edge_loss(m) == loss
        @test gradient(x -> edge_loss(x), m) isa Tuple
    end

    @testset "Chamfer Loss" begin
        m = deepcopy(_mesh)
        loss = chamfer_distance(m, m)
        @test all(isapprox.(loss, 0, rtol = 1e-5, atol = 1e-2))
        @test gradient(x -> chamfer_distance(x, x), m) isa Tuple

        # naive chamfer distance

        function naive_chamfer(x, y)
            xx = sum(x .^ 2, dims = 1)
            yy = sum(y .^ 2, dims = 1)
            zz = Flux3D.Flux.batched_mul(permutedims(x, (2, 1, 3)), y)
            rx = reshape(xx, size(xx, 2), 1, :)
            ry = reshape(yy, 1, size(yy, 2), :)
            P = (rx .+ ry) .- (2 .* zz)
            nn_for_x = minimum(P; dims = 2)
            nn_for_y = minimum(P; dims = 1)
            nn_for_x = mean(nn_for_x; dims = 2)
            nn_for_y = mean(nn_for_y; dims = 2)
            loss = mean(nn_for_x) + mean(nn_for_y)
            return loss
        end

        x = rand(Float32, 3, 1000, 2)
        y = rand(Float32, 3, 500, 2)
        @test isapprox(chamfer_distance(x, y), naive_chamfer(x, y))
        grad1 = gradient(x -> chamfer_distance(x, y), x)[1]
        grad2 = gradient(x -> naive_chamfer(x, y), x)[1]
        @test isapprox(grad1, grad2, atol = 1e-2, rtol = 1e-3)
    end
end


function ichampf(t, fft_u, fft_v, A1, A2, Ntaumm, matr, conjmatr, Ktaubis, llambda, sigma)

    u = ifft(exp(1im * t * A1) * fft_u) * matr
    v = ifft(exp(1im * t * A1) * fft_v) * matr

    z = (u + conj(v)) / 2

    fz1 = abs(z) .^ (2 * sigma) * z
    champu = -1im * llambda * A2 * exp(-1im * t * A1) * fft(conjmatr * fz1)
    champv = -1im * llambda * A2 * exp(-1im * t * A1) * fft(conjmatr * conj(fz1))

    champu_fft = fft(champu, dims=1)
    champv_fft = fft(champv, dims=1)

    champumoy = champu_fft[0, :] / Ntaumm
    champvmoy = champv_fft[0, :] / Ntaumm

    champu_fft[0, :] = 0 * champu_fft[0, :]
    champv_fft[0, :] = 0 * champv_fft[0, :]
    champu = ifft(champu_fft / (1im * Ktaubis), dims=1)
    champv = ifft(champv_fft / (1im * Ktaubis), dims=1)

    return champu, champv, champumoy, champvmoy

end