# This file is part of the TaylorSeries.jl Julia package, MIT license
#
# Luis Benet & David P. Sanders
# UNAM
#
# MIT Expat license
#

# Functions
for T in (:Taylor1, :TaylorN)
    @eval begin
        function exp(a::$T)
            order = a.order
            aux = exp(constant_term(a))
            aa = one(aux) * a
            c = $T( aux, order )
            for k in eachindex(a)
                exp!(c, aa, k)
            end
            return c
        end

        function log(a::$T)
            iszero(constant_term(a)) && throw(DomainError(a,
                    """The 0-th order coefficient must be non-zero in order to expand `log` around 0."""))

            order = a.order
            aux = log(constant_term(a))
            aa = one(aux) * a
            c = $T( aux, order )
            for k in eachindex(a)
                log!(c, aa, k)
            end
            return c
        end

        sin(a::$T) = sincos(a)[1]
        cos(a::$T) = sincos(a)[2]
        function sincos(a::$T)
            order = a.order
            aux = sin(constant_term(a))
            aa = one(aux) * a
            s = $T( aux, order )
            c = $T( cos(constant_term(a)), order )
            for k in eachindex(a)
                sincos!(s, c, aa, k)
            end
            return s, c
        end

        function tan(a::$T)
            order = a.order
            aux = tan(constant_term(a))
            aa = one(aux) * a
            c = $T(aux, order)
            c2 = $T(aux^2, order)
            for k in eachindex(a)
                tan!(c, aa, c2, k)
            end
            return c
        end

        function asin(a::$T)
            a0 = constant_term(a)
            a0^2 == one(a0) && throw(DomainError(a,
                """Series expansion of asin(x) diverges at x = ±1."""))

            order = a.order
            aux = asin(a0)
            aa = one(aux) * a
            c = $T( aux, order )
            r = $T( sqrt(1 - a0^2), order )
            for k in eachindex(a)
                asin!(c, aa, r, k)
            end
            return c
        end

        function acos(a::$T)
            a0 = constant_term(a)
            a0^2 == one(a0) && throw(DomainError(a,
            """Series expansion of asin(x) diverges at x = ±1."""))

            order = a.order
            aux = acos(a0)
            aa = one(aux) * a
            c = $T( aux, order )
            r = $T( sqrt(1 - a0^2), order )
            for k in eachindex(a)
                acos!(c, aa, r, k)
            end
            return c
        end

        function atan(a::$T)
            order = a.order
            a0 = constant_term(a)
            aux = atan(a0)
            aa = one(aux) * a
            c = $T( aux, order)
            r = $T(one(aux) + a0^2, order)
            iszero(constant_term(r)) && throw(DomainError(a,
                """Series expansion of atan(x) diverges at x = ±im."""))

            for k in eachindex(a)
                atan!(c, aa, r, k)
            end
            return c
        end

        function atan(a::$T, b::$T)
            c = atan(a/b)
            c[0] = atan(constant_term(a), constant_term(b))
            return c
        end

        sinh(a::$T) = sinhcosh(a)[1]
        cosh(a::$T) = sinhcosh(a)[2]
        function sinhcosh(a::$T)
            order = a.order
            aux = sinh(constant_term(a))
            aa = one(aux) * a
            s = $T( aux, order)
            c = $T( cosh(constant_term(a)), order)
            for k in eachindex(a)
                sinhcosh!(s, c, aa, k)
            end
            return s, c
        end

        function tanh(a::$T)
            order = a.order
            aux = tanh( constant_term(a) )
            aa = one(aux) * a
            c = $T( aux, order)
            c2 = $T( aux^2, order)
            for k in eachindex(a)
                tanh!(c, aa, c2, k)
            end
            return c
        end

        function asinh(a::$T)
            order = a.order
            a0 = constant_term(a)
            aux = asinh(a0)
            aa = one(aux) * a
            c = $T( aux, order )
            r = $T( sqrt(a0^2 + 1), order )
            iszero(constant_term(r)) && throw(DomainError(a,
                """Series expansion of asinh(x) diverges at x = ±im."""))
            for k in eachindex(a)
                asinh!(c, aa, r, k)
            end
            return c
        end

        function acosh(a::$T)
            a0 = constant_term(a)
            a0^2 == one(a0) && throw(DomainError(a,
                """Series expansion of acosh(x) diverges at x = ±1."""))

            order = a.order
            aux = acosh(a0)
            aa = one(aux) * a
            c = $T( aux, order )
            r = $T( sqrt(a0^2 - 1), order )
            for k in eachindex(a)
                acosh!(c, aa, r, k)
            end
            return c
        end

        function atanh(a::$T)
            order = a.order
            a0 = constant_term(a)
            aux = atanh(a0)
            aa = one(aux) * a
            c = $T( aux, order)
            r = $T(one(aux) - a0^2, order)
            iszero(constant_term(r)) && throw(DomainError(a,
                """Series expansion of atanh(x) diverges at x = ±1."""))

            for k in eachindex(a)
                atanh!(c, aa, r, k)
            end
            return c
        end

    end
end

# Recursive functions (homogeneous coefficients)
for T in (:Taylor1, :TaylorN)
    @eval begin
        @inline function identity!(c::$T{T}, a::$T{T}, k::Int) where {T<:Number}
            @inbounds c[k] = identity(a[k])
            return nothing
        end

        @inline function zero!(c::$T{T}, a::$T{T}, k::Int) where {T<:Number}
            @inbounds c[k] = zero(a[k])
            return nothing
        end

        @inline function one!(c::$T{T}, a::$T{T}, k::Int) where {T<:Number}
            if k == 0
                @inbounds c[0] = one(a[0])
            else
                @inbounds c[k] = zero(a[k])
            end
            return nothing
        end

        @inline function abs!(c::$T{T}, a::$T{T}, k::Int) where {T<:Number}
            z = zero(constant_term(a))
            if constant_term(constant_term(a)) > constant_term(z)
                return add!(c, a, k)
            elseif constant_term(constant_term(a)) < constant_term(z)
                return subst!(c, a, k)
            else
                throw(DomainError(a,
                """The 0th order coefficient must be non-zero
                (abs(x) is not differentiable at x=0)."""))
            end
            return nothing
        end

        @inline abs2!(c::$T{T}, a::$T{T}, k::Int) where {T<:Number} = sqr!(c, a, k)

        @inline function exp!(c::$T{T}, a::$T{T}, k::Int) where {T<:Number}
            if k == 0
                @inbounds c[0] = exp(constant_term(a))
                return nothing
            end
            if $T == Taylor1
                @inbounds c[k] = k * a[k] * c[0]
            else
                @inbounds mul!(c[k], k * a[k], c[0])
            end
            @inbounds for i = 1:k-1
                if $T == Taylor1
                    c[k] += (k-i) * a[k-i] * c[i]
                else
                    mul!(c[k], (k-i) * a[k-i], c[i])
                end
            end
            @inbounds c[k] = c[k] / k
            return nothing
        end

        @inline function log!(c::$T{T}, a::$T{T}, k::Int) where {T<:Number}
            if k == 0
                @inbounds c[0] = log(constant_term(a))
                return nothing
            elseif k == 1
                @inbounds c[1] = a[1] / constant_term(a)
                return nothing
            end

            if $T == Taylor1
                @inbounds c[k] = (k-1) * a[1] * c[k-1]
            else
                @inbounds mul!(c[k], (k-1)*a[1], c[k-1])
            end
            @inbounds for i = 2:k-1
                if $T == Taylor1
                    c[k] += (k-i) * a[i] * c[k-i]
                else
                    mul!(c[k], (k-i)*a[i], c[k-i])
                end
            end
            @inbounds c[k] = (a[k] - c[k]/k) / constant_term(a)
            return nothing
        end

        @inline function sincos!(s::$T{T}, c::$T{T}, a::$T{T}, k::Int) where {T<:Number}
            if k == 0
                a0 = constant_term(a)
                @inbounds s[0], c[0] = sin( a0 ), cos( a0 )
                return nothing
            end

            x = a[1]
            if $T == Taylor1
                @inbounds s[k] = x * c[k-1]
                @inbounds c[k] = -x * s[k-1]
            else
                mul!(s[k], x, c[k-1])
                mul!(c[k], -x, s[k-1])
            end
            @inbounds for i = 2:k
                x = i * a[i]
                if $T == Taylor1
                    s[k] += x * c[k-i]
                    c[k] -= x * s[k-i]
                else
                    mul!(s[k], x, c[k-i])
                    mul!(c[k], -x, s[k-i])
                end
            end

            @inbounds s[k] = s[k] / k
            @inbounds c[k] = c[k] / k
            return nothing
        end

        @inline function tan!(c::$T{T}, a::$T{T}, c2::$T{T}, k::Int) where {T<:Number}
            if k == 0
                @inbounds aux = tan( constant_term(a) )
                @inbounds c[0] = aux
                @inbounds c2[0] = aux^2
                return nothing
            end

            if $T == Taylor1
                @inbounds c[k] = (k-1) * a[k-1] * c2[1]
            else
                @inbounds mul!(c[k], (k-1) * a[k-1], c2[1])
            end
            @inbounds for i = 1:k-1
                if $T == Taylor1
                    c[k] += (k-i) * a[k-i] * c2[i]
                else
                    mul!(c[k], (k-i) * a[k-i], c2[i])
                end
            end
            @inbounds c[k] = a[k] + c[k]/k
            sqr!(c2, c, k)

            return nothing
        end

        @inline function asin!(c::$T{T}, a::$T{T}, r::$T{T}, k::Int) where {T<:Number}
            if k == 0
                a0 = constant_term(a)
                @inbounds c[0] = asin( a0 )
                @inbounds r[0] = sqrt( 1 - a0^2 )
                return nothing
            end

            if $T == Taylor1
                @inbounds c[k] = (k-1) * r[1] * c[k-1]
            else
                @inbounds mul!(c[k], (k-1) * r[1], c[k-1])
            end
            @inbounds for i in 2:k-1
                if $T == Taylor1
                    c[k] += (k-i) * r[i] * c[k-i]
                else
                    mul!(c[k], (k-i) * r[i], c[k-i])
                end
            end
            sqrt!(r, 1-a^2, k)
            @inbounds c[k] = (a[k] - c[k]/k) / constant_term(r)
            return nothing
        end

        @inline function acos!(c::$T{T}, a::$T{T}, r::$T{T}, k::Int) where {T<:Number}
            if k == 0
                a0 = constant_term(a)
                @inbounds c[0] = acos( a0 )
                @inbounds r[0] = sqrt( 1 - a0^2 )
                return nothing
            end

            if $T == Taylor1
                @inbounds c[k] = (k-1) * r[1] * c[k-1]
            else
                @inbounds mul!(c[k], (k-1) * r[1], c[k-1])
            end
            @inbounds for i in 2:k-1
                if $T == Taylor1
                    c[k] += (k-i) * r[i] * c[k-i]
                else
                    mul!(c[k], (k-i) * r[i], c[k-i])
                end
            end
            sqrt!(r, 1-a^2, k)
            @inbounds c[k] = -(a[k] + c[k]/k) / constant_term(r)
            return nothing
        end

        @inline function atan!(c::$T{T}, a::$T{T}, r::$T{T}, k::Int) where {T<:Number}
            if k == 0
                a0 = constant_term(a)
                @inbounds c[0] = atan( a0 )
                @inbounds r[0] = 1 + a0^2
                return nothing
            end

            if $T == Taylor1
                @inbounds c[k] = (k-1) * r[1] * c[k-1]
            else
                @inbounds mul!(c[k], (k-1) * r[1], c[k-1])
            end
            @inbounds for i in 2:k-1
                if $T == Taylor1
                    c[k] += (k-i) * r[i] * c[k-i]
                else
                    mul!(c[k], (k-i) * r[i], c[k-i])
                end
            end
            @inbounds sqr!(r, a, k)
            @inbounds c[k] = (a[k] - c[k]/k) / constant_term(r)
            return nothing
        end

        @inline function sinhcosh!(s::$T{T}, c::$T{T}, a::$T{T}, k::Int) where {T<:Number}
            if k == 0
                @inbounds s[0] = sinh( constant_term(a) )
                @inbounds c[0] = cosh( constant_term(a) )
                return nothing
            end

            x = a[1]
            if $T == Taylor1
                @inbounds s[k] = x * c[k-1]
                @inbounds c[k] = x * s[k-1]
            else
                @inbounds mul!(s[k], x, c[k-1])
                @inbounds mul!(c[k], x, s[k-1])
            end
            @inbounds for i = 2:k
                x = i * a[i]
                if $T == Taylor1
                    s[k] += x * c[k-i]
                    c[k] += x * s[k-i]
                else
                    mul!(s[k], x, c[k-i])
                    mul!(c[k], x, s[k-i])
                end
            end
            s[k] = s[k] / k
            c[k] = c[k] / k
            return nothing
        end

        @inline function tanh!(c::$T{T}, a::$T{T}, c2::$T{T}, k::Int) where {T<:Number}
            if k == 0
                @inbounds aux = tanh( constant_term(a) )
                @inbounds c[0] = aux
                @inbounds c2[0] = aux^2
                return nothing
            end

            if $T == Taylor1
                @inbounds c[k] = k * a[k] * c2[0]
            else
                @inbounds mul!(c[k], k * a[k], c2[0])
            end
            @inbounds for i = 1:k-1
                if $T == Taylor1
                    c[k] += (k-i) * a[k-i] * c2[i]
                else
                    mul!(c[k], (k-i) * a[k-i], c2[i])
                end
            end
            @inbounds c[k] = a[k] - c[k]/k
            sqr!(c2, c, k)

            return nothing
        end

        @inline function asinh!(c::$T{T}, a::$T{T}, r::$T{T}, k::Int) where {T<:Number}
            if k == 0
                a0 = constant_term(a)
                @inbounds c[0] = asinh( a0 )
                @inbounds r[0] = sqrt( a0^2 + 1 )
                return nothing
            end

            if $T == Taylor1
                @inbounds c[k] = (k-1) * r[1] * c[k-1]
            else
                @inbounds mul!(c[k], (k-1) * r[1], c[k-1])
            end
            @inbounds for i in 2:k-1
                if $T == Taylor1
                    c[k] += (k-i) * r[i] * c[k-i]
                else
                    mul!(c[k], (k-i) * r[i], c[k-i])
                end
            end
            sqrt!(r, a^2+1, k)
            @inbounds c[k] = (a[k] - c[k]/k) / constant_term(r)
            return nothing
        end

        @inline function acosh!(c::$T{T}, a::$T{T}, r::$T{T}, k::Int) where {T<:Number}
            if k == 0
                a0 = constant_term(a)
                @inbounds c[0] = acosh( a0 )
                @inbounds r[0] = sqrt( a0^2 - 1 )
                return nothing
            end

            if $T == Taylor1
                @inbounds c[k] = (k-1) * r[1] * c[k-1]
            else
                @inbounds mul!(c[k], (k-1) * r[1], c[k-1])
            end
            @inbounds for i in 2:k-1
                if $T == Taylor1
                    c[k] += (k-i) * r[i] * c[k-i]
                else
                    mul!(c[k], (k-i) * r[i], c[k-i])
                end
            end
            sqrt!(r, a^2-1, k)
            @inbounds c[k] = (a[k] - c[k]/k) / constant_term(r)
            return nothing
        end

        @inline function atanh!(c::$T{T}, a::$T{T}, r::$T{T}, k::Int) where {T<:Number}
            if k == 0
                a0 = constant_term(a)
                @inbounds c[0] = atanh( a0 )
                @inbounds r[0] = 1 - a0^2
                return nothing
            end

            if $T == Taylor1
                @inbounds c[k] = (k-1) * r[1] * c[k-1]
            else
                @inbounds mul!(c[k], (k-1) * r[1], c[k-1])
            end
            @inbounds for i in 2:k-1
                if $T == Taylor1
                    c[k] += (k-i) * r[i] * c[k-i]
                else
                    mul!(c[k], (k-i) * r[i], c[k-i])
                end
            end
            @inbounds sqr!(r, a, k)
            @inbounds c[k] = (a[k] + c[k]/k) / constant_term(r)
            return nothing
        end

    end
end

@doc doc"""
    inverse(f)

Return the Taylor expansion of ``f^{-1}(t)``, of order `N = f.order`,
for `f::Taylor1` polynomial if the first coefficient of `f` is zero.
Otherwise, a `DomainError` is thrown.

The algorithm implements Lagrange inversion at ``t=0`` if ``f(0)=0``:
```math
\begin{equation*}
f^{-1}(t) = \sum_{n=1}^{N} \frac{t^n}{n!} \left.
    \frac{{\rm d}^{n-1}}{{\rm d} z^{n-1}}\left(\frac{z}{f(z)}\right)^n
    \right\vert_{z=0}.
\end{equation*}
```

""" inverse

function inverse(f::Taylor1{T}) where {T<:Number}
    if f[0] != zero(T)
        throw(DomainError(f,
        """
        Evaluation of Taylor1 series at 0 is non-zero. For high accuracy, revert
        a Taylor1 series with first coefficient 0 and re-expand about f(0).
        """))
    end
    z = Taylor1(T,f.order)
    zdivf = z/f
    zdivfpown = zdivf
    S = eltype(zdivf)
    coeffs = zeros(S,f.order+1)

    @inbounds for n in 1:f.order
        coeffs[n+1] = zdivfpown[n-1]/n
        zdivfpown *= zdivf
    end
    Taylor1(coeffs, f.order)
end

# Documentation for the recursion relations
@doc doc"""
    exp!(c, a, k) --> nothing

Update the `k-th` expansion coefficient `c[k+1]` of `c = exp(a)`
for both `c` and `a` either `Taylor1` or `TaylorN`.

The coefficients are given by

```math
\begin{equation*}
c_k = \frac{1}{k} \sum_{j=0}^{k-1} (k-j) a_{k-j} c_j.
\end{equation*}
```

""" exp!

@doc doc"""
    log!(c, a, k) --> nothing

Update the `k-th` expansion coefficient `c[k+1]` of `c = log(a)`
for both `c` and `a` either `Taylor1` or `TaylorN`.

The coefficients are given by

```math
\begin{equation*}
c_k = \frac{1}{a_0} \big(a_k - \frac{1}{k} \sum_{j=0}^{k-1} j a_{k-j} c_j \big).
\end{equation*}
```

""" log!

@doc doc"""
    sincos!(s, c, a, k) --> nothing

Update the `k-th` expansion coefficients `s[k+1]` and `c[k+1]`
of `s = sin(a)` and `c = cos(a)` simultaneously, for `s`, `c` and `a`
either `Taylor1` or `TaylorN`.

The coefficients are given by

```math
\begin{eqnarray*}
s_k &=&  \frac{1}{k}\sum_{j=0}^{k-1} (k-j) a_{k-j} c_j ,\\
c_k &=& -\frac{1}{k}\sum_{j=0}^{k-1} (k-j) a_{k-j} s_j.
\end{eqnarray*}
```

""" sincos!

@doc doc"""
    tan!(c, a, p, k::Int) --> nothing

Update the `k-th` expansion coefficients `c[k+1]` of `c = tan(a)`,
for `c` and `a` either `Taylor1` or `TaylorN`; `p = c^2` and
is passed as an argument for efficiency.

The coefficients are given by

```math
\begin{equation*}
c_k = a_k + \frac{1}{k} \sum_{j=0}^{k-1} (k-j) a_{k-j} p_j.
\end{equation*}
```

""" tan!

@doc doc"""
    asin!(c, a, r, k)

Update the `k-th` expansion coefficients `c[k+1]` of `c = asin(a)`,
for `c` and `a` either `Taylor1` or `TaylorN`; `r = sqrt(1-c^2)` and
is passed as an argument for efficiency.

```math
\begin{equation*}
c_k = \frac{1}{ \sqrt{r_0} }
    \big( a_k - \frac{1}{k} \sum_{j=1}^{k-1} j r_{k-j} c_j \big).
\end{equation*}
```

""" asin!

@doc doc"""
    acos!(c, a, r, k)

Update the `k-th` expansion coefficients `c[k+1]` of `c = acos(a)`,
for `c` and `a` either `Taylor1` or `TaylorN`; `r = sqrt(1-c^2)` and
is passed as an argument for efficiency.

```math
\begin{equation*}
c_k = - \frac{1}{ r_0 }
    \big( a_k - \frac{1}{k} \sum_{j=1}^{k-1} j r_{k-j} c_j \big).
\end{equation*}
```

""" acos!

@doc doc"""
    atan!(c, a, r, k)

Update the `k-th` expansion coefficients `c[k+1]` of `c = atan(a)`,
for `c` and `a` either `Taylor1` or `TaylorN`; `r = 1+a^2` and
is passed as an argument for efficiency.

```math
\begin{equation*}
c_k = \frac{1}{r_0}\big(a_k - \frac{1}{k} \sum_{j=1}^{k-1} j r_{k-j} c_j\big).
\end{equation*}
```

""" atan!

@doc doc"""
    sinhcosh!(s, c, a, k)

Update the `k-th` expansion coefficients `s[k+1]` and `c[k+1]`
of `s = sinh(a)` and `c = cosh(a)` simultaneously, for `s`, `c` and `a`
either `Taylor1` or `TaylorN`.

The coefficients are given by

```math
\begin{eqnarray*}
s_k &=& \frac{1}{k} \sum_{j=0}^{k-1} (k-j) a_{k-j} c_j, \\
c_k &=& \frac{1}{k} \sum_{j=0}^{k-1} (k-j) a_{k-j} s_j.
\end{eqnarray*}
```

""" sinhcosh!

@doc doc"""
    tanh!(c, a, p, k)

Update the `k-th` expansion coefficients `c[k+1]` of `c = tanh(a)`,
for `c` and `a` either `Taylor1` or `TaylorN`; `p = a^2` and
is passed as an argument for efficiency.

```math
\begin{equation*}
c_k = a_k - \frac{1}{k} \sum_{j=0}^{k-1} (k-j) a_{k-j} p_j.
\end{equation*}
```

""" tanh!

@doc doc"""
    asinh!(c, a, r, k)

Update the `k-th` expansion coefficients `c[k+1]` of `c = asinh(a)`,
for `c` and `a` either `Taylor1` or `TaylorN`; `r = sqrt(1-c^2)` and
is passed as an argument for efficiency.

```math
\begin{equation*}
c_k = \frac{1}{ \sqrt{r_0} }
    \big( a_k - \frac{1}{k} \sum_{j=1}^{k-1} j r_{k-j} c_j \big).
\end{equation*}
```

""" asinh!

@doc doc"""
    acosh!(c, a, r, k)

Update the `k-th` expansion coefficients `c[k+1]` of `c = acosh(a)`,
for `c` and `a` either `Taylor1` or `TaylorN`; `r = sqrt(c^2-1)` and
is passed as an argument for efficiency.

```math
\begin{equation*}
c_k = \frac{1}{ r_0 }
    \big( a_k - \frac{1}{k} \sum_{j=1}^{k-1} j r_{k-j} c_j \big).
\end{equation*}
```

""" acosh!

@doc doc"""
    atanh!(c, a, r, k)

Update the `k-th` expansion coefficients `c[k+1]` of `c = atanh(a)`,
for `c` and `a` either `Taylor1` or `TaylorN`; `r = 1-a^2` and
is passed as an argument for efficiency.

```math
\begin{equation*}
c_k = \frac{1}{r_0}\big(a_k + \frac{1}{k} \sum_{j=1}^{k-1} j r_{k-j} c_j\big).
\end{equation*}
```

""" atanh!


module NCP_Exact

include("../Misc.jl")
include("NCP_MC.jl")
include("NCP_VB.jl")

using .Misc
import .NCP_MC, .NCP_VB

using Distributions, SpecialFunctions, Einsum, Combinatorics
import Base.Iterators: product

export allocations, full, generate
export log_marginal, alloc_dist, EP_dist
export log_posterior_T, log_PT

allocations(X::Array{ℜ,N}, H::Vararg{Ƶ}) where {ℜ<:Real, Ƶ<:Integer, N} = Channel(ctype=Array{ℜ,N+length(H)}) do c
    I, R = size(X), prod(H)
    for partition = product(map(Xᵢ -> combinations(1:(Xᵢ+R-1),R-1), X)...)
        push!(c,reshape([sᵢ for (i,pᵢ) ∈ enumerate(partition)
                        for sᵢ = diff([0, pᵢ..., X[i]+R]) .- 1],H...,I...))
    end
end

full(X::Array{ℜ,N}, T::ℜ) where {ℜ<:Real, N} = Channel(ctype=Array{ℜ,N}) do c
    miss = findall(isnan,X)
    M = length(miss)
    Tₘ = Int.(T - nansum(X))
    for partition ∈ combinations(1:(Tₘ+M-1),M-1)
        X_full = copy(X)
        X_full[miss] .= diff([0, partition..., Tₘ+M]) .- 1
        push!(c,X_full)
    end
end

function generate(T::Ƶ, R::Ƶ, I::Vararg{Ƶ}; a::ℜ=1.0) where {ℜ<:Real, Ƶ<:Integer}
    α_I = (a/R) ./ I
    α_R = a/R

    θ_R = rand(Dirichlet(R,α_R))
    θ_I = [rand(Dirichlet(Iₙ,α_Iₙ),R) for (Iₙ,α_Iₙ) in zip(I,α_I)]

    D_R = rand(Categorical(θ_R),T)
    D_I = [rand(Categorical(θ_Iₙ[:,r])) for r ∈ D_R, θ_Iₙ ∈ θ_I]

    return cooccurrences(D_I,I), (θ_R, θ_I...)
end

function log_marginal(X::Array{ℜ,N}, R::Ƶ; a=1.0, b=a/nansum(X)) where {ℜ<:Real, Ƶ<:Integer, N}
    I = size(X)
    α_I = (a/R) ./ I
    α_R = a/R

    S_R, S₊ = Array{ℜ}(undef,R), sum(X)
    S_I = [Array{ℜ}(undef,R,Iₙ) for Iₙ in I]

    log_PX = -Inf
    log_C = a*log(b) - (a+S₊)*log(b + 1) + lgamma(a + S₊) - lgamma(a)

    for S ∈ allocations(X,R)
        log_PS = log_C - sum(lgamma, S .+ 1)
        for (n, Iₙ) ∈ enumerate(I)
            S_I[n] .= sum(S,(2:n...,n+2:N+1...))
            log_PS += sum(lbeta(α_I[n] .+ S_I[n];dims=2)) - R*lbeta(α_I[n],Iₙ)
        end
        S_R .= sum(S,(2:N+1...,))
        log_PS += lbeta(α_R .+ S_R) - lbeta(α_R,R)
        log_PX = logsumexp([log_PX,log_PS])
    end
    return log_PX
end

function alloc_dist(X::Array{ℜ,N}, R::Ƶ; a=1.0, b=a/nansum(X)) where {ℜ<:Real, Ƶ<:Integer, N}
    I = size(X)
    α_I = (a/R) ./ I
    α_R = a/R

    S_R, S₊ = Array{ℜ}(undef,R), sum(X)
    S_I = [Array{ℜ}(undef,R,Iₙ) for Iₙ in I]

    nS = prod(x -> binomial(Ƶ(x+R-1),R-1), X)
    log_P = zeros(nS)

    log_C = a*log(b) - (a+S₊)*log(b + 1) + lgamma(a + S₊) - lgamma(a)

    for (s, S) ∈ enumerate(allocations(X,R))
        log_P[s] = log_C - sum(lgamma, S .+ 1)
        for (n, Iₙ) ∈ enumerate(I)
            S_I[n] .= sum(S,(2:n...,n+2:N+1...))
            log_P[s] += sum(lbeta(α_I[n] .+ S_I[n];dims=2)) - R*lbeta(α_I[n],Iₙ)
        end
        S_R .= sum(S,(2:N+1...,))
        log_P[s] += lbeta(α_R .+ S_R) - lbeta(α_R,R)
    end
    return log_P
end

function EP_dist(X::Array{ℜ,N}, R::Ƶ) where {ℜ<:Real, Ƶ<:Integer, N}
    I = size(X)

    S_R, S₊ = Array{ℜ}(undef,R), sum(X)
    S_I = [Array{ℜ}(undef,R,Iₙ) for Iₙ in I]

    nS = prod(x -> binomial(Ƶ(x+R-1),R-1), X)
    dₑₚ = zeros(nS)

    for (s, S) ∈ enumerate(allocations(X,R))
        for (n, Iₙ) ∈ enumerate(I)
            S_I[n] .= sum(S,(2:n...,n+2:N+1...))
            dₑₚ[s] += sum(!iszero, S_I[n])
        end
        S_R .= sum(S,(2:N+1...,))
        dₑₚ[s] += (N-1) * sum(!iszero, S_R) + !iszero(S₊)
    end
    return dₑₚ
end

function log_posterior_T(X::Array{ℜ,N}, T::ℜ , R::Ƶ; a=1.0, b=a/nansum(X),
    smc::Bool=false, elbo::Bool=false, particles::Int=1000, EPOCHS::Int=100, M::Int=100) where {ℜ<:Real, Ƶ<:Integer, N}
    if smc
        log_PX = [NCP_MC.smc_weight(X_full,R,particles; a=a, b=b)[1] for X_full ∈ full(X,T)]
        return logsumexp(log_PX)
    elseif elbo
        log_PX = nanmax([NCP_VB.standard_VB(X_full,R; a=a, b=b, EPOCHS=EPOCHS)[1][end] for m ∈ 1:M, X_full ∈ full(X,T)],1)
        return logsumexp(log_PX)
    else
        log_PX = [log_marginal(X_full,R; a=a, b=b) for X_full ∈ full(X,T)]
        return logsumexp(log_PX)
    end
end

end


module TestInputNoise
using Test
using LinearAlgebra
using GaussianDistributions
using CovarianceFunctions

@testset "input noise" begin
    # set up GP
    l = 1/4
    k = CovarianceFunctions.Lengthscale(CovarianceFunctions.EQ(), l)
    G = Gaussian(k)

    # set up data
    n = 32 # number of observations
    x = randn(n) # uncorrupted input

    σx = 5e-2 # standard deviation of input noise
    σy = 1e-2 # standard deviation of ouput noise
    εx = σx .* randn(n) # input noise
    εy = σy .* randn(n) # output noise

    f(x) = sin(2π*x)
    y = @. f(x + εx) + εy # output

    C = G | (x, y, σx^2) # canonical conditional distribution
    Ci = G | (x, σx^2, y, σy^2) # NIGP conditional distribution

    xs = 2maximum(abs, x) # outside of data range, where regular posterior mean is constant,
    @test var(Ci)(xs) ≈ var(C)(xs) # NIGP variance is unchanged to regular variance
    xs = 2rand(n) .- 1
    nincrease = sum(var(C, xs) .< var(Ci, xs))
    @test nincrease > n * 3/4 # NIGP tends to increase uncertainty in data-rich regions

    # using Plots
    # plot(Ci, label = "NIGP")
    # plot!(C, label = "GP")
    # scatter!(x + εx, y + εy, label = "data")
    # xs = range(minimum(x) - std(x), maximum(x) + std(x), length = 128)
    # plot!(xs, f, label = "f(x)")
    # gui()
end

end

# plotting uncertainty
# xs = range(-3, 3, length = 256)
# plot(xs, var(Ci), label = "NIGP")
# plot!(xs, var(C), label = "GP")
# scatter!(x + εx, zero(x), label = "data")

# d = var(Ci, xs) .- var(C, xs)
# println([t for t in zip(x, d)])
# i = findfirst(<(0), d)
# println(xs[i])
# println(var(Ci)(xs[i]))
# println(var(C)(xs[i]))
# scatter!([xs[i]], [var(Ci)(xs[i])])
# gui()


function foo()
    lines = readlines(open("data/6.txt", "r"))
    points = zeros(Int, length(lines), 2)
    maxx = 0
    for (i, line) in enumerate(lines)
        point = parse.(Int, split(line, ','))
        if any(point .> maxx)
            maxx = maximum(point)
        end
        points[i, :] = point
    end
    areas = Dict{Int, Int}()

    bad_idx = Set{Int}()
    stop = maxx
    start = 0
    for i in start:stop
        for j in start:stop
            dist = abs.(points[:, 1] .- i) .+ abs.(points[:, 2] .- j)
            min_dist = minimum(dist)
            if sum(dist .== min_dist) == 1
                idx = argmin(dist)
                if !(idx in keys(areas))
                    areas[idx] = 0
                end
                areas[idx] += 1
                if i == start || i == stop || j == start || j == stop
                    push!(bad_idx, idx)
                end
            end
        end
    end
    for idx in bad_idx
        delete!(areas, idx)
    end
    return maximum(values(areas))
end

println(foo())


module TestTreeDP

using Test
import Printf: @sprintf
import TreeLas
import TreeLas.TreeDP
import TreeLas.Dual: dual
import GraphIdx
import GraphIdx.Tree: ChildrenIndex, hierarchy, hierarchy_with, hierarchy_string

include(joinpath(dirname(pathof(TreeLas)), "cxx.jl"))

mean(x) = sum(x) / length(x)

@testset "TreeDP                         " begin

@testset "tree_dp: tree.mini.h5          " begin
    t = TreeDP.Tree([1, 1, 2, 1, 2, 3])
    y = [0.53 0.45 0.02
         0.17 0.96 0.97]
    y = Array(y')
    lam = 0.2
    x = TreeDP.tree_dp(y, t, lam)
    @test size(x) == size(y)
    xt = [0.4  0.4  0.4
          0.37 0.76 0.77]
    xt = Array(xt')
    if false
        hierarchy_with(ChildrenIndex(t.parent)) do i
            @sprintf("%.2f:%.2f", y[i], xt[i])
        end
    end
    @test x ≈ xt
end

@testset "tree_dp: 10                    " begin
    t = TreeDP.Tree(Int[1, 1, 2, 3, 4, 1, 8, 9, 4, 9])
    y = Float64[8.2, 7.0, 9.5, 6.8, 5.8, 6.3, 4.3, 2.2, 1.2, 2.8]
    lam = 1.0

    x = TreeDP.tree_dp(copy(y), t, lam)
    @test size(x) == size(y)
    expect3x = [22.7, 22.7, 22.7, 18.9, 18.9, 21.9,  9.9,  8.2,  8.2,  8.2]
    @test 3x ≈ expect3x
    local cidx = ChildrenIndex(t.parent)
    if true
        let z = expect3x/3 - y
            local alpha = dual(z, t.parent, cidx)
            # @info "alpha=" *string(round.(alpha, digits=4))
            @debug hierarchy_string(
                cidx, i -> @sprintf("%.1f≈%+.1f", -3y[i]+expect3x[i], -3alpha[i])
            )
        end
    end
end

@testset "MGT     : demo 3x7             " begin
    y = [0.62 0.73 0.71 1.5 1.17 0.43 1.08
         0.62 1.73 0.95 1.46 1.6 1.16 0.38
         0.9 0.32 -0.48 0.95 1.08 0.02 0.4]
    local n = length(y)
    local root = 1
    edges = Tuple{Int64,Int64}[
        (1,  2),  (4,  5),  (7, 8), (10, 11), (13, 14), (16, 17),
        (19, 20), (2,  3),  (5, 6), (8,   9), (11, 12), (14, 15),
        (17, 18), (20, 21), (1, 4), (4,   7), (7,  10), (10, 13),
        (13, 16), (16, 19), (2, 5), (5,   8), (8,  11), (11, 14),
        (14, 17), (17, 20), (3, 6), (6,   9), (9,  12), (12, 15),
        (15, 18), (18, 21)
    ]

    γ = [-abs(y[u] - y[v]) for (u, v) in edges]

    parent =
        [1, 5, 6, 1, 4, 5, 8, 9, 6, 7, 8, 9, 10, 15, 18, 13, 16, 17, 20, 17, 18]

    local tree = GraphIdx.Tree.RootedTree(root, parent)
    local dp_mem = TreeDP.TreeDPMem(n)
    local x = fill(NaN, size(y)...)
    local tree_alpha = fill(NaN, n)
    let mu = TreeDP.Ones{Float64}(), z = copy(y)
        TreeDP.tree_dp!(x,
                        z,
                        tree,
                        TreeDP.Ones{Float64}(),
                        mu,
                        dp_mem)
        @test !any(isnan.(x))
    end

    @testset "analyse solution" begin
        @test length(unique(x)) == 2
        @test mean(y) ≈ mean(x)
    end

    @testset "compare to cxx tree_dp" begin
        cxx_x = Cxx.cxx_tree_dp(y, parent, root, 1.0)
        @test cxx_x ≈ x
    end

    @testset "dual!" begin
        @test dp_mem.proc_order[end] == root
        @test dp_mem.proc_order[1:end-1] ==
            [12, 11, 19, 20, 21, 14, 15, 18, 17, 16, 13, 10, 7, 8, 9, 3, 6, 2, 5, 4]
        TreeLas.Dual.dual!(tree_alpha, x, copy(y), dp_mem.proc_order, parent)
        local alpha = vec(x) - vec(y)
        for i in @view dp_mem.proc_order[1:end-1]
            alpha[parent[i]] += alpha[i]
        end
        if !(alpha[2:end] ≈ tree_alpha[2:end])
            @show alpha - tree_alpha
            @test alpha[2:end] ≈ tree_alpha[2:end]
        end
        exceed = findall(abs.(tree_alpha) .> 1 + 1e12)
        if exceed != []
            for ei in exceed
                @show ei
                @show alpha[ei]
            end
        end
    end

    @testset "compare to tree_dual" begin
        cxx_alpha = Cxx.cxx_tree_dual!(x - copy(y), parent, root)
        @test abs.(tree_alpha[2:end]) ≈ abs.(cxx_alpha[2:end])
        sign01(x) = Int((sign(x) + 1) / 2)
        if !(sign01.(tree_alpha[2:end]) == sign01.(cxx_alpha[2:end]))
            hierarchy(ChildrenIndex(parent))
            @test sign01.(tree_alpha[2:end]) == sign01.(cxx_alpha[2:end])
        end
        @test tree_alpha[2:end] ≈ cxx_alpha[2:end]
    end

    local wtree = GraphIdx.Tree.WeightedTree(tree, TreeDP.Ones{Float64}())
    let gam = fill(NaN, n)
        TreeLas.MGT.gap_vec!(gam, x, tree_alpha, wtree)
        @test isnan(gam[root])
        gam[root] = 0.0
        wrong_ids = [i for i in 1:n if abs(gam[i]) > 1e-16]
        if wrong_ids != []
            for i in wrong_ids
                @show i, parent[i]
                @show x[i], x[parent[i]]
                @show tree_alpha[i]
                @show gam[i]
            end
            hierarchy(ChildrenIndex(parent))
            @test wrong_ids == []
        end
    end
end

end
end


(F::FqFiniteField)(coeffs::Array{T,1}) where {T<:Union{Integer,fmpz}} = begin
    g = gen(F)
    x = zero(F)
    for (i, c) in enumerate(coeffs)
        x += c * g^(i-1)
    end
    x
end

(F::FqNmodFiniteField)(coeffs::Array{T,1}) where {T<:Union{Integer,fmpz}} = begin
    g = gen(F)
    x = zero(F)
    for (i, c) in enumerate(coeffs)
        x += c * g^(i-1)
    end
    x
end

sqrt(a::FinFieldElem) = root(a, 2)

# http://trac.sagemath.org/ticket/7931
# http://sagenb.org/src/rings/finite_rings/element_base.pyx
root(a::FinFieldElem, n::Integer) = root(a, fmpz(n))
root(a::FinFieldElem, n::fmpz) = begin
    if iszero(a)
        n <= 0 && throw(DomainError())
        return a
    end

    K = parent(a)
    q = order(K)

    if n < 0
        a = inv(a)
        n = -n
    elseif n == 0
        a == 1 || throw(DomainError())
        return a
    end

    if isone(a)
        GCD = gcd(n, q-1)
        GCD == 1 && return a

        g = gen(K) # TODO: The generator is guaranteed to be a multiplicative generator only if the field is generated by a Conway polynomial.
        q1overn = (q-1) ÷ GCD
        nthroot = g^q1overn
        return nthroot
    end

    m = n % (q-1)
    m == 0 && error("$(a) has no $(n)th root in $(K)")

    # GCD = α*m + β*(q-1), so 1/m = α/GCD (mod q-1)
    GCD, α, β = gcdx(m, q-1)
    GCD == 1 && return a^α

    m = GCD
    q1overn = (q-1) ÷ m
    a^q1overn != 1 && error("$(a) has no $(n)th root in $(K)")

    b = a^α
    F = [(fmpz(p), e) for (p, e) in factor(BigInt(m))]

    g = gen(K) # TODO: The generator is guaranteed to be a multiplicative generator only if the field is generated by a Conway polynomial.
    for (r, v) in F
        k, h = remove(q-1, r)
        z = h * invmod(-h, r^v)::typeof(h)
        x = (1 + z) ÷ (r^v)
        if k == 1
            b = b^x
        else
            t = log(b^h, g^(r^v * h), r^(k-v))
            b = b^x * g^(-z*t)
        end
    end
    b
end

sqrts(a::FinFieldElem) = roots(a, 2)

# http://trac.sagemath.org/ticket/7931
# http://sagenb.org/src/rings/finite_rings/element_base.pyx
roots(a::FinFieldElem, n::Integer) = roots(a, fmpz(n))
roots(a::T, n::fmpz) where {T<:FinFieldElem} = begin
    if iszero(a)
        n <= 0 && throw(DomainError())
        return [a]
    end

    K = parent(a)
    q = order(K)

    if n < 0
        a = inv(a)
        n = -n
    elseif n == 0
        a == 1 || throw(DomainError())
        e = elements(K)
        e[e .!= 0]
    end

    if isone(a)
        GCD = gcd(n, q-1)
        GCD == 1 && return [a]

        g = gen(K) # TODO: The generator is guaranteed to be a multiplicative generator only if the field is generated by a Conway polynomial.
        q1overn = (q-1) ÷ GCD
        nthroot = g^q1overn
        return [nthroot^i for i in 0:(GCD-1)]
    end

    m = n % (q-1)
    m == 0 && return T[]

    # GCD = α*m + β*(q-1), so 1/m = α/GCD (mod q-1)
    GCD, α, β = gcdx(m, q-1)
    GCD == 1 && return [a^α]

    m = GCD
    q1overn = (q-1) ÷ m
    a^q1overn != 1 && return T[]

    b = a^α
    F = [(fmpz(p), e) for (p, e) in factor(BigInt(m))]

    g = gen(K) # TODO: The generator is guaranteed to be a multiplicative generator only if the field is generated by a Conway polynomial.
    for (r, v) in F
        k, h = remove(q-1, r)
        z = h * invmod(-h, r^v)::typeof(h)
        x = (1 + z) ÷ (r^v)
        if k == 1
            b = b^x
        else
            t = log(b^h, g^(r^v * h), r^(k-v))
            b = b^x * g^(-z*t)
        end
    end

    nthroot = g^q1overn
    L = [b]
    for _ in 1:(m-1)
        b *= nthroot
        push!(L, b)
    end
    L
end

elements(F::Union{FqFiniteField,FqNmodFiniteField}) = begin
    p = characteristic(F)
    k = degree(F)

    e = typeof(zero(F))[]
    @forcartesian c [p for _ in 1:k] begin
        push!(e, F(c .- 1))
    end
    e
end

rand(F::Union{FqFiniteField,FqNmodFiniteField}) = begin
    p = characteristic(F)
    k = degree(F)

    F([rand(0:(p-1)) for _ in 1:k])
end


# Dimensionality operations for 4D Tensors
# Author: Philipp Witte, pwitte3@gatech.edu
# Date: January 2020

export squeeze, unsqueeze, wavelet_squeeze, wavelet_unsqueeze, tensor_split, tensor_cat

####################################################################################################
# Squeeze and unsqueeze

"""
    Y = squeeze(X; pattern="column")

 Reshape input image such that each spatial dimension is reduced by a factor
 of 2, while the number of channels is increased by a factor of 4.

 *Input*:

 - `X`: 4D input tensor of dimensions `nx` x `ny` x `n_channel` x `batchsize`

 - `pattern`: Squeezing pattern

        1 2 3 4        1 1 3 3        1 3 1 3
        1 2 3 4        1 1 3 3        2 4 2 4
        1 2 3 4        2 2 4 4        1 3 1 3
        1 2 3 4        2 2 4 4        2 4 2 4

        column          patch       checkerboard

 *Output*:

 - `Y`: Reshaped tensor of dimensions `nx/2` x `ny/2` x `n_channel*4` x `batchsize`

 See also: [`unsqueeze`](@ref), [`wavelet_squeeze`](@ref), [`wavelet_unsqueeze`](@ref)
"""
function squeeze(X::AbstractArray{T,4}; pattern="column") where T

    # Dimensions
    nx_in, ny_in, nc_in, batchsize = size(X)
    if mod(nx_in, 2) == 1 || mod(ny_in, 2) == 1
        throw("Input dimensions must be multiple of 2")
    end
    nx_out = Int(round(nx_in/2))
    ny_out = Int(round(ny_in/2))
    nc_out = Int(round(nc_in*4))
    Y = cuzeros(X, nx_out, ny_out, nc_out, batchsize)

    if pattern == "column"
        Y = reshape(X, nx_out, ny_out, nc_out, batchsize)
    elseif pattern == "patch"
        Y[:, :, 1: nc_in, :] = X[1:nx_out, 1:ny_out, :, :]
        Y[:, :, nc_in+1: 2*nc_in, :] = X[nx_out+1:end, 1:ny_out, :, :]
        Y[:, :, 2*nc_in+1: 3*nc_in, :] = X[1:nx_out, ny_out+1:end, :, :]
        Y[:, :, 3*nc_in+1: 4*nc_in, :] = X[nx_out+1:end, ny_out+1:end, :, :]
    elseif pattern == "checkerboard"
        Y[:, :, 1: nc_in, :] = X[1:2:nx_in, 1:2:ny_in, 1:nc_in, :]
        Y[:, :, nc_in+1: 2*nc_in, :] = X[2:2:nx_in, 1:2:ny_in, 1:nc_in, :]
        Y[:, :, 2*nc_in+1: 3*nc_in, :] = X[1:2:nx_in, 2:2:ny_in, 1:nc_in, :]
        Y[:, :, 3*nc_in+1: 4*nc_in, :] = X[2:2:nx_in, 2:2:ny_in, 1:nc_in, :]
    else
        throw("Specified pattern not defined.")
    end
    return Y
end

function squeeze(X::AbstractArray{T,5}; pattern="column") where T

    # Dimensions
    nx_in, ny_in, nz_in, nc_in, batchsize = size(X)
    if mod(nx_in, 2) == 1 || mod(ny_in, 2) == 1
        throw("Input dimensions must be multiple of 2")
    end
    nx_out = Int(round(nx_in/2))
    ny_out = Int(round(ny_in/2))
    nz_out = Int(round(nz_in/2))
    nc_out = Int(round(nc_in*8))
    Y = cuzeros(X, nx_out, ny_out, nz_out, nc_out, batchsize)

    if pattern == "column"
        Y = reshape(X, nx_out, ny_out, nz_out, nc_out, batchsize)
    elseif pattern == "patch"
        Y[:, :, :, 1: nc_in, :] = X[1:nx_out, 1:ny_out, 1:nz_out, :, :]
        Y[:, :, :, nc_in+1: 2*nc_in, :] = X[nx_out+1:end, 1:ny_out, 1:nz_out, :, :]
        Y[:, :, :, 2*nc_in+1: 3*nc_in, :] = X[1:nx_out, ny_out+1:end, 1:nz_out, :, :]
        Y[:, :, :, 3*nc_in+1: 4*nc_in, :] = X[nx_out+1:end, ny_out+1:end, 1:nz_out, :, :]
        Y[:, :, :, 4*nc_in+1: 5*nc_in, :] = X[1:nx_out, 1:ny_out, nz_out+1:end, :, :]
        Y[:, :, :, 5*nc_in+1: 6*nc_in, :] = X[nx_out+1:end, 1:ny_out, nz_out+1:end, :, :]
        Y[:, :, :, 6*nc_in+1: 7*nc_in, :] = X[1:nx_out, ny_out+1:end, nz_out+1:end, :, :]
        Y[:, :, :, 7*nc_in+1: 8*nc_in, :] = X[nx_out+1:end, ny_out+1:end, nz_out+1:end, :, :]
    elseif pattern == "checkerboard"
        throw("Currently not implemented for 5D tensors.")
    else
        throw("Specified pattern not defined.")
    end
    return Y
end

"""
    X = unsqueeze(Y; pattern="column")

 Undo squeezing operation by reshaping input image such that each spatial dimension is
 increased by a factor of 2, while the number of channels is decreased by a factor of 4.

 *Input*:

 - `Y`: 4D input tensor of dimensions `nx` x `ny` x `n_channel` x `batchsize`

 - `pattern`: Squeezing pattern

            1 2 3 4        1 1 3 3        1 3 1 3
            1 2 3 4        1 1 3 3        2 4 2 4
            1 2 3 4        2 2 4 4        1 3 1 3
            1 2 3 4        2 2 4 4        2 4 2 4

            column          patch       checkerboard

 *Output*:

 - `X`: Reshaped tensor of dimensions `nx*2` x `ny*2` x `n_channel/4` x `batchsize`

 See also: [`squeeze`](@ref), [`wavelet_squeeze`](@ref), [`wavelet_unsqueeze`](@ref)
"""
function unsqueeze(Y::AbstractArray{T,4}; pattern="column") where T

    # Dimensions
    nx_in, ny_in, nc_in, batchsize = size(Y)
    if mod(nx_in, 2) == 1 || mod(ny_in, 2) == 1
        throw("Input dimensions must be multiple of 2")
    end
    nx_out = Int(round(nx_in*2))
    ny_out = Int(round(ny_in*2))
    nc_out = Int(round(nc_in/4))
    X = cuzeros(Y, nx_out, ny_out, nc_out, batchsize)

    if pattern == "column"
        X = reshape(Y, nx_out, ny_out, nc_out, batchsize)
    elseif pattern == "patch"
        X[1:nx_in, 1:ny_in, :, :] = Y[:, :, 1: nc_out, :]
        X[nx_in+1:end, 1:ny_in, :, :] = Y[:, :, nc_out+1: 2*nc_out, :]
        X[1:nx_in, ny_in+1:end, :, :] = Y[:, :, 2*nc_out+1: 3*nc_out, :]
        X[nx_in+1:end, ny_in+1:end, :, :] = Y[:, :, 3*nc_out+1: 4*nc_out, :]
    elseif pattern == "checkerboard"
        X[1:2:nx_out, 1:2:ny_out, 1:nc_out, :] = Y[:, :, 1: nc_out, :]
        X[2:2:nx_out, 1:2:ny_out, 1:nc_out, :] = Y[:, :, nc_out+1: 2*nc_out, :]
        X[1:2:nx_out, 2:2:ny_out, 1:nc_out, :] = Y[:, :, 2*nc_out+1: 3*nc_out, :]
        X[2:2:nx_out, 2:2:ny_out, 1:nc_out, :] = Y[:, :, 3*nc_out+1: 4*nc_out, :]
    else
        throw("Specified pattern not defined.")
    end
    return X
end

function unsqueeze(Y::AbstractArray{T,5}; pattern="column") where T

    # Dimensions
    nx_in, ny_in, nz_in, nc_in, batchsize = size(Y)
    if mod(nx_in, 2) == 1 || mod(ny_in, 2) == 1
        throw("Input dimensions must be multiple of 2")
    end
    nx_out = Int(round(nx_in*2))
    ny_out = Int(round(ny_in*2))
    nz_out = Int(round(ny_in*2))
    nc_out = Int(round(nc_in/8))
    X = cuzeros(Y, nx_out, ny_out, nz_out, nc_out, batchsize)

    if pattern == "column"
        X = reshape(Y, nx_out, ny_out, nz_out, nc_out, batchsize)
    elseif pattern == "patch"
        X[1:nx_in, 1:ny_in, 1:nz_in, :, :] = Y[:, :, :, 1: nc_out, :]
        X[nx_in+1:end, 1:ny_in, 1:nz_in, :, :] = Y[:, :, :, nc_out+1: 2*nc_out, :]
        X[1:nx_in, ny_in+1:end, 1:nz_in, :, :] = Y[:, :, :, 2*nc_out+1: 3*nc_out, :]
        X[nx_in+1:end, ny_in+1:end, 1:nz_in, :, :] = Y[:, :, :, 3*nc_out+1: 4*nc_out, :]
        X[1:nx_in, 1:ny_in, nz_in+1:end, :, :] = Y[:, :, :, 4*nc_out+1: 5*nc_out, :]
        X[nx_in+1:end, 1:ny_in, nz_in+1:end, :, :] = Y[:, :, :, 5*nc_out+1: 6*nc_out, :]
        X[1:nx_in, ny_in+1:end, nz_in+1:end, :, :] = Y[:, :, :, 6*nc_out+1: 7*nc_out, :]
        X[nx_in+1:end, ny_in+1:end, nz_in+1:end, :, :] = Y[:, :, :, 7*nc_out+1: 8*nc_out, :]
    elseif pattern == "checkerboard"
        throw("Specified pattern not defined.")
    else
        throw("Specified pattern not defined.")
    end
    return X
end

function unsqueeze(X::AbstractArray{T,4}, Y::AbstractArray{T,4}; pattern="column") where T
    return unsqueeze(X; pattern=pattern), unsqueeze(Y; pattern=pattern)
end

function unsqueeze(X::AbstractArray{T,5}, Y::AbstractArray{T,5}; pattern="column") where T
    return unsqueeze(X; pattern=pattern), unsqueeze(Y; pattern=pattern)
end

####################################################################################################
# Squeeze and unsqueeze using the wavelet transform

"""
    Y = wavelet_squeeze(X; type=WT.db1)

 Perform a 1-level channelwise 2D wavelet transform of X and squeeze output of each
 transform into 4 channels (per 1 input channel).

 *Input*:

 - `X`: 4D input tensor of dimensions `nx` x `ny` x `n_channel` x `batchsize`

 - `type`: Wavelet filter type. Possible values are `WT.haar` for Haar wavelets,
    `WT.coif2`, `WT.coif4`, etc. for Coiflet wavelets, or `WT.db1`, `WT.db2`, etc.
    for Daubechies wavetlets. See *https://github.com/JuliaDSP/Wavelets.jl* for a
    full list.

 *Output*:

 - `Y`: Reshaped tensor of dimensions `nx/2` x `ny/2` x `n_channel*4` x `batchsize`

 See also: [`wavelet_unsqueeze`](@ref), [`squeeze`](@ref), [`unsqueeze`](@ref)
"""
function wavelet_squeeze(X::AbstractArray{T,4}; type=WT.db1) where T
    nx_in, ny_in, nc_in, batchsize = size(X)
    nx_out = Int(round(nx_in/2))
    ny_out = Int(round(ny_in/2))
    nc_out = Int(round(nc_in*4))
    Y = cuzeros(X, nx_out, ny_out, nc_out, batchsize)
    for i=1:batchsize
        for j=1:nc_in
            Ycurr = dwt(X[:,:,j,i], wavelet(type), 1)
            Y[:, :, (j-1)*4 + 1: j*4, i] = squeeze(reshape(Ycurr, nx_in, ny_in, 1, 1); pattern="patch")
        end
    end
    return Y
end

function wavelet_squeeze(X::AbstractArray{T,5}; type=WT.db1) where T
    nx_in, ny_in, nz_in, nc_in, batchsize = size(X)
    nx_out = Int(round(nx_in/2))
    ny_out = Int(round(ny_in/2))
    nz_out = Int(round(nz_in/2))
    nc_out = Int(round(nc_in*8))
    Y = cuzeros(X, nx_out, ny_out, nz_out, nc_out, batchsize)
    for i=1:batchsize
        for j=1:nc_in
            Ycurr = dwt(X[:,:,:,j,i], wavelet(type), 1)
            Y[:, :, :, (j-1)*8 + 1: j*8, i] = squeeze(reshape(Ycurr, nx_in, ny_in, nz_in, 1, 1); pattern="patch")
        end
    end
    return Y
end

"""
    X = wavelet_unsqueeze(Y; type=WT.db1)

 Perform a 1-level inverse 2D wavelet transform of Y and unsqueeze output.
 This reduces the number of channels by 4 and increases each spatial
 dimension by a factor of 2. Inverse operation of `wavelet_squeeze`.

 *Input*:

 - `Y`: 4D input tensor of dimensions `nx` x `ny` x `n_channel` x `batchsize`

 - `type`: Wavelet filter type. Possible values are `haar` for Haar wavelets,
  `coif2`, `coif4`, etc. for Coiflet wavelets, or `db1`, `db2`, etc. for Daubechies
  wavetlets. See *https://github.com/JuliaDSP/Wavelets.jl* for a full list.

 *Output*:

 - `X`: Reshaped tensor of dimenions `nx*2` x `ny*2` x `n_channel/4` x `batchsize`

 See also: [`wavelet_squeeze`](@ref), [`squeeze`](@ref), [`unsqueeze`](@ref)
"""
function wavelet_unsqueeze(Y::AbstractArray{T,4}; type=WT.db1) where T
    nx_in, ny_in, nc_in, batchsize = size(Y)
    nx_out = Int(round(nx_in*2))
    ny_out = Int(round(ny_in*2))
    nc_out = Int(round(nc_in/4))
    X = cuzeros(Y, nx_out, ny_out, nc_out, batchsize)
    for i=1:batchsize
        for j=1:nc_out
            Ycurr = unsqueeze(Y[:, :, (j-1)*4 + 1: j*4, i:i]; pattern="patch")[:, :, 1, 1]
            Xcurr = idwt(Ycurr, wavelet(type), 1)
            X[:, :, j, i] = reshape(Xcurr, nx_out, ny_out, 1, 1)
        end
    end
    return X
end

function wavelet_unsqueeze(Y::AbstractArray{T,5}; type=WT.db1) where T
    nx_in, ny_in, nz_in, nc_in, batchsize = size(Y)
    nx_out = Int(round(nx_in*2))
    ny_out = Int(round(ny_in*2))
    nz_out = Int(round(nz_in*2))
    nc_out = Int(round(nc_in/8))
    X = cuzeros(Y, nx_out, ny_out, nz_out, nc_out, batchsize)
    for i=1:batchsize
        for j=1:nc_out
            Ycurr = unsqueeze(Y[:, :, :, (j-1)*8 + 1: j*8, i:i]; pattern="patch")[:, :, :, 1, 1]
            Xcurr = idwt(Ycurr, wavelet(type), 1)
            X[:, :, :, j, i] = reshape(Xcurr, nx_out, ny_out, nz_out, 1, 1)
        end
    end
    return X
end

####################################################################################################
# Split and concatenate

"""
    Y, Z = tensor_split(X)

 Split 4D input tensor in half along the channel dimension. Inverse operation
 of `tensor_cat`.

 *Input*:

 - `X`: 4D input tensor of dimensions `nx` x `ny` x `n_channel` x `batchsize`

 *Output*:

 - `Y`, `Z`: 4D output tensors, each of dimensions `nx` x `ny` x `n_channel/2` x `batchsize`

 See also: [`tensor_cat`](@ref)
"""
function tensor_split(X::AbstractArray{T,4}; split_index=nothing) where T
    if isnothing(split_index)
        k = Int(round(size(X, 3)/2))
    else
        k = split_index
    end
    return X[:, :, 1:k, :], X[:, :, k+1:end, :]
end

function tensor_split(X::AbstractArray{T,5}; split_index=nothing) where T
    if isnothing(split_index)
        k = Int(round(size(X, 4)/2))
    else
        k = split_index
    end
    return X[:, :, :, 1:k, :], X[:, :, :, k+1:end, :]
end

function tensor_split(X::AbstractArray{T,1}; split_index=nothing) where T
    if isnothing(split_index)
        k = Int(round(size(X, 1)/2))
    else
        k = split_index
    end
    return X[1:k], X[k+1:end]
end

"""
    X = tensor_cat(Y, Z)

 Concatenate 4D input tensors along the channel dimension. Inverse operation
 of `tensor_split`.

 *Input*:

 - `Y`, `Z`: 4D input tensors, each of dimensions `nx` x `ny` x `n_channel` x `batchsize`

 *Output*:

 - `X`: 4D output tensor of dimensions `nx` x `ny` x `n_channel*2` x `batchsize`

 See also: [`tensor_split`](@ref)
"""
function tensor_cat(X::AbstractArray{T,4}, Y::AbstractArray{T,4}) where T
    if size(X, 3) == 0
        return Y
    elseif size(Y, 3) == 0
        return X
    else
        return cat(X, Y; dims=3)
    end
end

function tensor_cat(X::AbstractArray{T,5}, Y::AbstractArray{T,5}) where T
    if size(X, 4) == 0
        return Y
    elseif size(Y, 4) == 0
        return X
    else
        return cat(X, Y; dims=4)
    end
end

function tensor_cat(X::AbstractArray{T,1}, Y::AbstractArray{T,1}) where T
    if size(X, 1) == 0
        return Y
    elseif size(Y, 1) == 0
        return X
    else
        return cat(X, Y; dims=1)
    end
end

tensor_cat(X::Tuple{AbstractArray{T,4}, AbstractArray{T,4}}) where T = tensor_cat(X[1], X[2])
tensor_cat(X::Tuple{AbstractArray{T,5}, AbstractArray{T,5}}) where T = tensor_cat(X[1], X[2])


export proxLLR!, normLLR

"""
    proxLLR!(x::Vector{T}, λ::Float64=1e-6; kargs...) where T

proximal map for LLR regularization using singular-value-thresholding

# Arguments
* `x::Vector{T}`                - Vector to apply proximal map to
* `λ::Float64`                  - regularization parameter
* `shape::Tuple{Int}=[]`        - dimensions of the image
* `blockSize::Tuple{Int}=[2;2]` - size of patches to perform singluar value thresholding on
* `randshift::Bool=true`        - randomly shifts the patches to ensure translation invariance
"""
function proxLLR!(x::Vector{T}, λ::Float64=1e-6; shape::NTuple=[],
   blockSize::Array{Int64,1}=[2; 2], randshift::Bool=true, kargs...) where T
  # xᵖʳᵒˣ = zeros(T,size(x))
  N = prod(shape)
  # K = floor(Int,length(x)/N)
  x[:] = vec( svt(x[:], shape, λ; blockSize=blockSize, randshift=randshift, kargs...) )

  return x
end

function svt(x::Vector{T}, shape::Tuple, λ::Float64=1e-6;
   blockSize::Array{Int64,1}=[2; 2], randshift::Bool=true, kargs...) where T

  x = reshape( x, tuple( shape...,floor(Int64, length(x)/prod(shape)) ) )

  Wy = blockSize[1]
  Wz = blockSize[2]

  if randshift
    # Random.seed!(1234)
    shift_idx = [rand(1:Wy) rand(1:Wz) 0]
    x = circshift(x, shift_idx)
  end

  ny, nz, K = size(x)

  # reshape into patches
  L = floor(Int,ny*nz/Wy/Wz) # number of patches, assumes that image dimensions are divisble by the blocksizes

  xᴸᴸᴿ = zeros(T,Wy*Wz,L,K)
  for i=1:K
    xᴸᴸᴿ[:,:,i] = im2colDistinct(x[:,:,i], (Wy,Wz))
  end
  xᴸᴸᴿ = permutedims(xᴸᴸᴿ,[1 3 2])

  # threshold singular values
  for i = 1:L
    if xᴸᴸᴿ[:,:,i] == zeros(T, Wy*Wz,K)
      continue
    end
    SVDec = svd(xᴸᴸᴿ[:,:,i])
    proxL1!(SVDec.S,λ)
    xᴸᴸᴿ[:,:,i] = SVDec.U*Matrix(Diagonal(SVDec.S))*SVDec.Vt
  end

  # reshape into image
  xᵗʰʳᵉˢʰ = zeros(T,size(x))
  for i = 1:K
    xᵗʰʳᵉˢʰ[:,:,i] = col2imDistinct( xᴸᴸᴿ[:,i,:], (Wy,Wz), (ny,nz) )
  end

  if randshift
    xᵗʰʳᵉˢʰ = circshift(xᵗʰʳᵉˢʰ, -1*shift_idx)
  end

  if !isempty(shape)
    xᵗʰʳᵉˢʰ = reshape( xᵗʰʳᵉˢʰ, prod(shape),floor( Int, length(xᵗʰʳᵉˢʰ)/prod(shape) ) )
  end

  return xᵗʰʳᵉˢʰ
end

"""
    normLLR(x::Vector{T}, λ::Float64; kargs...) where T

returns the value of the LLR-regularization term.
Arguments are the same is in `proxLLR!`
"""
function normLLR(x::Vector{T}, λ::Float64; shape::NTuple=[], L=1, blockSize::Array{Int64,1}=[2; 2], randshift::Bool=true, kargs...) where T

  N = prod(shape)
  K = floor(Int,length(x)/(N*L))
  normᴸᴸᴿ = 0.
  for i = 1:L
    normᴸᴸᴿ +=  blockNuclearNorm(x[(i-1)*N*K+1:i*N*K], shape; blockSize=blockSize, randshift=randshift, kargs...)
  end

  return λ*normᴸᴸᴿ
end

function blockNuclearNorm(x::Vector{T}, shape::Tuple; blockSize::Array{Int64,1}=[2; 2],
      randshift::Bool=true, kargs...) where T
    x = reshape( x, tuple( shape...,floor(Int64, length(x)/prod(shape)) ) )

    Wy = blockSize[1]
    Wz = blockSize[2]

    if randshift
      srand(1234)
      shift_idx = [rand(1:Wy) rand(1:Wz) 0]
      x = circshift(x, shift_idx)
    end

    ny, nz, K = size(x)

    # reshape into patches
    L = floor(Int,ny*nz/Wy/Wz) # number of patches, assumes that image dimensions are divisble by the blocksizes

    xᴸᴸᴿ = zeros(T,Wy*Wz,L,K)
    for i=1:K
      xᴸᴸᴿ[:,:,i] = im2colDistinct(x[:,:,i], (Wy,Wz))
    end
    xᴸᴸᴿ = permutedims(xᴸᴸᴿ,[1 3 2])

    # L1-norm of singular values
    normᴸᴸᴿ = 0.
    for i = 1:L
      SVDec = svd(xᴸᴸᴿ[:,:,i])
      normᴸᴸᴿ += norm(SVDec.S,1)
    end

    return normᴸᴸᴿ
end


using Distributions, Turing, Random, Parameters
ProjDir = @__DIR__
cd(ProjDir)
include("LNR.jl")

@model model(data, Nr) = begin
    μ = Array{Real,1}(undef,Nr)
    μ ~ [Normal(0,3)]
    s ~ Uniform(0, pi/2)
    σ = tan(s)
    #σ ~ Truncated(Cauchy(0, 1), 0, Inf)
    for i in 1:length(data)
        data[i] ~ LNR(μ, σ, 0.0)
    end
end

Random.seed!(343)
Nreps = 50
error_count = fill(0.0, Nreps)
Nr = 3
Nobs = 10

for i in 1:Nreps
    μ = rand(Normal(0, 3), Nr)
    σ = rand(Uniform(.2, 1))
    dist = LNR(μ=μ, σ=σ, ϕ=0.0)
    data = rand(dist, Nobs)
    chain = psample(model(data, Nr), NUTS(1000, .8), 2000, discard_adapt = false, progress = false)
    error_count[i] = sum(get(chain,:numerical_error)[1])
end

describe(error_count)


import Theta: ellipsoid_pointwise, ellipsoid_uniform

@testset "Ellipsoid" begin
    T = [1 0; 0 1];
    shift = [0; 0];
    @test Set(ellipsoid_pointwise(T, 2, shift)) == Set([[0,-1], [-1,0], [1,0], [0,1], [0,0]])
    @test issubset(ellipsoid_pointwise(T, 2, shift), ellipsoid_uniform(T, 2))
end


"""
    animate_triweight(x, μ, σ)

This function generates an animation depicting the probability density and the cumulative
    probability functions of the Normal distribution, as implemented in Distributions.jl.

```julia
x = collect(-5:0.01:5);
μ = collect(-4:0.5:4);
σ = collect(0.5:0.1:10);
animate_triweight(x, μ, σ)
```

"""

function animate_triweight(x, μ, σ, path = nothing)

    y = zeros(length(x), length(μ), length(σ), 2);

    for i = 1:length(μ), j = 1:length(σ)
        y[:, i, j, 1] .= pdf.(Triweight(μ[i], σ[j]), x);
        y[:, i, j, 2] .= cdf.(Triweight(μ[i], σ[j]), x);
    end

    ## animation
    indices = vcat(1:length(μ), reverse(1:length(μ)));

    anim = Animation()

    for i=indices
        p1 = Plots.plot(
            x, y[:, i, :, 1],
            legend=false,
            lw=3, ylims=(0, 2),
            xlabel="Sampling space", ylabel="Density",
            title="Tri-weight(μ, σ) - Probability density"
        )
        annotate!(median(x), 6-6/4, Plots.text(string.("μ = ", μ[i]), :black, :center, 12))

        p2 = Plots.plot(
            x, y[:, i, :, 2],
            legend=:outertopright, label = σ', legendtitle="σ",
            lw=3, ylims=(0, 1),
            xlabel="Sampling space", ylabel="Probability",
            title="Tri-weight(μ, σ) - Cumulative distribution"
        )

        plt = Plots.plot(p1, p2, size=(1000, 500))

        frame(anim, plt)
    end

    println("Saving GIF...")
    gif(anim, path)

end


@testset "Modified Cholesky decomposition" begin
    @testset "SPD matrix" begin
        dim = 10
        @testset for ntest in 1:15
            a = rand(test_rng, Float64, (dim, dim))
            a = a' * a
            achol = mcholesky!(copy(a))
            # For SPD matrices, must be the same as the usual Cholesky decomposition
            @test isapprox(achol.U' * achol.U, a; rtol = 1e-4)
        end
    end

    @testset "Generic symmetric matrix" begin
        dim = 10
        @testset for ntest in 1:15
            a = rand(test_rng, Float64, (dim, dim))
            a .+= a'
            achol = mcholesky!(copy(a))
            modcholprod = achol.U' * achol.U
            # checking that the decomposition differs from the original matrix
            # only in diagonal elements
            modcholprod .-= Diagonal(modcholprod)
            a .-= Diagonal(a)
            @test isapprox(UpperTriangular(modcholprod), UpperTriangular(a); rtol = 1e-4)
        end
    end
end


using Test
using ThreeBodyDecay

@testset "jp_str macro" begin
    @test jp"3+" == str2jp("3+")
    @test jp"3^+" == jp"3+"
    @test jp"1/2+" == jp(1//2, '+')
    @test jp"3/2-" == jp(3//2, '-')
    @test jp"3+" == jp(3, '+')
end

@testset "jp ⊗ jp" begin
    Swave = jp(1//2,'+') ⊗ jp(1,'-')
    Pwave = [sw ⊗ jp(1,'-') for sw in Swave]
    Dwave = [sw ⊗ jp(2,'+') for sw in Swave]
    #
    @test length(Swave) == 2
    @test length(vcat(Pwave...)) == 5
    @test length(Set(vcat(Pwave...))) == 3
end

let
    @test length(possible_ls(jp"3/2-", jp"3-"; jp=jp"1/2+")) == 4
    #
    lsLSv = possible_lsLS(1, jp"1/2+",
        [jp"1+", jp"1/2+", jp"0-", jp"1/2-"])
    @test size(lsLSv) == (1,2)
    #
    lsLSv = possible_lsLS(1, 1, '+',
        ThreeBodySpins(2,1,0; two_h0=1),
        ThreeBodyParities('+', '+', '-'; P0='-'))
    @test size(lsLSv) == (1,2)
end

@testset "x2 over2" begin
    @test x2(1//2) == 1
    @test x2(3) == 6
    @test x2((2,3//2)) == (4,3)
    #
    @test over2(3) == 3//2
    @test over2(2) == 1
    @test over2((2,3)) == (1,3//2)
end


@testset "Test interpolation of outgoing_translation_matrix" begin

    particle_medium = Acoustic(3; ρ=10.0, c=80.0);
    medium = Acoustic(3; ρ=1.0, c=1.0);

    R = 3.0
    r = 1.0

    separation_ratio = 1.02
    ka = 0.1
    k = ka / r
    vol_fraction = 0.3

    basis_order = 2

    ω = k * real(medium.c)

    s1 = Specie(
        particle_medium, Sphere(r);
        number_density = vol_fraction / volume(Sphere(r)),
        exclusion_distance = separation_ratio
    );
    species = [s1]
    material = Material(Sphere(R),species);

    a12 = 2.0 * minimum(outer_radius(s) * s.exclusion_distance for s in material.species)

    tol = 1e-2

    Uinter = outgoing_translation_matrix(ω, medium, material;  basis_order = basis_order, tol = tol)

    rs = LinRange(a12,2R-a12,13);
    φs = LinRange(-pi,pi,13);
    θs = LinRange(0.0,pi,13);

    r2x = radial_to_cartesian_coordinates

    error = [
        begin
            U1 = outgoing_translation_matrix(medium, basis_order, ω, r2x([r,θ,φ]))
            norm(Uinter(r2x([r,θ,φ])) - U1) / norm(U1)
        end
    for r in rs, θ in θs, φ in φs];

    @test maximum(error) < 5 * tol

    @test mean(error) < tol

end

@testset "Legendre polynomial approximation" begin

    P = Legendre{Float64}()

    polynomial_order = 25
    mesh_points = 3*(polynomial_order + 1)
    # mesh_points = (polynomial_order + 1)^2

    rs = LinRange(0,5,mesh_points)
    r2s = LinRange(0,5,4*mesh_points)

    r1_max = maximum(rs);
    rbars = 2.0 .* rs ./ r1_max .- (1.0);
    r2bars = 2.0 .* r2s ./ r1_max .- (1.0);

    data = cos.(rs) .* (1.0 .+ 4 .* exp.(-10 .* rs.^4))
    data2 = cos.(r2s) .* (1.0 .+ 4 .* exp.(-10 .* r2s.^4))
    # plot(rs,data)
    # plot!(r2s,data2)

    Pmat = P[rbars, 1:(polynomial_order + 1)];
    P2mat = P[r2bars, 1:(polynomial_order + 1)];

    # pls_arr * transpose(Pmat) ~ data
    # Pmat * transpose(pls_arr) ~ transpose(data)
    pls_arr = Pmat \ data

    ws = (0:polynomial_order) .+ 1/2
    σs = integration_scheme(rbars; scheme = :trapezoidal)
    tls_arr = ws .* (transpose(Pmat) * (σs .* data))

    λ = 1e-6
    cls_arr = inv(transpose(Pmat) * Pmat + λ*I) * transpose(Pmat) * data

    # norm(Pmat * pls_arr - data) / norm(data)
    # norm(Pmat * tls_arr - data) / norm(data)
    # norm(Pmat * cls_arr - data) / norm(data)
    @test norm(P2mat * pls_arr - data2) / norm(data2) < 5e-3
    @test norm(P2mat * tls_arr - data2) / norm(data2) < 0.6
    @test norm(P2mat * cls_arr - data2) / norm(data2) < 5e-3

# using Plots
#
#     scatter(rs,data)
#     plot!(rs,Pmat * pls_arr, linestyle = :dash)
#     plot!(r2s,P2mat * pls_arr, linestyle = :dash)
#     plot!(r2s,P2mat * tls_arr, linestyle = :dash)
#     plot!(r2s,P2mat * cls_arr, linestyle = :dash)

end


# For comparison with ffr-LFDFT

push!(LOAD_PATH, "../LF_common/")
using m_LF3d
include("../LF_common/sparse_LF3d.jl")

const ⊗ = kron

function test_main( NN::Array{Int64} )

  println("NN = ", NN)

  Nx = NN[1]
  Ny = NN[2]
  Nz = NN[3]
  Npoints = Nx*Ny*Nz

  AA = [0.0, 0.0, 0.0]
  BB = [16.0, 16.0, 16.0]

  # Initialize LF
  LF = init_LF3d_c( NN, AA, BB, verbose=true )

  ∇2 = get_Laplacian3d_kron( LF )

  nzval = ∇2.nzval
  rowval = ∇2.rowval
  for ip = 1:size(nzval)[1]
    @printf("%5d %5d %18.10f\n", ip, rowval[ip], nzval[ip])
  end

  v = 1.1*ones(Npoints)
  Av = ∇2 * v

  for ip = 1:Npoints
    @printf("%5d %18.10f %18.10f\n", ip, v[ip], Av[ip])
  end

end

test_main( [3, 2, 4])


#=
second-order cone
cuts are (1, -y/||y||_2) for (z, y) outside the cone
TODO extended formulation
=#
function get_init_cuts(cone::MOI.SecondOrderCone)
    cuts = Vector{Float64}[]
    dim = MOI.dimension(cone)

    # z is nonnegative
    cut = zeros(dim)
    cut[1] = 1.0
    push!(cuts, cut)

    # z is at least absolute value of each y_i
    for i in 2:dim
        cut = zeros(dim)
        cut[1] = 1.0
        cut[i] = 1.0
        push!(cuts, cut)
        cut = copy(cut)
        cut[i] = -1.0
        push!(cuts, cut)
    end

    return cuts
end

function get_sep_cuts(x::Vector{Float64}, cone::MOI.SecondOrderCone, tol::Float64)
    dim = MOI.dimension(cone)
    y_norm = norm(x[i] for i in 2:dim)

    if y_norm - x[1] > 0.1 * tol
        cut = similar(x)
        cut[1] = 1
        for i in 2:dim
            cut[i] = -x[i] / y_norm
        end

        if dot(x, cut) < -tol # TODO tolerance option
            return [cut]
        end
    end

    return Vector{Float64}[]
end


using packageDevTest1
using Test

packageDevTest1.test_dx(2,1)

@testset "packageDevTest1.jl" begin
    @test packageDevTest1.newton(sin,cos,0.01) ≈ 0.0
    @test packageDevTest1.test_fun(2,3) == 7
    @test packageDevTest1.test_fun(1,1) == 3
end


S = [:healthy, :in_situ, :invasive, :death]
A = [:wait, :test, :treat]
O = [:pos, :neg]
γ = 0.99
s0 = Deterministic(:healthy)
term = Set([:death])

function T(s, a, sp)
    if s == :healthy
        if sp == :in_situ
            return 0.02
        elseif sp == s
            return 0.98
        else
            return 0.0
        end
    elseif s == :in_situ
        if a == :treat
            if sp == :healthy
                return 0.60
            elseif sp == s
                return 0.40
            else
                return 0.0
            end
        else #a == :test || a == :wait
            if sp == :invasive
                return 0.10
            elseif sp == s
                return 0.90
            else
                return 0.0
            end
        end
    elseif s == :invasive
        if a == :treat
            if sp == :healthy
                return 0.20
            elseif sp == :death
                return 0.20
            elseif sp == s
                return 0.60
            else
                return 0.0
            end
        else
            if sp == :death
                return 0.60
            elseif sp == s
                return 0.40
            else
                return 0.0
            end
        end
    else # s == :death
        return 0.25
    end
end

function Z(a, sp, o)
    if a == :test
        if sp == :healthy
            if o == :pos
                return 0.05
            else
                return 0.95
            end
        elseif sp == :in_situ
            if o == :pos
                return 0.80
            else
                return 0.20
            end
        else
            if o == :pos
                return 1.0
            else
                return 0.0
            end
        end
    elseif a == :treat
        if sp == :in_situ || sp == :invasive
            if o == :pos
                return 1.0
            else
                return 0.0
            end
        else
            if o == :pos
                return 0.0
            else
                return 1.0
            end
        end
    else #a == :wait
        if o == :pos
            return 0.0
        else
            return 1.0
        end
    end
end

function R(s, a)
    if s == :death
        return 0.0
    elseif a == :wait
        return 1.0
    elseif a == :test
        return 0.80
    else
        return 0.10
    end
end

c = DiscreteExplicitPOMDP(S, A, O, T, Z, R, γ, s0, terminals=term);

function evalSolver(m)
    results = Dict{String, Float64}()
    for (key, solver) in ["SARSOP"=>SARSOPSolver()]
        policy = solve(solver, m)
        N = 10000
        rsum = 0.0
        for i in 1:N
            rsum += simulate(RolloutSimulator(max_steps=500), m, policy)
        end
        results[key] = rsum/N
    end;
    return results
end

results_c = evalSolver(c);
@show results_c["SARSOP"]
@test results_c["SARSOP"] >= 62.5


"""
    module BandTools

Utility functions for manipulating band structures. These need to be defined early in the code so other Modules have acccess to them.
"""
module BandTools

using LinearAlgebra
using SpecialFunctions

function gaussian(de, smearing=0.01)
    return 0.5 * (erfc.(de/smearing))
end

function gaussian_derivative(de, smearing=0.01)
    return (-0.5 / sqrt(pi)) * exp.(-1.0*(de/smearing).^2)
end

"""
    function calc_fermi(eigs, weights, nelec, smearing = 0.01)

calculate fermi energy using bisection
"""
function calc_fermi(eigs, weights, nelec, smearing = 0.01)

#    println("calc_fermi eigs ", eigs)

    efermi_max = maximum(eigs[:])+0.1
    efermi_min = minimum(eigs[:])-0.1

    efermi = (efermi_max+efermi_min)/2.0

    norm = sum(weights)

    n2= Int64(round(nelec/2))

#    println(size(eigs), " n2 $n2")

    if  abs(nelec - n2*2) < 1e-7
        if length(size(eigs)) == 2 && size(eigs)[2] > 1
            gap  = minimum(eigs[:,n2+1]) - maximum(eigs[:,n2])   #large gaps become numerically unstable with small smearing
        else
            gap  = minimum(eigs[n2+1]) - maximum(eigs[n2])   #large gaps become numerically unstable with small smearing
        end

        if gap > 0.05  #if gap is positive/large and we have an integer number of electrons
            smearing = max(smearing, gap / 10.0)
        end
        #    println("gap $gap smearing $smearing n2 $n2")

    end

    for iter = 1:35

        efermi = (efermi_max+efermi_min)/2.0

        occ = gaussian.(eigs.-efermi, smearing)

        n = sum(sum(occ, dims=2).*weights)/norm
#        println(iter, "  ", n,"  " ,  efermi)
        if n > (nelec/2.0)
            efermi_max = efermi
        else
            efermi_min = efermi
        end
    end

    return efermi

end

"""
    function band_energy(eigs, weights, nelec, smearing = 0.01; returnk=false, returnocc=false, returnef=false, returnboth=false)

Calculate band energy. Has options for additional return variables. Calculates fermi energy internally.
"""
function band_energy(eigs, weights, nelec, smearing = 0.01; returnk=false, returnocc=false, returnef=false, returnboth=false)

    efermi = calc_fermi(eigs, weights, nelec, smearing)
    norm = sum(weights)

    occ = gaussian.(eigs.-efermi, smearing)

    ek = sum(eigs.*occ, dims=2).*weights * 2.0 / norm
    #    energy = sum(sum(eigs.*occ, dims=2).*weights) * 2.0 / norm
    energy=sum(ek)

    if returnk
        return energy,ek
    elseif returnocc
        return energy, occ
    elseif returnef
        return energy, efermi
    elseif returnboth
        return energy, efermi, occ
    else
        return energy
    end

end

"""
    function smearing_energy(eigs, weights, efermi, smearing = 0.01)

Smearing contribution to total energy. If you don't take this int account in metals or
small gap semiconductors, your energy is not variational
in the correct way, and things like forces and stresses become wrong.
"""
function smearing_energy(eigs, weights, efermi, smearing = 0.01)

#    efermi = calc_fermi(eigs, weights, nelec, smearing)
    norm = sum(weights)

    dgauss = gaussian_derivative(eigs.-efermi, smearing)

    ek = smearing * sum(dgauss, dims=2).*weights * 2.0 / norm

    energy=sum(ek)

    return energy

end

end


using CartoonLike
points = create_circle(.4)
image, xLoc, yLoc = given_points(points, ninterp=100*100,ε=.001)
heatmap(xLoc, yLoc, image)

points = [-.5  -.9   0.0 .25   .5 .75;
          -.75 -.125 .9  .5  -.25 -.75]'
image, xLoc, yLoc = given_points(points, xSize=samples, ySize=samples)
heatmap(xLoc, yLoc, image')


#=
                       ------ MAIN INTERFACE -----
This code is a companion to the paper Surrogate Monte Carlo (SMC) -arxiv 2021. This is the main
code used to generate SMC data. However the default setting here are selected
for the code to run withing minutes (<10 min).
Authors: Andrei Da Silva & A.Christian Silva
e-mail: csilva@idatafactory.com
version on 10/2020
=#

using Random
using CSV
using DataFrames
using StatsBase

include("CDFJulia.jl")  # random numbers generation
include("utilities.jl") # opt functions
##
# Parameters
cgoal = 1e-3  # opt goal
nlag  = 50    # max lag for autocorrelation
n2lag = nlag  # max lag for autocorrelation of abs()
maxN  = 1e6   # max numb of terms in loop
##
# Data: downloaded from yahoo finance
sp = CSV.read("spyData.csv", DataFrame)
##
# Algorithm
# Step 1: Generate random numbers that have the same unconditional probability distribution
Random.seed!(28)
rs = getrnd(sp.R, length(sp.R))
#Step 2: Calculate the components of the optimization function
aR  = autocor(sp.R, 1:nlag)
bR  = autocor(abs.(sp.R), 1:n2lag)
cR  = crosscor(sp.R, abs.(sp.R), -nlag:nlag)
bR2 = autocor((sp.R).^2, 1:n2lag)
# Step 3: Calculate the optimization function. This is the starting point of the optimization
cmin = efun6(aR, bR, cR, bR2, rs, nlag, n2lag)
# Step 4: Perform optimization. Random permutation of rs should create final time series
# which minimizes the objective function and it is consequenty dynamically similar to the original
# time series
# ra: artificial time series
# E:  optimization function value
@time ra,E = annelOpt(aR,bR,cR,bR2,cmin,rs,nlag,n2lag,cgoal,maxN)
##
# How good is the opt ?
# Example:
#(1) Compare the autocorrelation of the abs(ra) with the "real" autocorrelation
B = autocor(abs.(ra), 1:n2lag)
plot(B)
plot!(bR)
#(2) Compare the crosscorrelation of the ra and abs(ra) with the measured crosscorrelation
C  = crosscor(ra, abs.(ra), -nlag:nlag)
plot(C)
plot!(cR)


using Random
using Distributions
using Plots

mutable struct Point
    label::Int64
    coordinates::Vector{Float64}
end

function DBSCAN(data, eps, min_cluster_size)
    cluster_counter = 0
    point_data = [Point(0, p) for p in eachcol(data)]

    for point in point_data
        if point.label == 0
            (neighbor_counter, neighbors) = find_neighbors(point_data, point, eps)
            if neighbor_counter < min_cluster_size
                point.label = -1
            else
                cluster_counter += 1
                point.label = cluster_counter
                for seed in neighbors
                    if seed.label < 1
                        seed.label = cluster_counter
                    end

                    (neighbor_counter, seed_neighbors) = find_neighbors(point_data, seed, eps)
                    if neighbor_counter >= min_cluster_size
                        union!(neighbors, seed_neighbors)
                    end
                end
            end
        end
    end

    return (point_data, cluster_counter)
end

function find_neighbors(data, point, eps)
    number_of_neighbors = 0
    neighbors = []

    for second_point in data
        if point != second_point
            if euclidian_distance_alt(point.coordinates, second_point.coordinates) < eps
                number_of_neighbors += 1
                append!(neighbors, [second_point])
            end
        end
    end

    return (number_of_neighbors, neighbors)
end

function euclidian_distance_alt(p1, p2)
    sum = 0

    for i in 1:length(p1)
        sum += abs2(p1[i] - p2[i])
    end

    return sqrt(sum)
end

function euclidian_distances(v1, v2)
    v1 = [v1[i, :] for i in 1:size(v1,1)]
    v2 = [v2[i, :] for i in 1:size(v2,1)]
    s = size(v1,1)
    d = zeros((length(v1[1]),length(v2[2])))

    for j in 1:length(v2[2])
        for k in 1:length(v1[1])
            summed = 0
            for i in 1:s
                summed += (v1[i][k] - v2[i][j])^2
            end
            d[k, j] = sqrt(summed)
        end
    end

    return d'
end

data = zeros(Float64, 3, 15000)
for i in 1:5000
    data[1, i] = rand(Normal(100, 10))
    data[2, i] = rand(Normal(40, 5))
    data[3, i] = rand(Normal(40, 10))
end

for i in 5000:10000
    data[1, i] = rand(Normal(20, 4))
    data[2, i] = rand(Normal(50, 6))
    data[3, i] = rand(Normal(50, 10))
end

for i in 10000:15000
    data[1, i] = rand(Normal(50, 5))
    data[2, i] = rand(Normal(90, 10))
    data[3, i] = rand(Normal(70, 8))
end

# display(data)
plotly()
scatter()
@time ( clustered, cluster_counter ) = DBSCAN(data, 4, 6)
for i in -1:cluster_counter
    cluster = [(x.coordinates[1], x.coordinates[2], x.coordinates[3]) for x in clustered if x.label == i]
    scatter!(cluster, markersize=1.5)
end
display(current())

# @time find_neighbors(data, data[:, 1], 10)


function allocate(datakind::Symbol, dtype::DataType, dims...)
    if datakind == :local
        return zeros(dtype, dims...)
    elseif datakind == :shared
        return SharedArray{dtype}(dims...)
    else
        ErrorException("Unknown kind: " * string(datakind)) |> throw
    end
end

#=
"""
This function checks if CFL criteria is satisfied which is required by Euler Forward Scheme. Explicitly,

 K Δt       1
------  <= ---
(Δz)^2      2

for every layer. This function returns true every layer is satisfied, returns false if any of the layers is not.

"""
function checkAllDiffusionStability(;
    Δzs:: AbstractArray{Float64, 1},
    K  :: Float64,
    Δt :: Float64,
)

    return all( Δzs .>= √(2.0 * K * Δt) )
end

function checkDiffusionStability(;
    Δz :: Float64,
    K  :: Float64,
    Δt :: Float64,
)

    return Δz >= √(2.0 * K * Δt)
end

function checkDiffusionStability(oc::OceanColumn; Δt)
    return checkAllDiffusionStability(Δzs=oc.Δzs, K=oc.K, Δt=Δt)
end

function minΔz(;
    K :: Float64,
    Δt:: Float64,
)
    return √(2.0 * K * Δt)
end

=#

function boundMLD(h_ML::Float64; h_ML_min::Float64, h_ML_max::Float64)
    return max(min(h_ML, h_ML_max), h_ML_min)
end

"""
    getTKE(fric_u)

# Description
This function returns the TKE (turbulent kinetic energy) `k = 0.5 * (v'^2)` of ML. This parameterization is given by Kim 1976: "A Generalized Bulk Model of the Oceanic Mixed Layer" in its equation (11)

"""
function getTKE(;
    fric_u :: Float64
)
    cm = max(3e-2, 3.0 * fric_u)
    return  0.5 * cm^2.0
end

function updateFLDO!(
        ocn :: Ocean,
    )

    @loop_hor ocn i j let
        OC_updateFLDO!(ocn, i, j)
    end
end

function OC_updateFLDO!(
        ocn :: Ocean,
        i   :: Integer,
        j   :: Integer,
    )
    ocn.FLDO[i, j] = getFLDO(zs=ocn.cols.zs[i, j], h_ML=ocn.h_ML[i, j], Nz=ocn.Nz[i, j])
end

"""

    Returns the FLDO. If mixed-layer depth is equal to the total depth
    of ocean column, -1 will be returned.

"""
function getFLDO(;
    zs   :: AbstractArray{Float64,1},
    h_ML :: Float64,
    Nz   :: Integer,
)
    for i = 1:Nz
        if h_ML < - zs[i+1]  # I don't use equality in order to avoid Δb = 0 during some initialization
            return i
        end
    end

    return -1
end

function getLayerFromDepth(;
    zs   :: AbstractArray{Float64,1},
    z    :: Float64,
    Nz   :: Integer,
)
    for i = 1:Nz
        if zs[i+1] < z  # I don't use equality in order to avoid Δb = 0 during some initialization
            return i
        end
    end

    return -1
end

function getWindStress(;
    u10::Float64
)

    return u10 * 1e-3 * ( (u10 < 25.0)
                    ? 2.7 + 0.142 * u10 + 0.0764 * u10^2.0
                    : u10 * (2.16 + 0.5406 * (1.0 - exp(- (u10 - 25.0) / 7.5)))
    )

end

function getFricU(;
    ua::Float64
)
    return √(getWindStress(u10=ua) / ρ)
end

function OC_updateB!(
    ocn :: Ocean,
    i   :: Integer,
    j   :: Integer,
)

    ocn.b_ML[i, j] = TS2b(ocn.T_ML[i, j], ocn.S_ML[i, j])
    for k=1:ocn.Nz[i, j]
        ocn.bs[k, i, j] = TS2b(ocn.Ts[k, i, j], ocn.Ss[k, i, j])
    end

end

function updateB!(ocn::Ocean)

    @loop_hor ocn i j let
        OC_updateB!(ocn, i, j)
    end

end


module RandomizedPreconditioners

using LinearAlgebra
using Random

include("utils.jl")
include("rangefinder.jl")
include("sketch.jl")
include("preconditioner.jl")

export NystromSketch, EigenSketch, RandomizedSVD, adaptive_sketch
export NystromPreconditioner, NystromPreconditionerInverse

end


using LinearAlgebra, Random
using Statistics, Distributions

"""
@article{Lewandowski2009a,
    Author = {Daniel Lewandowski and Dorota Kurowicka and Harry Joe},
    Journal = {Journal of Multivariate Analysis},
    Number = {9},
    Pages = {1989 - 2001},
    Title = {Generating random correlation matrices based on vines and extended onion method},
    Volume = {100},
    Year = {2009}}
https://doi.org/10.1016/j.jmva.2009.04.008
"""
function onion(n; η = 1.0)
    S = Symmetric(ones(n, n))
    β = η + (n - 2) / 2
    S.data[1, 2] = 2 * rand(Beta(β, β)) - 1
    for k = 2:n-1
        β -= 0.5
        y = rand(Beta(k / 2, β))
        u = normalize!(Random.randn(k))
        w = sqrt(y) * u
        F = cholesky(Symmetric(S[1:k, 1:k]))
        S.data[1:k, k+1] = F.L * w
    end
    return S
end

function randncm(n; seed = 0, γ = 0.0, p = 0.5, gaussian_noise = false)
    rng = Random.seed!(seed)

    # Random target matrix
    U = onion(n)

    # Random sparse H
    if gaussian_noise
        # Hij ∈ [1, 10] for Hij ≠ 0
        Htmp = [rand() < p ? 1.0 + 9.0 * rand() : 0.0 for i = 1:n, j = 1:n]
    else
        Htmp = [rand() < p ? rand() : 0.0 for i = 1:n, j = 1:n]
    end
    H = Symmetric(triu(Htmp, 1) + I)

    # Random noise
    if gaussian_noise
        # G = U + E./H
        G = copy(U)
        for j = 1:n
            for i = 1:j-1
                Hij = H.data[i, j]
                if Hij > 0.0
                    G.data[i, j] += γ * Random.randn() / Hij
                else
                    G.data[i, j] = 0.0
                end
            end
        end
    else
        Etmp = 2 * rand(n, n) .- 1
        E = Symmetric(triu(Etmp, 1) + I)
        Gtmp = (1 - γ) .* U .+ γ .* E
        G = Symmetric(triu(Gtmp, 1) + I)
    end

    Random.seed!(rng)

    return U, G, H
end


"""
    bvp1.py solves the 1D Poisson equation using FEniCS
    Copyright (C) 2013  Greg von Winckel
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.
    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.
    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
   Solve the 1D Poisson equation -u''=-1 with u(0)=0, u'(1)=1
   The exact solution is x^2/2
"""
module bvp1

using FenicsPy

# Number of elements
nel = 20

# Left end point
xmin = 0

# Right end point
xmax = 1

# Polynomial order of trial/test functions
p = 2

# Create mesh and function space
mesh = IntervalMesh(nel,xmin,xmax)

# Define function space for this mesh using Continuous Galerkin
# (Lagrange) functions of order p on each element
V = FunctionSpace(mesh,"CG",p)

# Define boundary boundary values
u0 = Expression("x[0]", degree=1)

# This imposes a Dirichlet condition at the point x=0
dirichlet_boundary = "on_boundary && abs(x[0])<1e-14"

# Enforce u=u0 at x=0
bc = DirichletBC(V,u0,dirichlet_boundary)

# Define variational problem
u = TrialFunction(V)
v = TestFunction(V)
f = Constant(-1)
g = Constant(1)
a = inner(grad(u),grad(v))*dx
L = f*v*dx+g*v*ds

# Compute solution
u = FeFunction(V)
solve(a == L, u, bc)

# plot solution
#plot(u)

# Dump solution to file in VTK format
file = File("bvp1/bvp1.pvd")
file << u

end # module bvp1


# This file is part of IntegerSequences.
# Copyright Peter Luschny. License is MIT.

(@__DIR__) ∉ LOAD_PATH && push!(LOAD_PATH, (@__DIR__))

module CarmichaelNumbers

using Nemo, NumberTheory, Counts

export ModuleCarmichaelNumbers
export isCarmichael, I002997, F002997, L002997
export isweakCarmichael, I225498, F225498, L225498

"""

* isCarmichael, I002997, F002997, L002997
* isweakCarmichael, I225498, F225498, L225498
"""
const ModuleCarmichaelNumbers = ""

"""

Is ``n`` a Carmichael (or Šimerka) number?
"""
function isCarmichael(n)
    (n == 1 || isEven(n) || isPrime(n)) && return false
    for f ∈ Factors(n)
        (f[2] > 1 || (n - 1) % (f[1] - 1) != 0) && return false
    end
    return true
end

"""

Iterate over the first ``n`` Carmichael (or Šimerka) numbers.
"""
I002997(n) = takeFirst(isCarmichael, n)

"""

Iterate over the Carmichael (or Šimerka) numbers which do not exceed ``n``.
"""
F002997(n) = filter(isCarmichael, 1:n)

"""

Return the first ``n`` Carmichael (or Šimerka) numbers in an array.
"""
L002997(n) = collect(I002997(n))

"""

Is ``n`` a weak Carmichael number?
"""
function isweakCarmichael(n)
    (n == 1 || isEven(n) || isPrime(n)) && return false
    for f ∈ Factors(n)
        (n - 1) % (f[1] - 1) != 0 && return false
    end
    return true
end

"""

Iterate over the first ``n`` weak Carmichael numbers.
"""
I225498(n) = takeFirst(isweakCarmichael, n)

"""

Iterate over the weak Carmichael numbers which do not exceed ``n``.
"""
F225498(n) = filter(isweakCarmichael, 1:n)

"""

Return the first ``n`` weak Carmichael numbers in an array.
"""
L225498(n) = collect(I225498(n))

#START-TEST-########################################################

using Test, SeqTests

function test()
    @testset "Carmichael" begin
        @test !isCarmichael(560)
        @test isCarmichael(561)
        @test !isCarmichael(563)
        @test isweakCarmichael(561)
        @test !isweakCarmichael(563)
        @test isweakCarmichael(625)

        if data_installed()
            L = [L002997, L225498]
            SeqTest(L, 'L')
        end

    end
end

function demo()
    for n ∈ 1:30000
        isCarmichael(n) && println(n)
    end

    println()
    for v ∈ I225498(10)
        println(v)
    end
end

"""

@time L = L002997(30)
0.967192 seconds (4.01 M allocations: 248.863 MiB, 28.97% gc time)
"""
function perf()
    @time L = L002997(30)
end

function main()
    test()
    demo()
    perf()
end

main()

end # module


inpt = """AlphaCentauri to Snowdin = 66
AlphaCentauri to Tambi = 28
AlphaCentauri to Faerun = 60
AlphaCentauri to Norrath = 34
AlphaCentauri to Straylight = 34
AlphaCentauri to Tristram = 3
AlphaCentauri to Arbre = 108
Snowdin to Tambi = 22
Snowdin to Faerun = 12
Snowdin to Norrath = 91
Snowdin to Straylight = 121
Snowdin to Tristram = 111
Snowdin to Arbre = 71
Tambi to Faerun = 39
Tambi to Norrath = 113
Tambi to Straylight = 130
Tambi to Tristram = 35
Tambi to Arbre = 40
Faerun to Norrath = 63
Faerun to Straylight = 21
Faerun to Tristram = 57
Faerun to Arbre = 83
Norrath to Straylight = 9
Norrath to Tristram = 50
Norrath to Arbre = 60
Straylight to Tristram = 27
Straylight to Arbre = 81
Tristram to Arbre = 90"""
places = Vector{ASCIIString}()
for line in split(inpt, '\n')
	p1 = split(line)[1]
	p2 = split(line)[3]
	if !(p1 in places)
		push!(places, p1)
	end
	if !(p2 in places)
		push!(places, p2)
	end
end

distances = zeros(Int64, (length(places), length(places)))
println(inpt)
println(places)
for line in split(inpt, '\n')
	place1 = split(line)[1]
	place2 = split(line)[3]
	dist = parse(Int, split(line)[end])
	println(place1, place2, dist)
	i = findfirst(places, place1)
	j = findfirst(places, place2)
	distances[i,j] = dist
	distances[j,i] = dist
end
println(distances)

shortest = (0, places)
for p in permutations(places)
	sum = 0
	for city in 2:length(p)
		i = findfirst(places, p[city])
		j = findfirst(places, p[city-1])
		sum += distances[i, j]
	end
	if sum > shortest[1]
		shortest = (sum, p)
	end
end

println(shortest)


v = rand(10)

a, b = 4, 8
c, d = 5, 9

M = randn(a, b)
N = rand(c, d)

p = rand(1:a, 6)
q = rand(1:c, 6)

r = rand(1:b, 10)
t = rand(1:d, 10)

kronprod = N ⊗ M
ikp = kronprod[p,q,r,t]

subsystem = kron(N, M)[a * (q .- 1) .+ p, b * (t .- 1) .+ r]
@test subsystem ≈ collect(ikp)

# result shortcut
#u = genvectrick(M, N, v, p, q, r, t)
u = ikp * v
# result naive
unaive = kron(N, M)[a * (q .- 1) .+ p, b * (t .- 1) .+ r] * v
@test all(u .≈ unaive)


# ---
# title: 1304. Find N Unique Integers Sum up to Zero
# id: problem1304
# author: Tian Jun
# date: 2020-10-31
# difficulty: Easy
# categories: Array
# link: <https://leetcode.com/problems/find-n-unique-integers-sum-up-to-zero/description/>
# hidden: true
# ---
#
# Given an integer `n`, return **any** array containing `n` **unique**  integers
# such that they add up to 0.
#
#
#
# **Example 1:**
#
#
#
#     Input: n = 5
#     Output: [-7,-1,1,3,4]
#     Explanation: These arrays also are accepted [-5,-1,1,2,3] , [-3,-1,2,-2,4].
#
#
# **Example 2:**
#
#
#
#     Input: n = 3
#     Output: [-1,0,1]
#
#
# **Example 3:**
#
#
#
#     Input: n = 1
#     Output: [0]
#
#
#
#
# **Constraints:**
#
#   * `1 <= n <= 1000`
#
#
## @lc code=start
using LeetCode

## add your code here:
## @lc code=end


function drawlimitset(generators::Vector{F}, z0::Number=0;
  iterations::Int=5) where F <: AbstractComplexInvertibleFunction

  SDDGraphics.newdrawing()

  #SDDGraphics.updatecolorarray(iterations)

  genset = GeneratorsSet(generators)

  SDDGraphics.updatecolorarray(size(genset))

  function procedure(zn::Number, genindex::Int, level::Int, levels::Int)
    if level == levels && SDDGraphics.insiderectregion(zn)
      SDDGraphics.color(genindex)
      SDDGraphics.drawpoint(zn)
    end
  end

  traversegroupfirstdepth(generators, iterations, z0, procedure)

  SDDGraphics.drawing()
end

const drawΛ = drawlimitset

function drawfixedpointslimitset(generators::Vector{F};
  iterations::Int=5) where F <: AbstractMobiusTransformation

  SDDGraphics.newdrawing()

  SDDGraphics.updatecolorarray(iterations)

  #genset = GeneratorSet(generators)

  #SDDGraphics.updatecolorarray(length(genset))

  function procedure(f::AbstractMobiusTransformation, genindex::Int, level::Int, levels::Int)
    #z1, z2 = fixedpoints(f)
    for z in fixedpoints(f) #[z1,z2]
      if SDDGraphics.insiderectregion(z)
        SDDGraphics.color(level)
        SDDGraphics.drawpoint(z)
      end
    end
  end

  traversegroupfirstdepthcomp(generators, iterations, procedure)

  SDDGraphics.drawing()
end

const drawfixedpointsΛ = drawfixedpointslimitset

function drawchaosgamelimitset(generators::Vector{F}, z0::Number=0;
  preiterations::Int=100, iterations::Int=100) where F <: AbstractComplexInvertibleFunction

  SDDGraphics.newdrawing()

  #SDDGraphics.updatecolorarray(iterations)

  genset = GeneratorsSet(generators)

  SDDGraphics.updatecolorarray(size(genset))

  previndex = 0

  function randT(z::Number)
    k = rand(1:size(genset))
    invindex = inverseindex(genset,k)
    while previndex == invindex
      k = rand(1:size(genset))
      invindex = inverseindex(genset,k)
    end
    previndex = k
    genset[k](z)
  end

  zn = z0

  for n in 1:preiterations
    zn = randT(zn)
  end

  for n in 1:iterations
    zn = randT(zn)
    if SDDGraphics.insiderectregion(zn)
      SDDGraphics.color(previndex)
      SDDGraphics.drawpoint(zn)
    end
  end

  SDDGraphics.drawing()
end

const drawchaosgameΛ = drawchaosgamelimitset

function drawtrappedpointslimitset(generators::Vector{F};
  distance::Function=(z1::Number,z2::Number) -> abs2(z1-z2),
  maxiterations::Int=50) where F <: AbstractMobiusTransformation

  # Veryfying if graphics backend supports functions
  SDDGraphics.supported(:drawpixel)

  # Verifying functions
  @assert typeof(distance(1., im)) <: Real

  SDDGraphics.newdrawing()

  SDDGraphics.updatecolorarray(maxiterations)

  finitefixedpoints = ComplexF64[]
  mobiusT = AbstractMobiusTransformation[]
  parabolic = Bool[]

  R = 0

  # ToDo: Handle parabolic transformations with fixed point at Inf.

  for f in generators
    z1, z2 = fixedpoints(f)
    if kind(f) == :parabolic
      if isfinite(z1)
        push!(finitefixedpoints, z1)
        push!(mobiusT, f)
        push!(parabolic, true)
        if R < abs2(z1)
          R = abs2(z1)
        end
      end
    else
      fder = derivative(f)
      for z in [z1,z2]
        if isfinite(z)
          push!(finitefixedpoints, z)
          push!(parabolic, false)
          if abs2(fder(z)) >= 1 # loxodromic or elliptic
            push!(mobiusT, f)
          else # loxodromic
            push!(mobiusT, inverse(f))
          end # if |f'(z)|>=1

          if R < abs2(z)
            R = abs2(z)
          end
        end # if |z|<Inf
      end # for z in [z1,z2]
    end # if else
  end # for generators

  #R *= √2

  # Quasi-PieceWise Transformation, using a Voronoi partition over fixed points
  function voronoiqpwt(z::Number)
    k = 1
    dist = Inf

    for n in 1:length(finitefixedpoints)
      distzfix = distance(z,finitefixedpoints[n])
      if dist > distzfix
        dist = distzfix
        k = n
      end
    end

    if parabolic[k]
      z1 = mobiusT[k](z)
      z2 = inverse(mobiusT[k])(z)
      if distance(z1,finitefixedpoints[k]) >= distance(z2,finitefixedpoints[k])
        return z1
      else
        return z2
      end
    end

    mobiusT[k](z)
  end

  @sweeprectregion SDDGraphics.xlims() SDDGraphics.ylims() SDDGraphics.canvassize() begin
      zn = complex(x,y)
      escapetime = maxiterations
      for n in 1:maxiterations
          if abs2(zn) > R
              escapetime = n
              break
          end # if hasescaped
          zn = voronoiqpwt(zn)
      end # for n maxiterations
      SDDGraphics.color(escapetime)
      SDDGraphics.drawpixel(i,j)
  end # Implemented algorithm

  SDDGraphics.drawing()

end

const drawtrappedpointsΛ = drawtrappedpointslimitset


const ManyMatrices = AbstractVecOrTuple{AbstractMatrix}

@inline inner(A,B) = dot(A,B)

@inline inner(A::NTuple{2,Operator}) = @inbounds inner(A[1],A[2])

@inline inner(A::Tuple{Vararg{Operator}},B::Operator) = inner_cs(data.(A),data(B))

function inner_cs(A,B::AbstractMatrix,s::BitSet=BitSet())
    # calculate |trace(A'*B)|cs (coherent subspaces)
    # this inner product cares only about relative phase between the subspaces contained in s
    x,y = _inner_cs_1(A,B,s)
    return _inner_cs_2(x,y)
end

function inner_cs_grad(Pj,JXj::AbstractMatrix,x,y,s::BitSet=BitSet())
    # calculate the partial derivative of |trace(A'*B)|cs (coherent subspaces)
    w,z = _inner_cs_1(Pj,JXj,s)
    res = real(w*x)
    @inbounds for i = 1:length(z)
        res += real(z[i]*y[i])
        res += real(w*normalize(x))*abs(y[i]) + real(z[i]*normalize(y[i]))*abs(x)
        for j = 1:i-1
            res += real(z[i]*normalize(y[i]))*abs(y[j]) + real(z[j]*normalize(y[j]))*abs(y[i])
        end
    end
    return res/_inner_cs_2(x,y)
end

# sub-routines

function _inner_cs_1(A::AbstractMatrix{T},B::AbstractMatrix{S},s::BitSet) where {T,S}
    m, n = size(A)
    size(B) == (m,n) || throw(DimensionMismatch("matrices must have the same dimensions"))
    x = zero(promote_type(T,S))
    y = Vector{typeof(x)}(undef,n-length(s))
    t = 1
    @inbounds for j = 1:n
        a = zero(x)
        for i = 1:m
            a += A[i,j]' * B[i,j]
        end
        if j in s
            x += a
        else
            y[t] = a
            t += 1
        end
    end
    return x, y
end

function _inner_cs_1(A::ManyMatrices,B::AbstractMatrix,::BitSet)
    # calculate |trace(A'*B)|cs (coherent subspaces)
    # this inner product cares only about relative phase between the subspaces contained in s
    x = inner.(A,(B,))
    return first(x), tail(x)
end

function _inner_cs_1(A::AbstractMatrix,B::ManyMatrices,::BitSet)
    # calculate |trace(A'*B)|cs (coherent subspaces)
    # this inner product cares only about relative phase between the subspaces contained in s
    x = inner.((A,),B)
    return first(x), tail(x)
end

function _inner_cs_2(x::Number,y::Union{Vector{<:Number},NTuple{M,<:Number} where M})
    res = abs2(x)
    @inbounds for i = 1:length(y)
        res += abs2(y[i])
        res += 2*abs(y[i])*abs(x)
        for j = 1:i-1
            res += 2*abs(y[i])*abs(y[j])
        end
    end
    return sqrt(res)
end

f_cs_2(x::Number,y::Union{Vector{<:Number},NTuple{M,<:Number} where M}) = sum(abs,y)+abs(x)


@doc doc"""
Résolution des problèmes de minimisation sous contraintes d'égalités

# Syntaxe
```julia
Lagrangien_Augmente(algo,fonc,contrainte,gradfonc,hessfonc,grad_contrainte,
			hess_contrainte,x0,options)
```

# Entrées
  * **algo** 		   : (String) l'algorithme sans contraintes à utiliser:
    - **"newton"**  : pour l'algorithme de Newton
    - **"cauchy"**  : pour le pas de Cauchy
    - **"gct"**     : pour le gradient conjugué tronqué
  * **fonc** 		   : (Function) la fonction à minimiser
  * **contrainte**	   : (Function) la contrainte [x est dans le domaine des contraintes ssi ``c(x)=0``]
  * **gradfonc**       : (Function) le gradient de la fonction
  * **hessfonc** 	   : (Function) la hessienne de la fonction
  * **grad_contrainte** : (Function) le gradient de la contrainte
  * **hess_contrainte** : (Function) la hessienne de la contrainte
  * **x0** 			   : (Array{Float,1}) la première composante du point de départ du Lagrangien
  * **options**		   : (Array{Float,1})
    1. **epsilon** 	   : utilisé dans les critères d'arrêt
    2. **tol**         : la tolérance utilisée dans les critères d'arrêt
    3. **itermax** 	   : nombre maximal d'itération dans la boucle principale
    4. **lambda0**	   : la deuxième composante du point de départ du Lagrangien
    5. **mu0,tho** 	   : valeurs initiales des variables de l'algorithme

# Sorties
* **xmin**		   : (Array{Float,1}) une approximation de la solution du problème avec contraintes
* **fxmin** 	   : (Float) ``f(x_{min})``
* **flag**		   : (Integer) indicateur du déroulement de l'algorithme
   - **0**    : convergence
   - **1**    : nombre maximal d'itération atteint
   - **(-1)** : une erreur s'est produite
* **niters** 	   : (Integer) nombre d'itérations réalisées

# Exemple d'appel
```julia
using LinearAlgebra
f(x)=100*(x[2]-x[1]^2)^2+(1-x[1])^2
gradf(x)=[-400*x[1]*(x[2]-x[1]^2)-2*(1-x[1]) ; 200*(x[2]-x[1]^2)]
hessf(x)=[-400*(x[2]-3*x[1]^2)+2  -400*x[1];-400*x[1]  200]
algo = "gct" # ou newton|gct
x0 = [1; 0]
options = []
contrainte(x) =  (x[1]^2) + (x[2]^2) -1.5
grad_contrainte(x) = [2*x[1] ;2*x[2]]
hess_contrainte(x) = [2 0;0 2]
output = Lagrangien_Augmente(algo,f,contrainte,gradf,hessf,grad_contrainte,hess_contrainte,x0,options)
```
"""
function Lagrangien_Augmente(algo,fonc::Function,contrainte::Function,gradfonc::Function,
	hessfonc::Function,grad_contrainte::Function,hess_contrainte::Function,x0,options)

	if options == []
		epsilon = 1e-8
		tol = 1e-5
		itermax = 50
		lambda0 = 2
		mu0 = 100
		tho = 2
	else
		epsilon = options[1]
		tol = options[2]
		itermax = options[3]
		lambda0 = options[4]
		mu0 = options[5]
		tho = options[6]
	end

	#Definitions des données
    n = length(x0)
    xmin = zeros(n)
	fxmin = 0
	flag = -2
	iter = 0
	alpha = 0.1
	beta = 0.9
	epsilon0 = 1/mu0
	ethachapeau = 0.1258925

	#Definitions des variables
	etha_k = ethachapeau / (mu0^alpha)
	x_k = x0
	mu_k = mu0
	lambda_k = lambda0
	epsilon_k = epsilon0

	La(x) = fonc(x) + transpose(lambda_k) * contrainte(x) + (mu_k/2) * norm(contrainte(x))^2
	grad_La(x) = gradfonc(x) + (lambda_k)'*grad_contrainte(x) + (mu_k)*grad_contrainte(x)*contrainte(x)
	Hess_La(x) = hessfonc(x) + (lambda_k)'*hess_contrainte(x) + (mu_k)*hess_contrainte(x)*contrainte(x) + (mu_k)*grad_contrainte(x)*grad_contrainte(x)'

	grad_Lac(x,y) = gradfonc(x) + (y)'*grad_contrainte(x)

while flag == -2

		iter += 1
		iter_while = 0

		if algo == "newton"
			while norm(grad_La(x_k)) > epsilon_k && iter_while <= itermax
				x_k,f_minN,flagN,iterN = Algorithme_De_Newton(La,grad_La,Hess_La,x_k,[])
				iter_while += 1
        	end

		elseif algo == "cauchy"
			while norm(grad_La(x_k)) > epsilon_k && iter_while <= itermax
				x_k,f_minC,flagC,iterC = Regions_De_Confiance("cauchy", La,grad_La,Hess_La,x_k,[])
				iter_while += 1
			end

		elseif algo == "gct"
			while norm(grad_La(x_k)) > epsilon_k && iter_while <= itermax
				x_k,f_minG,flagG,iterG = Regions_De_Confiance("gct", La,grad_La,Hess_La,x_k,[])
				iter_while += 1
			end
		end

        if iter_while == itermax
			flag = -1
		elseif  norm(grad_Lac(x_k, lambda_k)) <= max(tol*norm(grad_Lac(x0, lambda0)), epsilon) && norm(contrainte(x_k)) <= max(tol*norm(contrainte(x0)), epsilon)
			flag = 0
		elseif iter == itermax
			flag = 1

		else
				if norm(contrainte(x_k)) <= etha_k
					lambda_k += mu_k*contrainte(x_k)
					epsilon_k = epsilon_k/mu_k
					etha_k = etha_k/(mu_k^(beta))
	  			else
		  			mu_k = tho*mu_k
		  			epsilon_k = epsilon0/mu_k
		  			etha_k = ethachapeau/(mu_k^(alpha))
	  			end

		end
end

xmin = x_k
fxmin = fonc(x_k)

	return xmin,fxmin,flag,iter
end