###########################################################
#
#   (non-weighted) sampling
#
###########################################################

### Algorithms for sampling with replacement

# Direct sampling
#
#   for each i, draw an index from 1:length(a), and
#   pick the corresponding element from a.
#
#   x[i] and x[j] are independent, when i != j
#
function direct_sample!(a::UnitRange, x::AbstractArray)
    s = RandIntSampler(length(a))
    b = a[1] - 1
    if b == 0
        for i = 1:length(x)
            @inbounds x[i] = rand(s)
        end
    else
        for i = 1:length(x)
            @inbounds x[i] = b + rand(s)
        end
    end
    return x
end

function direct_sample!(a::AbstractArray, x::AbstractArray)
    s = RandIntSampler(length(a))
    for i = 1:length(x)
        @inbounds x[i] = a[rand(s)]
    end
    return x
end

# Expanded Multinomial sampling
#
#   for each element in a, we draw the number of its
#   occurrences, and fill this element to x for
#   this number of times.
#
function xmultinom_sample!(a::AbstractArray, x::AbstractArray)
    n = length(a)
    k = length(x)
    offset = 0
    i = 1
    while offset < k
        rk = k - offset
        if i == n
            @inbounds ai = a[i]
            for j = 1:rk
                @inbounds x[offset + j] = ai
            end
            offset = k
        else
            m = rand_binom(rk, 1.0 / (n - i + 1))
            if m > 0
                @inbounds ai = a[i]
                for j = 1:m
                    @inbounds x[offset + j] = a[i]
                end
                offset += m
            end
            i += 1
        end
    end
    x
end

### draw a pair of distinct integers in [1:n]

function samplepair(n::Int)
    i1 = randi(n)
    i2 = randi(n-1)
    return (i1, ifelse(i2 == i1, n, i2))
end

function samplepair(a::AbstractArray)
    i1, i2 = samplepair(length(a))
    return a[i1], a[i2]
end

### Algorithm for sampling without replacement

# Knuth's Algorithm S
#
#   Reference: The Art of Computer Programming, Vol 2, 3.4.2, p.142
#
function knuths_sample!(a::AbstractArray, x::AbstractArray; initshuffle::Bool=true)
    n = length(a)
    k = length(x)
    k <= n || error("length(x) should not exceed length(a)")

    # initialize
    for i = 1:k
        @inbounds x[i] = a[i]
    end
    if initshuffle
        @inbounds for j = 1:k
            l = randi(j, k)
            if l != j
                t = x[j]
                x[j] = x[l]
                x[l] = t
            end
        end
    end

    # scan remaining
    s = RandIntSampler(k)
    for i = k+1:n
        if rand() * i < k  # keep it with probability k / i
            @inbounds x[rand(s)] = a[i]
        end
    end
    return x
end

# Fisher-Yates sampling
#
#   create an array of index inds = [1:n]
#
#   for i = 1:k
#       swap inds[i] with a random one in inds[i:n]
#       set x[i] = a[inds[i]]
#   end
#
#   O(n) for initialization + O(k) for random shuffling
#
function fisher_yates_sample!(a::AbstractArray, x::AbstractArray)
    n = length(a)
    k = length(x)
    k <= n || error("length(x) should not exceed length(a)")

    inds = Array(Int, n)
    for i = 1:n
        @inbounds inds[i] = i
    end

    @inbounds for i = 1:k
        j = randi(i, n)
        t = inds[j]
        inds[j] = inds[i]
        inds[i] = t
        x[i] = a[t]
    end
    return x
end

# Self-avoid sampling
#
#   Use a set to maintain the index that has been sampled,
#
#   each time draw a new index, if the index has already
#   been sampled, redraw until it draws an unsampled one.
#
function self_avoid_sample!(a::AbstractArray, x::AbstractArray)
    n = length(a)
    k = length(x)
    k <= n || error("length(x) should not exceed length(a)")

    s = Set{Int}()
    sizehint(s, k)
    rgen = RandIntSampler(n)

    # first one
    idx = rand(rgen)
    x[1] = a[idx]
    push!(s, idx)

    # remaining
    for i = 2:k
        idx = rand(rgen)
        while idx in s
            idx = rand(rgen)
        end
        x[i] = a[idx]
        push!(s, idx)
    end
    return x
end

# Random subsequence sampling
#
#  References:
#
#   Jeffrey Scott Vitter.
#   "Faster Methods for Random Sampling"
#   Communications of the ACM, 27 (7), July 1984
#
#
# This paper presents three algorithms, respectively
# named Algorithm A, C, and D.
#
# These algorithms are for sequential sampling, and the
# outputs are ordered. They are implemented below
#

## Algorithm A (page 714)
#
#  Require O(n) random numbers.
#
function seqsample_a!(a::AbstractArray, x::AbstractArray)
    n = length(a)
    k = length(x)
    k <= n || error("length(x) should not exceed length(a)")

    i = 0
    j = 0
    while k > 1
        u = rand()
        q = (n - k) / n
        while q > u  # skip
            i += 1
            n -= 1
            q *= (n - k) / n
        end
        @inbounds x[j+=1] = a[i+=1]
        n -= 1
        k -= 1
    end

    if k > 0  # checking k > 0 is necessary: x can be empty
        s = itrunc(n * rand())
        x[j+1] = a[i+(s+1)]
    end
    return x
end

## Algorithm C (page 715)
#
#  Require O(k^2) random numbers
#
function seqsample_c!(a::AbstractArray, x::AbstractArray)
    n = length(a)
    k = length(x)
    k <= n || error("length(x) should not exceed length(a)")

    i = 0
    j = 0
    while k > 1
        l = n - k + 1
        minv = l
        u = n
        while u >= l
            v = u * rand()
            if v < minv
                minv = v
            end
            u -= 1
        end
        s = itrunc(minv) + 1
        x[j+=1] = a[i+=s]
        n -= s
        k -= 1
    end

    if k > 0
        s = itrunc(n * rand())
        x[j+1] = a[i+(s+1)]
    end
    return x
end

## TODO: implement Algorithm D (page 716 - 717)

### Interface functions (poly-algorithms)

sample(a::AbstractArray) = a[randi(length(a))]

function sample!(a::AbstractArray, x::AbstractArray; replace::Bool=true, ordered::Bool=false)
    n = length(a)
    k = length(x)
    k == 0 && return x

    if replace  # with replacement
        if ordered
            if k > 10 * n
                xmultinom_sample!(a, x)
            else
                sort!(direct_sample!(a, x))
            end
        else
            direct_sample!(a, x)
        end

    else  # without replacement
        k <= n || error("Cannot draw more samples without replacement.")

        if ordered
            if n > 10 * k * k
                seqsample_c!(a, x)
            else
                seqsample_a!(a, x)
            end
        else
            if k == 1
                @inbounds x[1] = sample(a)
            elseif k == 2
                @inbounds (x[1], x[2]) = samplepair(a)
            elseif n < k * 24
                fisher_yates_sample!(a, x)
            else
                self_avoid_sample!(a, x)
            end
        end
    end
    return x
end

function sample{T}(a::AbstractArray{T}, n::Integer; replace::Bool=true, ordered::Bool=false)
    sample!(a, Array(T, n); replace=replace, ordered=ordered)
end

function sample{T}(a::AbstractArray{T}, dims::Dims; replace::Bool=true, ordered::Bool=false)
    sample!(a, Array(T, dims); replace=replace, ordered=ordered)
end

################################################################
#
#  Weighted sampling
#
################################################################

function sample(wv::WeightVec)
    t = rand() * sum(wv)
    w = values(wv)
    n = length(w)
    i = 1
    cw = w[1]
    while cw < t && i < n
        i += 1
        @inbounds cw += w[i]
    end
    return i
end

sample(a::AbstractArray, wv::WeightVec) = a[sample(wv)]

function direct_sample!(a::AbstractArray, wv::WeightVec, x::AbstractArray)
    n = length(a)
    length(wv) == n || throw(DimensionMismatch("Inconsistent lengths."))
    for i = 1:length(x)
        x[i] = a[sample(wv)]
    end
    return x
end

function make_alias_table!(w::AbstractVector{Float64}, wsum::Float64,
                           a::AbstractVector{Float64},
                           alias::AbstractVector{Int})
    # Arguments:
    #
    #   w [in]:         input weights
    #   wsum [in]:      pre-computed sum(w)
    #
    #   a [out]:        acceptance probabilities
    #   alias [out]:    alias table
    #
    # Note: a and w can be the same way, then that away will be
    #       overriden inplace by acceptance probabilities
    #
    # Returns nothing
    #

    n = length(w)
    length(a) == length(alias) == n ||
        throw(DimensionMismatch("Inconsistent array lengths."))

    ac = n / wsum
    for i = 1:n
        @inbounds a[i] = w[i] * ac
    end

    larges = Array(Int, n)
    smalls = Array(Int, n)
    kl = 0  # actual number of larges
    ks = 0  # actual number of smalls

    for i = 1:n
        @inbounds ai = a[i]
        if ai > 1.0
            larges[kl+=1] = i  # push to larges
        elseif ai < 1.0
            smalls[ks+=1] = i  # push to smalls
        end
    end

    while kl > 0 && ks > 0
        s = smalls[ks]; ks -= 1  # pop from smalls
        l = larges[kl]; kl -= 1  # pop from larges
        @inbounds alias[s] = l
        @inbounds al = a[l] = (a[l] - 1.0) + a[s]
        if al > 1.0
            larges[kl+=1] = l  # push to larges
        else
            smalls[ks+=1] = l  # push to smalls
        end
    end

    # this loop should be redundant, except for rounding
    for i = 1:ks
        @inbounds a[smalls[i]] = 1.0
    end
    nothing
end

function alias_sample!(a::AbstractArray, wv::WeightVec, x::AbstractArray)
    n = length(a)
    length(wv) == n || throw(DimensionMismatch("Inconsistent lengths."))

    # create alias table
    ap = Array(Float64, n)
    alias = Array(Int, n)
    make_alias_table!(values(wv), sum(wv), ap, alias)

    # sampling
    s = RandIntSampler(n)
    for i = 1:length(x)
        j = rand(s)
        x[i] = rand() < ap[j] ? a[j] : a[alias[j]]
    end
    return x
end

function xmultinom_sample!(a::AbstractArray, wv::WeightVec, x::AbstractArray)
    n = length(a)
    length(wv) == n || throw(DimensionMismatch("Inconsistent lengths."))
    k = length(x)
    offset = 0
    i = 1
    wsum = sum(wv)
    w = values(wv)

    while offset < k
        rk = k - offset
        wi = w[i]

        if i == n || wi >= wsum
            for j = 1 : rk
                @inbounds x[offset + j] = a[i]
            end
            offset = k
        else
            m = rand_binom(rk, wi / wsum)
            for j = 1 : m
                @inbounds x[offset + j] = a[i]
            end
            i += 1
            wsum -= wi
            offset += m
        end
    end
    return x
end

function naive_wsample_norep!(a::AbstractArray, wv::WeightVec, x::AbstractArray)
    n = length(a)
    length(wv) == n || throw(DimensionMismatch("Inconsistent lengths."))
    k = length(x)

    w = Array(Float64, n)
    copy!(w, values(wv))
    wsum = sum(wv)

    for i = 1:k
        u = rand() * wsum
        j = 1
        c = w[1]
        while c < u && j < n
            @inbounds c += w[j+=1]
        end
        @inbounds x[i] = a[j]

        @inbounds wsum -= w[j]
        @inbounds w[j] = 0.0
    end
    return x
end

function sample!(a::AbstractArray, wv::WeightVec, x::AbstractArray;
                 replace::Bool=true, ordered::Bool=false)
    n = length(a)
    k = length(x)

    if replace
        if ordered
            if k < n && k < 256
                sort!(direct_sample!(a, wv, x))
            else
                xmultinom_sample!(a, wv, x)
            end
        else
            if k > 3 * n
                shuffle!(xmultinom_sample!(a, wv, x))
            else
                if n < 40
                    direct_sample!(a, wv, x)
                else
                    t = ifelse(n < 500, 64, 32)
                    if k < t
                        direct_sample!(a, wv, x)
                    else
                        alias_sample!(a, wv, x)
                    end
                end
            end
        end
    else
        naive_wsample_norep!(a, wv, x)
        if ordered
            sort!(x)
        end
    end
    return x
end

sample{T}(a::AbstractArray{T}, wv::WeightVec, n::Integer; replace::Bool=true, ordered::Bool=false) =
    sample!(a, wv, Array(T, n); replace=replace, ordered=ordered)

sample{T}(a::AbstractArray{T}, wv::WeightVec, dims::Dims; replace::Bool=true, ordered::Bool=false) =
    sample!(a, wv, Array(T, dims); replace=replace, ordered=ordered)

# wsample interface

wsample(w::RealVector) = sample(weights(w))
wsample(a::AbstractArray, w::RealVector) = sample(a, weights(w))

wsample!(a::AbstractArray, w::RealVector, x::AbstractArray; replace::Bool=true, ordered::Bool=false) =
    sample!(a, weights(w), x; replace=replace, ordered=ordered)

wsample{T}(a::AbstractArray{T}, w::RealVector, n::Integer; replace::Bool=true, ordered::Bool=false) =
    wsample!(a, w, Array(T, n); replace=replace, ordered=ordered)

wsample{T}(a::AbstractArray{T}, w::RealVector, dims::Dims; replace::Bool=true, ordered::Bool=false) =
    wsample!(a, w, Array(T, dims); replace=replace, ordered=ordered)


# In this tutorial, we will learn
#  - How to implement multi-field PDEs
#  - How to build div-conforming FE spaces
#  - How to impose boundary conditions in multi-field problems
#
# ## Problem statement
#
# In this tutorial, we show how to solve a multi-field PDE in Gridap. As a model problem, we consider the Darcy equations with Dirichlet and Neumann boundary conditions. The PDE we want to solve is: find the flux vector $u$, and the pressure $p$ such that
#
# ```math
#    \left\lbrace
#    \begin{aligned}
#       \kappa^{-1} u + \nabla p = 0  \ &\text{in} \ \Omega,\\
#       \nabla \cdot u = f  \ &\text{in} \ \Omega,\\
#       u \cdot n = g \ &\text{on}\  \Gamma_{\rm D},\\
#       p = h \ &\text{on}\ \Gamma_{\rm N},\\
#    \end{aligned}
#    \right.
# ```
#
# being $n$ the outwards unit normal vector to the boundary $\partial\Omega$.  In this particular tutorial, we consider the unit square $\Omega \doteq (0,1)^2$ as the computational domain, the Neumann boundary $\Gamma_{\rm N}$ is the right and left sides of $\Omega$, and $\Gamma_{\rm D}$ is the bottom and top sides of $\Omega$. We consider $f = g \doteq 0$ and $h(x) \doteq x_1$, i.e., $h$ equal to 0 on the left side and 1 on the right side. The inverse of the permeability tensor, namely $\kappa^{-1}(x)$, is chosen equal to
#
# ```math
# \begin{pmatrix}
#   100 & 90 \\
#   90 & 100
# \end{pmatrix}
# \text{ for } \ x \in [0.4,0.6]^2, \text{ and }
# \begin{pmatrix}
#   1 & 0 \\
#   0 & 1
# \end{pmatrix}
# \ \text	{otherwise.}
# ```
#
# In order to state this problem in weak form, we introduce the following Sobolev spaces. $H(\mathrm{div};\Omega)$ is the space of vector fields in $\Omega$, whose components and divergence are in $L^2(\Omega)$. On the other hand, $H_g(\mathrm{div};\Omega)$ and $H_0(\mathrm{div};\Omega)$ are the subspaces of functions in $H(\mathrm{div};\Omega)$ such that their normal traces are equal to $g$ and $0$ respectively almost everywhere in $\Gamma_{\rm D}$. With these notations, the weak form reads: find $(u,p)\in H_g(\mathrm{div};\Omega)\times L^2(\Omega)$ such that $a((u,q),(v,q)) = b(v,q)$ for all $(v,q)\in H_0(\mathrm{div};\Omega)\times L^2(\Omega)$, where
#
# ```math
# \begin{aligned}
# a((u,q),(v,q)) &\doteq \int_{\Omega}  v \cdot \left(\kappa^{-1} u\right) \ {\rm d}\Omega - \int_{\Omega} (\nabla \cdot v)\ p \ {\rm d}\Omega + \int_{\Omega} q\ (\nabla \cdot u) \ {\rm d}\Omega,\\
# b(v,q) &\doteq \int_{\Omega} q\ f \ {\rm  d}\Omega - \int_{\Gamma_{\rm N}} (v\cdot n)\ h  \ {\rm  d}\Gamma.
# \end{aligned}
# ```
#
#
#  ## Numerical scheme
#
# In this tutorial, we use the Raviart-Thomas (RT)  space for the flux approximation [1]. On a reference square with sides aligned with the Cartesian axes, the \ac{rt} space of order $k$ is represented as $Q_{(k+1,k)} \times Q_{(k,k+1)}$, being the polynomial space defined as follows. The component  $w_\alpha$ of a vector field $w$ in $Q_{(k+1,k)} \times Q_{(k,k+1)}$ is obtained as the tensor product of univariate polynomials of order $k+1$ in direction $\alpha$ times univariate polynomials of order $k$ on the other directions. That is, $\nabla\cdot w \in Q_k$, where $Q_k$ is the multivariate polynomial space of degree at most $k$ in each of the spatial coordinates. Note that the definition of the \ac{rt} space also applies to arbitrary dimensions. The global FE space for the flux $V$ is obtained by mapping the cell-wise RT space into the physical space using the Piola transformation and enforcing continuity of normal traces across cells (see [1] for specific details).
#
#  We consider the subspace  $V_0$ of functions in $V$ with zero normal trace on $\Gamma_{\rm D}$, and the subspace $V_g$ of functions in $V$ with normal trace equal to the projection of $g$ onto the space of traces of $V$ on $\Gamma_{\rm D}$. With regard to the pressure, we consider the discontinuous space of cell-wise polynomials in $Q_k$.
#
# ## Discrete model
#
# We start the driver loading the Gridap package and constructing the geometrical model. We generate a $100\times100$ structured mesh for the domain $(0,1)^2$.

using Gridap
domain = (0,1,0,1)
partition = (100,100)
model = CartesianDiscreteModel(domain,partition)

# ## Multi-field FE spaces
#
# Next, we build the FE spaces. We consider the first order RT space for the flux and the discontinuous pressure space as described above.  This mixed FE pair satisfies the inf-sup condition and, thus, it is stable.

order = 1

V = TestFESpace(
  reffe=:RaviartThomas, order=order, valuetype=VectorValue{2,Float64},
  conformity=:HDiv, model=model, dirichlet_tags=[5,6])

Q = TestFESpace(
  reffe=:QLagrangian, order=order, valuetype=Float64,
  conformity=:L2, model=model)

# Note that the Dirichlet boundary for the flux are the bottom and top sides of the squared domain (identified with the boundary tags 5, and 6 respectively), whereas no Dirichlet data can be imposed on the pressure space. We select `conformity=:HDiv` for the flux (i.e., shape functions with $H^1(\mathrm{div};\Omega)$ regularity) and `conformity=:L2` for the pressure (i.e. discontinuous shape functions).
#
# From these objects, we construct the trial spaces. Note that we impose homogeneous boundary conditions for the flux.

uD = VectorValue(0.0,0.0)
U = TrialFESpace(V,uD)
P = TrialFESpace(Q)

# When the singe-field spaces have been designed, the multi-field test and trial spaces are expressed as arrays of single-field ones in a natural way.

Y = MultiFieldFESpace([V, Q])
X = MultiFieldFESpace([U, P])

# ## Numerical integration
#
# In this example we need to integrate in the interior of $\Omega$ and on the Neumann boundary $\Gamma_{\rm N}$. For the volume integrals, we extract the triangulation from the geometrical model and define the corresponding cell-wise quadrature of degree of exactness at least 2 as follows.

trian = Triangulation(model)
degree = 2
quad = CellQuadrature(trian,degree)

# In order to integrate the Neumann boundary condition, we only need to build an integration mesh for the right side of the domain (which is the only part of $\Gamma_{\rm N}$, where the Neumann function $h$ is different from zero). Within the model, the right side of $\Omega$ is identified with the boundary tag 8. Using this identifier, we extract the corresponding surface triangulation and create a quadrature with the desired degree of exactness.

neumanntags = [8,]
btrian = BoundaryTriangulation(model,neumanntags)
bquad = CellQuadrature(btrian,degree)

# ## Weak form
#
# We start by defining the permeability tensors inverses commented above using the `@law` macro.

const kinv1 = TensorValue(1.0,0.0,0.0,1.0)
const kinv2 = TensorValue(100.0,90.0,90.0,100.0)
@law function σ(x,u)
   if ((abs(x[1]-0.5) <= 0.1) && (abs(x[2]-0.5) <= 0.1))
      return kinv2⋅u
   else
      return kinv1⋅u
   end
end

# With this definition, we can express the integrand of the bilinear form as follows.

px = get_physical_coordinate(trian)

function a(x,y)
   v, q = y
   u, p = x
   v⋅σ(px,u) - (∇⋅v)*p + q*(∇⋅u)
end

# The arguments `x` and `y` of previous function represent a trial and a test function in the multi-field test and trial spaces `X` and `Y` respectively. In the first lines in the function definition, we unpack the single-field test and trial functions from the multi-field ones. E.g., `v` represents a test function for the flux and `q` for the pressure. These quantities can also be written as `y[1]` and `y[2]` respectively. From the single-field functions, we write the different terms of the bilinear form as we have done in previous tutorials.
#
# In a similar way, we can define the forcing term related to the Neumann boundary condition.

nb = get_normal_vector(btrian)
h = -1.0
function b_ΓN(y)
  v, q = y
  (v⋅nb)*h
end

# ## Multi-field FE problem
#
# Finally, we can assemble the FE problem and solve it. Note that we build the `AffineFEOperator` object using the multi-field trial and test spaces `Y` and `X`.

t_Ω = LinearFETerm(a,trian,quad)
t_ΓN = FESource(b_ΓN,btrian,bquad)
op = AffineFEOperator(X,Y,t_Ω,t_ΓN)
xh = solve(op)
uh, ph = xh

# Since this is a multi-field example, the `solve` function returns a multi-field solution `xh`, which can be unpacked in order to finally recover each field of the problem. The resulting single-field objects can be visualized as in previous tutorials (see next figure).

writevtk(trian,"darcyresults",cellfields=["uh"=>uh,"ph"=>ph])

# ![](../assets/darcy/darcy_results.png)
#
# ## References
#
# [1] F. Brezzi and M. Fortin. *Mixed and hybrid finite element methods*. Springer-Verlag, 1991.
#


@testset "complex_powers" begin
    complex_powers_comparison(z, m, T=Float64) = (
        complex_powers(Complex{T}(z), m),
        z.^collect(0:m),
        eps(T)*m
    )

    for T in [Float64, Float32, Float16]
        fudge = one(T) + 2 * sqrt(eps(T))
        for k in 0:25
            z = exp(k*big(π)im/10)
            for m in [0, 1, 2, 3, 4, 1_000]
                mine, theirs, ϵ = complex_powers_comparison(z, m, T)
                @test mine ≈ theirs rtol=2ϵ
                @test_throws DomainError complex_powers(Complex{T}(z*fudge), m)
                @test_throws DomainError complex_powers(Complex{T}(z/fudge), m)
            end
        end
    end
end


using DelimitedFiles
using RollingFunctions

input = joinpath(@__DIR__, "input")
entries = readdlm(input, Int)

function depth(entries)
    return length(filter(n -> n > 0, diff(entries)))
end

p1 = depth(vec(entries))
p2 = depth(rolling(sum, vec(entries), 3))

println("-----------------------------------------------------------------------")
println("sonar sweep -- part one :: $p1")
println("sonar sweep -- part two :: $p2")
println("-----------------------------------------------------------------------")

@assert(p1 == 1154)
@assert(p2 == 1127)


using ZigZagBoomerang
using LinearAlgebra
using Random
using SparseArrays
using Distributions
using Statistics
using ZigZagBoomerang: idot, idot_moving!
using ReverseDiff
using Test
using FileIO
using GLM

println("Sparse logistic regression")

# Design matrix
Random.seed!(2)

include("exampledesign.jl")
# create mock design matrix with 2 categorical explanatory variables
# and their interaction effects and 2 continuous explanatory variable
A = example_design_matrix(;num_rows = 50_000)

n, p = size(A)
@show n, p

const m = 1 # m Bernoulli samples for each set of covariates
println("Av. number of regressors per column ", mean(sum(A .!= 0, dims=1)), ", row ", mean(sum(A .!= 0, dims=2)))

# Data from the model
xtrue = randn(p) .* (rand(p) .< 0.2)
sigmoid(x) = inv(one(x) + exp(-x))
lsigmoid(x) = -log(one(x) + exp(-x))
y_ = Matrix(rand(n, m) .< sigmoid.(A*xtrue))
y = sum(y_, dims=2)[:]

# gaussian mixture prior and target potential
pri(x, w1 = 0.5, a=(0.1)^2, w2 = 0.5, b=4.0^2) = -log(w1*exp(-x^2/(2*a))/sqrt(2pi*a) + w2*exp(-x^2/(2*b))/sqrt(2pi*b))
∇pri(x, w1 = 0.5, a=(0.1)^2, w2 = 0.5, b=4.0^2) = -(-(w1*x*exp(-x^2/(2*a)))/a^(3/2) - (w2*x*exp(-x^2/(2*b)))/b^(3/2))/(w1*exp(-x^2/(2*a))/sqrt(a) + w2*exp(-x^2/(2*b))/sqrt(b))
ϕ(x, A, y::Vector) = pri.(x) - sum(y .* lsigmoid.(A*x) + (m .- y) .* lsigmoid.(-A*x))

sigmoidn(x) = sigmoid(-x)
nsigmoid(x) = -sigmoid(x)

# Gradient (found using http://www.matrixcalculus.org/ S. Laue, M. Mitterreiter, and J. Giesen. A Simple and Efficient Tensor Calculus, AAAI 2020.)
∇ϕ(x, A, y) = ∇pri.(x) - A'*(y .* sigmoidn.(A*x)) - A'*((m .- y).*nsigmoid.(A*x))
function ∇ϕ!(out, x, A, y)
    out .= ∇pri.(x) - A'*(y .* sigmoidn.(A*x)) - A'*((m .- y).*nsigmoid.(A*x))
    out
end

function ∇ϕ(x, i, A, y) # partial derivative
    s = ∇pri(x[i])
    n, p = size(A)
    for j in 1:n
        z = sum(A[j,r]*x[r] for r in 1:p)
        s += -A[j,i]*(y[j]*sigmoidn(z) + (m - y[j])*nsigmoid(z))
    end
    return s
end

function ∇ϕr(x, i, A, y, control, μ, bias, k) # with subsampling and optional control variate
    s = ∇pri(x[i]) + control*(bias[i] - pri(μ[i]))
    n, p = size(A)
    sampler = Random.SamplerRangeNDL(1:n)
    for _ in 1:k
        j = rand(sampler)
        z = sum(A[j,r]*x[r] for r in 1:p)
        s += -n/k*A[j,i]*(y[j]*sigmoidn(z) + (m - y[j])*nsigmoid(z))
        if control
            z2 = sum(A[j,r]*μ[r] for r in 1:p)
            s -= -n/k*A[j,i]*(y[j]*sigmoidn(z2) + (m - y[j])*nsigmoid(z2))
        end
    end
    return s
end

# Some Newton steps towards the mode as starting point
# The mode can be used as control variate
println("Newton steps")
x0 = zeros(p)
@time for i in 1:8
    println(i)
    global x0
    H = ReverseDiff.jacobian(x -> ∇ϕ(x, A, y), x0)
    x0 = x0 - cholesky(Symmetric(I + H))\∇ϕ(x0, A, y)
end
@show norm(xtrue - x0)

μ = copy(x0)
Γfull = ReverseDiff.jacobian(x -> ∇ϕ(x, A, y), μ)
Γ = sparse(1.0I, p, p)
bias = ∇ϕ(μ, A, y)
@show norm(bias)

xtest = randn(p)
i = 8
@test ∇ϕ(xtest, A, y)[i] ≈ ∇ϕ(xtest, i, A, y)
#@test ∇ϕ(xtest, i, A, y) ≈ mean(∇ϕr(xtest, i, A, y, false, μ, bias, 100) for k in 1:20000)
#@test ∇ϕ(xtest, i, A, y) ≈ mean(∇ϕr(xtest, i, A, y, true, μ, bias, 100) for k in 1:20000)

# Initial values
t0 = 0.0

# Define ZigZag

#σ = sqrt.(diag(inv(Matrix(Γfull))))
#extrema(σ)
σ = 0.1*(Vector(diag(Γ))).^(-0.5) # cheaper

Z = ZigZag(Γ, μ, σ)

θ0 = rand([-1.0,1.0], p) .* σ

# Run sparse ZigZag for T time units and collect trajectory
println("Distance starting point")
@show norm(x0 - xtrue)

println("Run pdmp")
T = 100.0
c = ones(Float64, p) # rejection bounds
c0 = copy(c)
traj, u, (acc,num), c = @time pdmp(∇ϕr, t0, x0, θ0, T, c, Z, A, y, true, μ, bias, 40; adapt=true)
@show acc, num, acc/num
@show extrema(c ./ c0)

dt = T/4000
x̂ = mean(x for (t,x) in discretize(traj, dt))
X = Float64[]
ts = Float64[]
for (t,x) in discretize(traj, dt)
    append!(ts, t)
    append!(X, x)
end
X = reshape(X, p, length(X)÷p)

[xtrue x̂ abs.(xtrue - x̂)]

@show norm(xtrue - μ)
@show norm(xtrue - x̂)

println("Plot")
using Makie
using Colors
using GoldenSequences

ps = [3, 22, 50, p-3, p-2, p-1]
p0 = length(ps)
cs = map(x->RGB(x...), (Iterators.take(GoldenSequence(3), p0)))
pis = []
for i in 1:p0
    p_i = lines(ts, X[ps[i], :], color=cs[i])
    lines!(p_i, [0, T], [xtrue[ps[i]], xtrue[ps[i]]], color=cs[i], linewidth=2.0)
    ylabel!(p_i, "var$(ps[i])")
    xlabel!(p_i, "t")
    push!(pis, p_i)
end
p1 = title(hbox(pis...), "Spike and slab Bayesian logistic regression n=$(m*n), p=$p", textsize=20)
save(joinpath("figures","spikeandslab$(typeof(Z).name).png"), p1)

p2 = title(vbox(hbox([text("$p", show_axis=false, textsize=10) for p in ps]...),
    [hbox([i==j ? lines(ts, X[i,:]) : lines(X[i, :], X[j, :]) for i in ps]...) for j in ps]...), "$ps")

r = -3:0.01:3
p3 = lines(r, ∇pri.(r))
lines!(p3, r, exp.(-pri.(r)))

p1


# If you follow the equations in the PDF in the /doc directory, you will see all functions here straightforward.

# TODO: I cound that "utils.jl" has been `include`ed in the main `VectorSphericalWaves` module. Is there a more neat way? when I `include("utils.jl")` here, I get lots of warnings.

using StaticArrays

# exporting main functions, separating real and imaginary parts
export M_mn_wave_SeparateRealImag
export N_mn_wave_SeparateRealImag

export M_mn_wave_SeparateRealImag_SMatrix
export N_mn_wave_SeparateRealImag_SMatrix

#############################################################################################
# The following functions are the same as before, but are desinged to be automatically differentiable (e.g., with Zygote).
# The main difference is that we don't use any complex numbers here. Complex numbers and vectors are represented as 1x2 and nx2 Matrices, respectively. The first and second columns represent the real and imaginary parts.
#############################################################################################
#############################################################################################
# calculate B(θ), C(θ), P(θ)
"""
    equation C.19, returns an array
It returns a 3x2 Matrix, the first and second columns are the real and imaginary parts, respectively.
"""
function B_mn_of_θ_SeparateRealImag(m::I, n::I, θ::R) where {R <: Real, I <: Integer}
    # equation C.19
    return hcat(
        vcat(zero(R), τₘₙ(m, n, θ), zero(R)),
        vcat(zero(R), zero(R), πₘₙ(m, n, θ))
    )
end

"""
    Same as B_mn_of_θ_SeparateRealImag, but return SMatrix
It returns a 3x2 SMatrix, the first and second columns are the real and imaginary parts, respectively.
"""
function B_mn_of_θ_SeparateRealImag_SMatrix(m::I, n::I, θ::R) where {R <: Real, I <: Integer}
    # equation C.19
    return hcat(
        SVector(zero(R), τₘₙ(m, n, θ), zero(R)),
        SVector(zero(R), zero(R), πₘₙ(m, n, θ))
    )
end

"""
    equation C.20, returns an array
It returns a 3x2 Matrix, the first and second columns are the real and imaginary parts, respectively.
"""
function C_mn_of_θ_SeparateRealImag(m::I, n::I, θ::R) where {R <: Real, I <: Integer}
    # equation C.20
    return hcat(
        vcat(zero(R), zero(R), -1 * τₘₙ(m, n, θ)),
        vcat(zero(R), πₘₙ(m, n, θ), zero(R))
    )
end

"""
    Same as C_mn_of_θ_SeparateRealImag, but return SMatrix
It returns a 3x2 SMatrix, the first and second columns are the real and imaginary parts, respectively.
"""
function C_mn_of_θ_SeparateRealImag_SMatrix(m::I, n::I, θ::R) where {R <: Real, I <: Integer}
    # equation C.20
    return hcat(
        SVector(zero(R), zero(R), -1 * τₘₙ(m, n, θ)),
        SVector(zero(R), πₘₙ(m, n, θ), zero(R))
)
end

"""
    equation C.21, returns an array
It returns a 3x2 Matrix, the first and second columns are the real and imaginary parts, respectively.
"""
function P_mn_of_θ_SeparateRealImag(m::I, n::I, θ::R) where {R <: Real, I <: Integer}
    # equation C.21
    return hcat(
        P_mn_of_θ(m, n, θ),
        vcat(zero(R), zero(R), zero(R))
)
end

"""
    Same as P_mn_of_θ_SeparateRealImag, returns SMatrix
It returns a 3x2 SMatrix, the first and second columns are the real and imaginary parts, respectively.
"""
function P_mn_of_θ_SeparateRealImag_SMatrix(m::I, n::I, θ::R) where {R <: Real, I <: Integer}
    # equation C.21
    return hcat(
        P_mn_of_θ_SVector(m, n, θ),
        SVector(zero(θ), zero(θ), zero(θ))
    )
end

function convert_from_fun_of_θ_to_fun_of_θ_ϕ(fun_tobe_converted::Function, m::I, n::I, θ::R, ϕ::R) where {R <: Real, I <: Integer}
    # equation C.16, C.17, and C.18
    coeff = Zygote.dropgrad((-one(R))^m * sqrt_factorial_n_plus_m_over_factorial_n_minus_m(m, n))
    B_of_θ_coef = convert(R, coeff) .* fun_tobe_converted(m, n, θ)
    B_of_θ_coef_real = B_of_θ_coef[:,1]
    B_of_θ_coef_imag = B_of_θ_coef[:,2]
    exp_imϕ_real = cos(Zygote.dropgrad(m) * ϕ)
    exp_imϕ_imag = sin(Zygote.dropgrad(m) * ϕ)

    return hcat(
        B_of_θ_coef_real .* exp_imϕ_real - B_of_θ_coef_imag .* exp_imϕ_imag,
        B_of_θ_coef_real .* exp_imϕ_imag + B_of_θ_coef_imag .* exp_imϕ_real
    )
end

#############################################################################################
# calculate B(θ,ϕ), C(θ,ϕ), P(θ,ϕ)
function B_mn_of_θ_ϕ_SeparateRealImag(m::I, n::I, θ::R, ϕ::R) where {R <: Real, I <: Integer}
    # equation C.16
    return convert_from_fun_of_θ_to_fun_of_θ_ϕ(B_mn_of_θ_SeparateRealImag, m, n, θ, ϕ)
end

function C_mn_of_θ_ϕ_SeparateRealImag(m::I, n::I, θ::R, ϕ::R) where {R <: Real, I <: Integer}
    # equation C.17
    return convert_from_fun_of_θ_to_fun_of_θ_ϕ(C_mn_of_θ_SeparateRealImag, m, n, θ, ϕ)
end

function P_mn_of_θ_ϕ_SeparateRealImag(m::I, n::I, θ::R, ϕ::R) where {R <: Real, I <: Integer}
    # equation C.18
    return convert_from_fun_of_θ_to_fun_of_θ_ϕ(P_mn_of_θ_SeparateRealImag, m, n, θ, ϕ)
end

# same as above, but returns *_SMatrix
function B_mn_of_θ_ϕ_SeparateRealImag_SMatrix(m::I, n::I, θ::R, ϕ::R) where {R <: Real, I <: Integer}
    # equation C.16
    return convert_from_fun_of_θ_to_fun_of_θ_ϕ(B_mn_of_θ_SeparateRealImag_SMatrix, m, n, θ, ϕ)
end

function C_mn_of_θ_ϕ_SeparateRealImag_SMatrix(m::I, n::I, θ::R, ϕ::R) where {R <: Real, I <: Integer}
    # equation C.17
    return convert_from_fun_of_θ_to_fun_of_θ_ϕ(C_mn_of_θ_SeparateRealImag_SMatrix, m, n, θ, ϕ)
end

function P_mn_of_θ_ϕ_SeparateRealImag_SMatrix(m::I, n::I, θ::R, ϕ::R) where {R <: Real, I <: Integer}
    # equation C.18
    return convert_from_fun_of_θ_to_fun_of_θ_ϕ(P_mn_of_θ_SeparateRealImag_SMatrix, m, n, θ, ϕ)
end

#############################################################################################
# calculate spherical Bessel and Hankel functions and their derivarive
function SeparateRealImag_for_Bessel_Hankel_and_derivatives(fun_tobe_converted::Function, n::I, x_r::R, x_i::R) where {R <: Real, I <: Integer}
    bessel_complex = fun_tobe_converted(n, complex(x_r, x_i))
    return hcat(real(bessel_complex), imag(bessel_complex))
end

function spherical_Bessel_j_n_SeparateRealImag(n::I, x_r::R, x_i::R) where {R <: Real, I <: Integer}
    return SeparateRealImag_for_Bessel_Hankel_and_derivatives(spherical_Bessel_j_n, n, x_r, x_i)
end

function spherical_Bessel_y_n_SeparateRealImag(n::I, x_r::R, x_i::R) where {R <: Real, I <: Integer}
    return SeparateRealImag_for_Bessel_Hankel_and_derivatives(spherical_Bessel_y_n, n, x_r, x_i)
end

function spherical_Hankel_h1_n_SeparateRealImag(n::I, x_r::R, x_i::R) where {R <: Real, I <: Integer}
    return SeparateRealImag_for_Bessel_Hankel_and_derivatives(spherical_Hankel_h1_n, n, x_r, x_i)
end

function one_over_x_by_∂_x_j_n_by_∂x_SeparateRealImag(n::I, x_r::R, x_i::R) where {R <: Real, I <: Integer}
    # n_over_x = complex_divide(n * ones(size(x_r)), zeros(size(x_r)), x_r, x_i) # TODO: find a way to make sure all inputs have the same size
    n_over_x = complex_divide(Zygote.dropgrad(n), zero(x_r), x_r, x_i)
    return (spherical_Bessel_j_n_SeparateRealImag(n - 1, x_r, x_i) - complex_multiply(n_over_x, spherical_Bessel_j_n_SeparateRealImag(n, x_r, x_i)))
end

function one_over_x_by_∂_x_y_n_by_∂x_SeparateRealImag(n::I, x_r::R, x_i::R) where {R <: Real, I <: Integer}
    # n_over_x = complex_divide(n * ones(size(x_r)), zeros(size(x_r)), x_r, x_i) # TODO: find a way to make sure all inputs have the same size
    n_over_x = complex_divide(Zygote.dropgrad(n), zero(x_r), x_r, x_i)
    return (spherical_Bessel_y_n_SeparateRealImag(n - 1, x_r, x_i) - complex_multiply(n_over_x, spherical_Bessel_y_n_SeparateRealImag(n, x_r, x_i)))
end

function one_over_x_by_∂_x_h_n_by_∂x_SeparateRealImag(n::I, x_r::R, x_i::R) where {R <: Real, I <: Integer}
    der_j = one_over_x_by_∂_x_j_n_by_∂x_SeparateRealImag(n, x_r, x_i)
    der_y = one_over_x_by_∂_x_y_n_by_∂x_SeparateRealImag(n, x_r, x_i)
    return der_j + complex_multiply(zero(R), one(R), der_y[1], der_y[2])
    # return (one_over_x_by_∂_x_j_n_by_∂x_SeparateRealImag(n, x_r, x_i) + complex_multiply(0, 1, one_over_x_by_∂_x_y_n_by_∂x_SeparateRealImag(n, x_r, x_i)...))
end

function get_radial_function_and_special_derivative_given_kind_SeparateRealImag(kind)
    """
    return the radial function that is appropriate with the type of VSWF
    """

    if kind in ["regular", "incoming", 1]
        radial_function = spherical_Bessel_j_n_SeparateRealImag
        radial_function_special_derivative = one_over_x_by_∂_x_j_n_by_∂x_SeparateRealImag
    elseif kind in ["irregular", "outgoing", 2]
        radial_function = spherical_Hankel_h1_n_SeparateRealImag
        radial_function_special_derivative = one_over_x_by_∂_x_h_n_by_∂x_SeparateRealImag
    else
        throw(DomainError(""" 'kind' has to be one of the following: ["regular", "incoming", "irregular", "outgoing"] """))
    end

    return radial_function, radial_function_special_derivative
end

"""Calculate vector spherical wave M
This version takes kr_r and kr_i as inputs (i.e., wavevector multiplied by r)
"""
function M_mn_wave_SeparateRealImag(m::I, n::I, kr_r::R, kr_i::R, θ::R, ϕ::R, kind="regular") where {R <: Real, I <: Integer}
    radial_function, _ = get_radial_function_and_special_derivative_given_kind_SeparateRealImag(kind)
    gamma_by_radial = convert(R, Zygote.dropgrad(γ_mn(m, n))) .* radial_function(Zygote.dropgrad(n), kr_r, kr_i)
    return complex_multiply(vcat((gamma_by_radial), (gamma_by_radial), (gamma_by_radial)), C_mn_of_θ_ϕ_SeparateRealImag(Zygote.dropgrad(m), Zygote.dropgrad(n), θ, ϕ))

    # TODO: write this in a more elegant way: vcat(gamma_by_radial, gamma_by_radial, gamma_by_radial)
end

"""Calculate vector spherical wave M
This version takes k_r, k_i and r as inputs.
"""
function M_mn_wave_SeparateRealImag(m::I, n::I, k_r::R, k_i::R, r::R, θ::R, ϕ::R, kind="regular") where {R <: Real, I <: Integer}
    return M_mn_wave_SeparateRealImag(m, n, k_r*r, k_i*r, θ, ϕ, kind)
end

"""
    Same as `M_mn_wave_SeparateRealImag`, but returns SMatrix
"""
function M_mn_wave_SeparateRealImag_SMatrix(m::I, n::I, kr_r::R, kr_i::R, θ::R, ϕ::R, kind="regular") where {R <: Real, I <: Integer}
    radial_function, _ = get_radial_function_and_special_derivative_given_kind_SeparateRealImag(kind)
    gamma_by_radial = convert(R, Zygote.dropgrad(γ_mn(m, n))) .* radial_function(Zygote.dropgrad(n), kr_r, kr_i)
    return complex_multiply_SMatrix(gamma_by_radial .* SMatrix{3,2}(ones(3, 2)), C_mn_of_θ_ϕ_SeparateRealImag_SMatrix(Zygote.dropgrad(m), Zygote.dropgrad(n), θ, ϕ))
    # TODO: write this in a more elegant way: vcat(gamma_by_radial, gamma_by_radial, gamma_by_radial)
end

function M_mn_wave_SeparateRealImag_SMatrix(m::I, n::I, k_r::R, k_i::R, r::R, θ::R, ϕ::R, kind="regular") where {R <: Real, I <: Integer}
    return M_mn_wave_SeparateRealImag_SMatrix(m, n, k_r*r, k_i*r, θ, ϕ, kind)
end

function N_mn_wave_SeparateRealImag(m::I, n::I, kr_r::R, kr_i::R, θ::R, ϕ::R, kind="regular") where {R <: Real, I <: Integer}
    radial_function, radial_function_special_derivative  = get_radial_function_and_special_derivative_given_kind_SeparateRealImag(kind)
    ceoff = complex_multiply(complex_divide(Zygote.dropgrad(n * (n + 1)), zero(kr_r), kr_r, kr_i), radial_function(Zygote.dropgrad(n), kr_r, kr_i))
    rad_fun_der = radial_function_special_derivative(Zygote.dropgrad(n), kr_r, kr_i)
    return convert(R, Zygote.dropgrad(γ_mn(m, n))) .* (
        complex_multiply([ceoff; ceoff; ceoff], P_mn_of_θ_ϕ_SeparateRealImag(Zygote.dropgrad(m), n, θ, ϕ))
        + complex_multiply([rad_fun_der; rad_fun_der; rad_fun_der], B_mn_of_θ_ϕ_SeparateRealImag(Zygote.dropgrad(m), n, θ, ϕ))
    )
    # TODO: write this in a more elegant way: [ceoff; ceoff; ceoff]
end

function N_mn_wave_SeparateRealImag(m::I, n::I, k_r::R, k_i::R, r::R, θ::R, ϕ::R, kind="regular") where {R <: Real, I <: Integer}
    return N_mn_wave_SeparateRealImag(m, n, k_r*r, k_i*r, θ, ϕ, kind)
end

"""
    Same as `N_mn_wave_SeparateRealImag`, but returns SMatrix
"""
function N_mn_wave_SeparateRealImag_SMatrix(m::I, n::I, kr_r::R, kr_i::R, θ::R, ϕ::R, kind="regular") where {R <: Real, I <: Integer}
    radial_function, radial_function_special_derivative  = get_radial_function_and_special_derivative_given_kind_SeparateRealImag(kind)
    ceoff = complex_multiply(complex_divide(n * (n + 1), zero(kr_r), kr_r, kr_i), radial_function(n, kr_r, kr_i))
    rad_fun_der = radial_function_special_derivative(n, kr_r, kr_i)
    return convert(R, γ_mn(m, n)) .* (
        complex_multiply_SMatrix(ceoff .* SMatrix{3,2}(ones(3, 2)), P_mn_of_θ_ϕ_SeparateRealImag_SMatrix(m, n, θ, ϕ))
        + complex_multiply_SMatrix(rad_fun_der .* SMatrix{3,2}(ones(3, 2)), B_mn_of_θ_ϕ_SeparateRealImag_SMatrix(m, n, θ, ϕ))
    )
    # TODO: write this in a more elegant way: [ceoff; ceoff; ceoff]
end

function N_mn_wave_SeparateRealImag_SMatrix(m::I, n::I, k_r::R, k_i::R, r::R, θ::R, ϕ::R, kind="regular") where {R <: Real, I <: Integer}
    return N_mn_wave_SeparateRealImag_SMatrix(m, n, k_r*r, k_i*r, θ, ϕ, kind)
end


"""There is the possibility of the created vector existing outside of the bounds specified at initialization.
 Because of this, we created a separate function i.e BoundEnsure that checks and corrects for this. In the event
 that one of these rogue points are found, we’ll simply move it to the nearest boundary """

function bound_ensure(m::SVector{N, T}, u::SVector{N, T}, X::IntervalBox{N, T}) where {N, T}

    for i in 1:length(X)
      if m[i] < X[i].lo
        m = setindex(m, u[i]+rand()*(X[i].lo-u[i]), i)
      elseif X[i].hi < m[i]
        m = setindex(m, u[i]+rand()*(X[i].hi-u[i]), i)
      end
    end
    return m
end


import Base.diag

function diag(a::Number)
	return a
end

function diagmmult(A,l)
    (n,m)=size(A)
    Anew=zeros(A)
    for i in 1:m
        Anew[:,i]=A[:,i]*l[i]
    end

    return Anew
end

function bdot(a,b)
    (n,m)=size(a)
    out=zeros(ComplexF64,m)
    for i in 1:m
        out[i]=dot(a[:,i],b[:,i])
    end

    return out
end

function zbicgstabBlock(A,b,x0,maxit,eps)
	#matrix multiply:

	(n,m)=size(x0)

        r=b-A*x0
	s=copy(r)
	y0=copy(r)

	x=copy(x0)

	#inner product:
	#delta=(y0'*r)[1]

	#delta=diag(y0'*r)
	delta=bdot(y0,r)

	#phi=(y0'*(A*s))[1]/delta

	for i in 1:maxit
        #println("    BCG $i")
		#matrix multiply:
		As=A*s
		#inner product
		#phi=(y0'*(As))[1]/delta
		#phi=diag(y0'*As) ./delta
		phi=bdot(y0,As) ./delta

		#omega=1/phi
		omega=1 ./phi
		w=r-As*spdiagm( 0 => omega)
		#w=r-diagmmult(As,omega)

		#matrix multiply:
		Aw=A*w
		#inner product:
		#chi=((Aw)'*w)[1]/norm(Aw)^2
		#chi=diag(Aw'*w) ./diag(Aw'*Aw)
        chi=bdot(Aw,w) ./bdot(Aw,Aw)

		#r=w-chi*Aw
		r=w-Aw*spdiagm( 0 => chi)
        #r=w-diagmmult(Aw,chi)

		rs=sqrt.(bdot(r,r))
		#println("   $i     $(maximum(abs.(rs)))   $(norm(Aw))")
		#if(norm(r)<eps)
		if(maximum(abs.(rs))<eps)
			#println("r finish: $(norm(r))")
			#println("     bicgstab its $i")
			return x
		end

		#x=x+s*omega+w*chi
		x=x+s*spdiagm( 0 =>  omega)+w*spdiagm( 0 =>  chi)
		#x=x+diagmmult(s,omega)+diagmmult(w,chi)

		deltaold=delta
		#inner product:
		#delta=(y0'*r)[1]
		delta=bdot(y0,r)#diag(y0'*r)
		#psi=-1.0*omega*delta/(deltaold*chi)
		psi=-1.0*omega.*delta ./(deltaold.*chi)
		s=r-(s-(As)*spdiagm( 0 =>  chi))*spdiagm( 0 =>  psi)
		#s=r-(s-diagmmult(As,chi.*psi))
		#print("     $(norm(s))\n")
		ss=sqrt.(bdot(s,s))
		if(maximum(abs.(ss))<eps)
		#if(norm(s)<eps)
			#println("s finish: $(norm(s))")
			#println("     bicgstab its $i")
			return x
		end

	end

	#println("          no converge: $(norm(r))")
	return x
end


for N in [Float64, Rational{Int}, Float32]
    # constructor from empty set
    E = EmptySet{N}()
    @test SymmetricIntervalHull(E) == E

    # normal constructor
    h = SymmetricIntervalHull(Ball1(N[2, 3], N(4)))

    # constructor attempt from an unbounded set
    @test_throws AssertionError SymmetricIntervalHull(HalfSpace(N[1], N(0)))

    # dimension
    @test dim(h) == 2

    # support vector
    d = N[1, 0]
    @test σ(d, h)[1] == N(6) && σ(-d, h)[1] == N(-6)
    d = N[0, 1]
    @test σ(d, h)[2] == N(7) && σ(-d, h)[2] == N(-7)

    # boundedness
    @test isbounded(h)

    # isempty
    @test !isempty(h)

    # an_element function
    @test an_element(h) == N[0, 0]

    # radius_hyperrectangle function (uncached and cached)
    h = SymmetricIntervalHull(Ball1(N[2, 3], N(4)))
    for i in 1:2
        @test radius_hyperrectangle(h, 1) == 6
        @test radius_hyperrectangle(h) == N[6, 7]
    end

    # high and low
    h = SymmetricIntervalHull(Singleton(N[1, 2]))
    @test high(h) == N[1, 2]
    @test low(h) == N[-1, -2]
end


include("dwt.jl")

function compressaudio(channels, p, l, h, threshold)
    # Split channels
    left = channels[:,1]
    right = channels[:,2]

    # Pad signals with zeros
    len = length(left)
    J = Integer(ceil(log2(length(left))))
    N = 2^J
    left = [left; zeros(N - len)]
    right = [right; zeros(N - len)]

    # Perfom transforms
    wleft, nstop = mallatdecomp(left,l,h,p)
    wright, _ = mallatdecomp(right,l,h,p)

    # Compress
    tleft = threshold * maximum(abs.(wleft))
    ileft = abs.(wleft) .> tleft
    tright = threshold * maximum(abs.(wright))
    iright = abs.(wright) .> tright

    # Generate return values
    wleft = wleft[ileft]
    wright = wright[iright]
    s = sizeof(wleft) + sizeof(wright) + sizeof(ileft) + sizeof(iright) +
        sizeof(nstop) + sizeof(len)
    return ((wleft, ileft, wright, iright, nstop, len), s)
end

function decompressaudio(compression, p, il, ih)
    # Unpack compression
    cleft, ileft, cright, iright, nstop, len = compression

    # Decompress left channel
    wleft = zeros(length(ileft))
    j = 1
    for i = 1:length(wleft)
        if ileft[i]
            wleft[i] = cleft[j]
            j += 1
        end
    end

    # Decompress right channel
    wright = zeros(length(iright))
    j = 1
    for i = 1:length(wright)
        if iright[i]
            wright[i] = cright[j]
            j += 1
        end
    end

    # Perform inverse transforms
    left = mallatrecon(wleft,nstop,il,ih)
    right = mallatrecon(wright,nstop,il,ih)

    # Unpad
    left = left[1:len]
    right = right[1:len]
    return [left right]
end


using PicoQuant
using LaTeXStrings
import Plots
#=
This script produces graphs showing the output distribution of a rqc converges
to the Porter-Thomas distribution as the depth of the rqc increases.
=#

# Set parameters of the rqc & number of samples to take for each value of depth.
depth = 30; n = 4; m = 4; samples = 100

# Allocate memory for entropy and inverse participation ratios.
average_entropy = zeros(depth); average_sqrd_entropy = zeros(depth)
IPR = zeros(5, depth); IPR_sqrd = zeros(5, depth)

# For each sample at each circuit depth, create a RQC, convert it to a
# TensorNetworkCircuit, contract the network to get the output state and
# calculate the entropy and inverse participation ratios.
for (i, d) in enumerate(1:depth)
    for _ = 1:samples
        rqc = create_RQC(n, m, d)

        # An InteractiveBackend is used by default
        rqc_tn = convert_qiskit_circ_to_network(rqc)
        add_input!(rqc_tn, "0"^length(rqc_tn.input_qubits))

        output_state = full_wavefunction_contraction!(rqc_tn, "vector")
        output_distribution = abs.(output_state).^2

        for k = 1:5
            ipr = 1/sum(output_distribution.^(2*k))
            IPR[k, d] += ipr
            IPR_sqrd[k, d] += ipr^2
        end

        entropy = -sum(output_distribution.*log.(output_distribution))
        average_entropy[i] += entropy
        average_sqrd_entropy[i] += entropy^2
    end
end

average_entropy /= samples
average_sqrd_entropy /= samples
IPR /= samples; IPR_sqrd /= samples

# Get the error in the average entropy.
error = sqrt.(abs.(average_sqrd_entropy - average_entropy.^2))./sqrt(samples)

# Plot the average entropy of the output distribution as a function of depth.
plt = Plots.plot(1:depth, average_entropy, xlabel = "Depth", ylabel = "Entropy",
                 label = "Average entropy of output distribution",
                 yerror = error,
                 title = "Convergence to Porter-Thomas Entropy")

N = 2^(n*m)
porter_thomas_entropy = log(N) - 1 + 0.5772156649

plt = Plots.plot!(plt, 1:depth, porter_thomas_entropy*ones(depth),
                  label = "Porter-Thomas entropy")
display(plt)
Plots.savefig("entropy_convergence.png")

# Divide by Porter-Thomas IPR and get the IPR error
N = BigFloat(N)
for k = 1:5
    IPR[k, :] *= factorial(2*k)/N^(2*k - 1)
    IPR_sqrd[k, :] *= (factorial(2*k)/(N^(2*k - 1)))^2
end
IPR_error = sqrt.(abs.(IPR_sqrd - IPR.^2))./sqrt(samples)

# Plot the average inverse participation ratio of the Porter-Thomas distribution
# divided by that of the RQC output distribution as a function of depth.
plt2 = Plots.plot(xlabel = "Depth", ylabel = L"N^{-k+1}k!/IPR^k",
                  title = "Convergence to Porter-Thomas IPR",
                  include_mathjax = "cdn")
for k = 1:5
    global plt2
    plt2 = Plots.plot!(plt2, 3:depth, IPR[k, 3:end],
                       label = "k = $(2*k)", yaxis = :log, color=k)
    plt2 = Plots.plot!(plt2, 3:depth, IPR[k, 3:end]+IPR_error[k, 3:end],
                       label = :none, yaxis = :log,
                       seriestype = :scatter, markershape = :hline, color=k)
end
display(plt2)
Plots.savefig("IPR_convergence.png")


# Conjugates related to normal distribution
#
#	Normal - Normal (with known sigma)
#	InverseGamma - Normal (with known mu)
#	NormalInverseGamma - Normal
#

#### Normal prior on μ (with σ^2 known)

function posterior_canon(pri::Normal, ss::NormalKnownSigmaStats)
	p0 = 1.0 / abs2(pri.σ)
	p1 = 1.0 / abs2(ss.σ)
	h = pri.μ * p0 + ss.sx * p1
	prec = p0 + ss.tw * p1
	NormalCanon(h, prec)
end

function posterior_canon(pri::(Normal, Float64), G::Type{Normal}, x::Array)
	μpri::Normal, σ::Float64 = pri
	posterior_canon(μpri, suffstats(NormalKnownSigma(σ), x))
end

function posterior_canon(pri::(Normal, Float64), G::Type{Normal}, x::Array, w::Array{Float64})
	μpri::Normal, σ::Float64 = pri
	posterior_canon(μpri, suffstats(NormalKnownSigma(σ), x, w))
end

function posterior(pri::(Normal, Float64), G::Type{Normal}, x::Array)
	convert(Normal, posterior_canon(pri, G, x))
end

function posterior(pri::(Normal, Float64), G::Type{Normal}, x::Array, w::Array{Float64})
	convert(Normal, posterior_canon(pri, G, x, w))
end

complete(G::Type{Normal}, pri::(Normal, Float64), μ::Float64) = Normal(μ, pri[2])

#### InverseGamma on σ^2 (with μ known)

function posterior_canon(pri::InverseGamma, ss::NormalKnownMuStats)
	α1 = pri.shape + 0.5 * ss.tw
	β1 = pri.scale + 0.5 * ss.s2
	return InverseGamma(α1, β1)
end

function posterior_canon(pri::(Float64, InverseGamma), G::Type{Normal}, x::Array)
	μ::Float64, σ2pri::InverseGamma = pri
	posterior_canon(σ2pri, suffstats(NormalKnownMu(μ), x))
end

function posterior_canon(pri::(Float64, InverseGamma), G::Type{Normal}, x::Array, w::Array{Float64})
	μ::Float64, σ2pri::InverseGamma = pri
	posterior_canon(σ2pri, suffstats(NormalKnownMu(μ), x, w))
end

posterior(pri::(Float64, InverseGamma), G::Type{Normal}, x::Array) = posterior_canon(pri, G, x)
posterior(pri::(Float64, InverseGamma), G::Type{Normal}, x::Array, w::Array{Float64}) = posterior_canon(pri, G, x, w)

complete(G::Type{Normal}, pri::(Float64, InverseGamma), σ2::Float64) = Normal(pri[1], sqrt(σ2))

#### Gamma on precision σ^(-2) (with μ known)

function posterior_canon(pri::Gamma, ss::NormalKnownMuStats)
	α1 = pri.shape + 0.5 * ss.tw
	β1 = pri.scale + 0.5 * ss.s2
	return Gamma(α1, β1)
end

function posterior_canon(pri::(Float64, Gamma), G::Type{Normal}, x::Array)
	μ::Float64, τpri::Gamma = pri
	posterior_canon(τpri, suffstats(NormalKnownMu(μ), x))
end

function posterior_canon(pri::(Float64, Gamma), G::Type{Normal}, x::Array, w::Array{Float64})
	μ::Float64, τpri::Gamma = pri
	posterior_canon(τpri, suffstats(NormalKnownMu(μ), x, w))
end

posterior(pri::(Float64, Gamma), G::Type{Normal}, x::Array) = posterior_canon(pri, G, x)
posterior(pri::(Float64, Gamma), G::Type{Normal}, x::Array, w::Array{Float64}) = posterior_canon(pri, G, x, w)

complete(G::Type{Normal}, pri::(Float64, Gamma), τ::Float64) = Normal(pri[1], 1.0 / sqrt(τ))

#### NormalInverseGamma on (μ, σ^2)

function posterior_canon(prior::NormalInverseGamma, ss::NormalStats)
    mu0 = prior.mu
    v0 = prior.v0
    shape0 = prior.shape
    scale0 = prior.scale

    # ss.tw contains the number of observations if weight wasn't used to
    # compute the sufficient statistics.

    vn_inv = 1./v0 + ss.tw
    mu = (mu0/v0 + ss.s) / vn_inv  # ss.s = ss.tw*ss.m = n*xbar
    shape = shape0 + 0.5*ss.tw
    scale = scale0 + 0.5*ss.s2 + 0.5/(vn_inv*v0)*ss.tw*(ss.m-mu0).^2

    return NormalInverseGamma(mu, 1./vn_inv, shape, scale)
end

complete(G::Type{Normal}, pri::NormalInverseGamma, t::(Float64, Float64)) = Normal(t[1], sqrt(t[2]))

#### NormalGamma on (μ, σ^(-2))

function posterior_canon(prior::NormalGamma, ss::NormalStats)
    mu0 = prior.mu
    nu0 = prior.nu
    shape0 = prior.shape
    rate0 = prior.rate

    # ss.tw contains the number of observations if weight wasn't used to
    # compute the sufficient statistics.

    nu = nu0 + ss.tw
    mu = (nu0*mu0 + ss.s) / nu
    shape = shape0 + 0.5*ss.tw
    rate = rate0 + 0.5*ss.s2 + 0.5*nu0/nu*ss.tw*(ss.m-mu0).^2

    return NormalGamma(mu, nu, shape, rate)
end

complete(G::Type{Normal}, pri::NormalGamma, t::(Float64, Float64)) = Normal(t[1], 1. / sqrt(t[2]))


# Tests for Profitability DEA Models
@testset "ProfitabilityDEAModel" begin

    ## Test Profitability DEA Model with Zofío and Prieto (2006) data
    X = [5 3; 2 4; 4 2; 4 8; 7 9]
    Y = [7 4; 10 8; 8 10; 5 4; 3 6]
    W = [2 1; 2 1; 2 1; 2 1; 2 1.0]
    P = [3 2; 3 2; 3 2; 3 2; 3 2.0]

    # alpha = 0.5 CRS equals Input Oriented CRS
    deaprofbl = deaprofitability(X, Y, W, P)

    @test typeof(deaprofbl) == ProfitabilityDEAModel

    @test nobs(deaprofbl) == 5
    @test ninputs(deaprofbl) == 2
    @test noutputs(deaprofbl) == 2

    @test efficiency(deaprofbl) ≈ [0.388;
                                   1.000;
                                   0.765;
                                   0.250;
                                   0.159] atol = 1e-3
    @test efficiency(deaprofbl, :CRS) ≈ [0.636;
                                1.000;
                                1.000;
                                0.250;
                                0.261] atol = 1e-3
    @test efficiency(deaprofbl, :VRS) ≈ [0.682;
                                1.000;
                                1.000;
                                0.250;
                                0.360] atol = 1e-3
    @test efficiency(deaprofbl, :Scale) ≈ [0.933;
                                1.000;
                                1.000;
                                1.000;
                                0.725] atol = 1e-3
    @test efficiency(deaprofbl, :Allocative) ≈ [0.610;
                                1.000;
                                0.765;
                                1.000;
                                0.609] atol = 1e-3

    @test efficiency(deaprofitability(targets(deaprofbl, :X), targets(deaprofbl, :Y), W, P)) ≈ ones(5) atol=1e-7

    # Check defaults
    @test efficiency(deaprofitability(X, Y, W, P, alpha = 0.5)) == efficiency(deaprofbl)
    @test efficiency(deaprofbl, :Economic) == efficiency(deaprofbl)

    # Print
    show(IOBuffer(), deaprofbl)

    # Test errors
    @test_throws DimensionMismatch deaprofitability([1; 2 ; 3], [4 ; 5], [1; 1; 1], [4; 5]) #  Different number of observations
    @test_throws DimensionMismatch deaprofitability([1; 2; 3], [4; 5; 6], [1; 2; 3; 4], [4; 5; 6]) # Different number of observation in input prices
    @test_throws DimensionMismatch deaprofitability([1; 2; 3], [4; 5; 6], [1; 2; 3], [4; 5; 6; 7]) # Different number of observation in output prices
    @test_throws DimensionMismatch deaprofitability([1 1; 2 2; 3 3], [4; 5; 6], [1 1 1; 2 2 2; 3 3 3], [4; 5; 6]) # Different number of input prices and inputs
    @test_throws DimensionMismatch deaprofitability([1; 2; 3], [4 4; 5 5; 6 6], [1; 2; 3], [4 4 4; 5 5 5; 6 6 6]) # Different number of oputput prices and outputs
    @test_throws ArgumentError efficiency(deaprofbl, :Error)
    @test_throws ArgumentError normfactor(deaprofitability(X, Y, W, P)) # ERROR: ProfitabilityDEAModel has no normalization factor

    # ------------------
    # Test Vector and Matrix inputs and outputs
    # ------------------

    # Inputs is Matrix, Outputs is Vector
    X = [2 2; 1 4; 4 1; 4 3; 5 5; 6 1; 2 5; 1.6	8]
    Y = [1; 1; 1; 1; 1; 1; 1; 1]
    W = [1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1]
    P = [1; 1; 1; 1; 1; 1; 1; 1]

    @test efficiency(deaprofitability(X, Y, W, P)) ≈ ( sum(Y .* P, dims = 2) ./ sum(X .* W, dims = 2) )  / 0.25 atol = 1e-5

    # Inputs is Vector, Output is Matrix
    X = [1; 1; 1; 1; 1; 1; 1; 1]
    Y = [7 7; 4 8; 8 4; 3 5; 3 3; 8 2; 6 4; 1.5 5]
    W = [1; 1; 1; 1; 1; 1; 1; 1]
    P = [1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1; 1 1]

    @test efficiency(deaprofitability(X, Y, W, P)) ≈ ( sum(Y .* P, dims = 2) ./ sum(X .* W, dims = 2) )  / 14 atol = 1e-5

    # Inputs is Vector, Output is Vector
    X = [2; 4; 8; 12; 6; 14; 14; 9.412]
    Y = [1; 5; 8; 9; 3; 7; 9; 2.353]
    W = [1; 1; 1; 1; 1; 1; 1; 1]
    P = [1; 1; 1; 1; 1; 1; 1; 1]

    @test efficiency(deaprofitability(X, Y, W, P)) ≈ ( sum(Y .* P, dims = 2) ./ sum(X .* W, dims = 2) )  / 1.25 atol = 1e-5

end


@testset "FiniteDiff comparison" begin
    q = [1,2,3,4.0]
    q = Quaternion(q ./ norm(q)...,true)
    @test norm(Dojo.dmrpdq(Dojo.vector(q)) -
        ForwardDiff.jacobian(Dojo.mrp, Dojo.vector(q)), Inf) < 1.0e-8
    @test norm(Dojo.daxisdq(Dojo.vector(q)) -
        ForwardDiff.jacobian(Dojo.axis, Dojo.vector(q))) < 1.0e-8
    @test norm(Dojo.drotation_vectordq(Dojo.vector(q)) -
        ForwardDiff.jacobian(Dojo.rotation_vector, Dojo.vector(q))) < 1.0e-8
    q = one(Quaternion{Float64})
    @test norm(Dojo.dmrpdq(Dojo.vector(q)) -
        ForwardDiff.jacobian(Dojo.mrp, Dojo.vector(q)), Inf) < 1.0e-8
    # TODO For zero rotation, ForwardDiff runs into a numeric singularity
    @test norm(Dojo.drotation_vectordq(Dojo.vector(q)) -
        FiniteDiff.finite_difference_jacobian(Dojo.rotation_vector, Dojo.vector(q))) < 1.0e-5
end

@testset "axes pair to quaternion" begin
    n1 = [0,0,0.0]
    n2 = [0,0,0.0]
    q = Dojo.axes_pair_to_quaternion(n1, n2)
    @test norm(n2 - Dojo.vector_rotate(n1, q), Inf) < 1e-8

    n1 = rand(3)
    n1 /= norm(n1)
    n2 = n1
    q = Dojo.axes_pair_to_quaternion(n1, n2)
    @test norm(n2 - Dojo.vector_rotate(n1, q), Inf) < 1e-8

    n1 = rand(3)
    n1 /= norm(n1)
    n2 = -n1
    q = Dojo.axes_pair_to_quaternion(n1, n2)
    @test norm(n2 - Dojo.vector_rotate(n1, q), Inf) < 1e-5

    n1 = rand(3)
    n1 /= norm(n1)
    n2 = rand(3)
    n2 /= norm(n2)
    q = Dojo.axes_pair_to_quaternion(n1, n2)
    @test norm(n2 - Dojo.vector_rotate(n1, q), Inf) < 1e-8
end


using JET
using BenchmarkTools
using Combinatorics
include("auxfun.jl")

function solution1(n=1000)
    s = 0
    for i in 1:n-1
        if mod(i, 3) == 0 || mod(i, 5) == 0
            s += i
        end
    end
    s
end

"""

"""
function solution2(n=4000000)
    a = 1
    b = 2
    s = 0
    while b <= n
        if iseven(b)
            s += b
        end
        a, b = b, a + b
    end
    s
end

function solution3(n=600851475143) # Pollard Rho因数分解
    k = 2
    while n != k
        if n % k == 0
            n = round(Int, n / k)
        else
            k += k == 2 ? 1 : 2
        end
    end
return k
end

function solution4(n=3)
    maxprod=1
    for i in 10^n - 1:-1:10^(n - 1)
        for j in i:-1:10^(n - 1)
            product = i * j
            str= "$product"
            if str == reverse(str) && product >maxprod
                maxprod=product
            end
        end
    end
    maxprod
end

function solution5(n=20)
    lcm(1:n)
end

solution6(n=100) = sum(1:n)^2 - sum((1:n).^2)
function solution7(n=10001)
    primes = zeros(Int, n)
    primes[1:2] .= [2,3]
    k = 3
    counter = 2
    while counter < n
        k += 2
        is_prime = true
        for i in primes[2:counter]
            if k % i == 0
                is_prime = false
                break
            end
        end
        if is_prime
            counter += 1
            primes[counter] = k
        end
    end
    k
end

function solution8(cont=13)
    str = """
    73167176531330624919225119674426574742355349194934
96983520312774506326239578318016984801869478851843
85861560789112949495459501737958331952853208805511
12540698747158523863050715693290963295227443043557
66896648950445244523161731856403098711121722383113
62229893423380308135336276614282806444486645238749
30358907296290491560440772390713810515859307960866
70172427121883998797908792274921901699720888093776
65727333001053367881220235421809751254540594752243
52584907711670556013604839586446706324415722155397
53697817977846174064955149290862569321978468622482
83972241375657056057490261407972968652414535100474
82166370484403199890008895243450658541227588666881
16427171479924442928230863465674813919123162824586
17866458359124566529476545682848912883142607690042
24219022671055626321111109370544217506941658960408
07198403850962455444362981230987879927244284909188
84580156166097919133875499200524063689912560717606
05886116467109405077541002256983155200055935729725
71636269561882670428252483600823257530420752963450
    """
    str = strip(replace(str, "\n" => ""))
    max_prod = 0
    for i in 1:length(str) - cont
        s = prod(parse(Int, num) for num in str[i:i + cont - 1])
        if s > max_prod
            max_prod = s
        end
    end
    max_prod
end

function solution9(n=1000)
    for i in 1:floor(Int,n/3)
        for j in i:floor(Int,(n-i)/2)
            if  i^2+j^2==(n-j-i)^2
                return i*j*(n-i-j)
            end
        end
    end
end

function solution10(n=2000000)
    return sum(sieve(n))
end

function solution11(n=4)
    mat=[
08 02 22 97 38 15 00 40 00 75 04 05 07 78 52 12 50 77 91 08
49 49 99 40 17 81 18 57 60 87 17 40 98 43 69 48 04 56 62 00
81 49 31 73 55 79 14 29 93 71 40 67 53 88 30 03 49 13 36 65
52 70 95 23 04 60 11 42 69 24 68 56 01 32 56 71 37 02 36 91
22 31 16 71 51 67 63 89 41 92 36 54 22 40 40 28 66 33 13 80
24 47 32 60 99 03 45 02 44 75 33 53 78 36 84 20 35 17 12 50
32 98 81 28 64 23 67 10 26 38 40 67 59 54 70 66 18 38 64 70
67 26 20 68 02 62 12 20 95 63 94 39 63 08 40 91 66 49 94 21
24 55 58 05 66 73 99 26 97 17 78 78 96 83 14 88 34 89 63 72
21 36 23 09 75 00 76 44 20 45 35 14 00 61 33 97 34 31 33 95
78 17 53 28 22 75 31 67 15 94 03 80 04 62 16 14 09 53 56 92
16 39 05 42 96 35 31 47 55 58 88 24 00 17 54 24 36 29 85 57
86 56 00 48 35 71 89 07 05 44 44 37 44 60 21 58 51 54 17 58
19 80 81 68 05 94 47 69 28 73 92 13 86 52 17 77 04 89 55 40
04 52 08 83 97 35 99 16 07 97 57 32 16 26 26 79 33 27 98 66
88 36 68 87 57 62 20 72 03 46 33 67 46 55 12 32 63 93 53 69
04 42 16 73 38 25 39 11 24 94 72 18 08 46 29 32 40 62 76 36
20 69 36 41 72 30 23 88 34 62 99 69 82 67 59 85 74 04 36 16
20 73 35 29 78 31 90 01 74 31 49 71 48 86 81 16 23 57 05 54
01 70 54 71 83 51 54 69 16 92 33 48 61 43 52 01 89 19 67 48]
    max_t2b=maximum(prod(mat[i:i+n-1,j])  for j in 1:size(mat,2) for i in 1:size(mat,1)-n+1)
    max_l2r=maximum(prod(mat[i,j:j+n-1]) for j in 1:size(mat,2)-n+1  for i in 1:size(mat,1))
    max_diag=maximum(prod(mat[i+k,j+k] for k in 0:n-1) for j in 1:size(mat,2)-n+1 for i in 1:size(mat,1)-n+1)
    max_subdiag=maximum(prod(mat[i+k,j-k] for k in 0:n-1) for j in n:size(mat,2) for i in 1:size(mat,1)-n+1)
    maximum([max_t2b,max_l2r,max_diag,max_subdiag])
end

function solution12(n=500)
    tri_num(n)=round(Int,n*(n+1)/2)
    i=2
    K=tri_num(i)
    N=num_of_factors(K)
    while N<n
        i+=1
        K=tri_num(i)
        N=num_of_factors(K)
    end
    K
end

function solution13(n=10)
    nums=[
        37107287533902102798797998220837590246510135740250
        46376937677490009712648124896970078050417018260538
        74324986199524741059474233309513058123726617309629
        91942213363574161572522430563301811072406154908250
        23067588207539346171171980310421047513778063246676
        89261670696623633820136378418383684178734361726757
        28112879812849979408065481931592621691275889832738
        44274228917432520321923589422876796487670272189318
        47451445736001306439091167216856844588711603153276
        70386486105843025439939619828917593665686757934951
        62176457141856560629502157223196586755079324193331
        64906352462741904929101432445813822663347944758178
        92575867718337217661963751590579239728245598838407
        58203565325359399008402633568948830189458628227828
        80181199384826282014278194139940567587151170094390
        35398664372827112653829987240784473053190104293586
        86515506006295864861532075273371959191420517255829
        71693888707715466499115593487603532921714970056938
        54370070576826684624621495650076471787294438377604
        53282654108756828443191190634694037855217779295145
        36123272525000296071075082563815656710885258350721
        45876576172410976447339110607218265236877223636045
        17423706905851860660448207621209813287860733969412
        81142660418086830619328460811191061556940512689692
        51934325451728388641918047049293215058642563049483
        62467221648435076201727918039944693004732956340691
        15732444386908125794514089057706229429197107928209
        55037687525678773091862540744969844508330393682126
        18336384825330154686196124348767681297534375946515
        80386287592878490201521685554828717201219257766954
        78182833757993103614740356856449095527097864797581
        16726320100436897842553539920931837441497806860984
        48403098129077791799088218795327364475675590848030
        87086987551392711854517078544161852424320693150332
        59959406895756536782107074926966537676326235447210
        69793950679652694742597709739166693763042633987085
        41052684708299085211399427365734116182760315001271
        65378607361501080857009149939512557028198746004375
        35829035317434717326932123578154982629742552737307
        94953759765105305946966067683156574377167401875275
        88902802571733229619176668713819931811048770190271
        25267680276078003013678680992525463401061632866526
        36270218540497705585629946580636237993140746255962
        24074486908231174977792365466257246923322810917141
        91430288197103288597806669760892938638285025333403
        34413065578016127815921815005561868836468420090470
        23053081172816430487623791969842487255036638784583
        11487696932154902810424020138335124462181441773470
        63783299490636259666498587618221225225512486764533
        67720186971698544312419572409913959008952310058822
        95548255300263520781532296796249481641953868218774
        76085327132285723110424803456124867697064507995236
        37774242535411291684276865538926205024910326572967
        23701913275725675285653248258265463092207058596522
        29798860272258331913126375147341994889534765745501
        18495701454879288984856827726077713721403798879715
        38298203783031473527721580348144513491373226651381
        34829543829199918180278916522431027392251122869539
        40957953066405232632538044100059654939159879593635
        29746152185502371307642255121183693803580388584903
        41698116222072977186158236678424689157993532961922
        62467957194401269043877107275048102390895523597457
        23189706772547915061505504953922979530901129967519
        86188088225875314529584099251203829009407770775672
        11306739708304724483816533873502340845647058077308
        82959174767140363198008187129011875491310547126581
        97623331044818386269515456334926366572897563400500
        42846280183517070527831839425882145521227251250327
        55121603546981200581762165212827652751691296897789
        32238195734329339946437501907836945765883352399886
        75506164965184775180738168837861091527357929701337
        62177842752192623401942399639168044983993173312731
        32924185707147349566916674687634660915035914677504
        99518671430235219628894890102423325116913619626622
        73267460800591547471830798392868535206946944540724
        76841822524674417161514036427982273348055556214818
        97142617910342598647204516893989422179826088076852
        87783646182799346313767754307809363333018982642090
        10848802521674670883215120185883543223812876952786
        71329612474782464538636993009049310363619763878039
        62184073572399794223406235393808339651327408011116
        66627891981488087797941876876144230030984490851411
        60661826293682836764744779239180335110989069790714
        85786944089552990653640447425576083659976645795096
        66024396409905389607120198219976047599490197230297
        64913982680032973156037120041377903785566085089252
        16730939319872750275468906903707539413042652315011
        94809377245048795150954100921645863754710598436791
        78639167021187492431995700641917969777599028300699
        15368713711936614952811305876380278410754449733078
        40789923115535562561142322423255033685442488917353
        44889911501440648020369068063960672322193204149535
        41503128880339536053299340368006977710650566631954
        81234880673210146739058568557934581403627822703280
        82616570773948327592232845941706525094512325230608
        22918802058777319719839450180888072429661980811197
        77158542502016545090413245809786882778948721859617
        72107838435069186155435662884062257473692284509516
        20849603980134001723930671666823555245252804609722
        53503534226472524250874054075591789781264330331690
        ]
    s=sum(nums)
    return parse(Int,"$s"[1:n])
end
function solution14(n=1000000)
    n>1 || return 1
    f(n)= n%2==0 ? round(Int,n/2) : 3n+1
    function l(n)
        n> 1 || return 1
        a=n
        l=1
        while f(a) >1
            a=f(a)
            l+=1
        end
        l+=1
    end
    lengths=l.(2:n)
    findfirst(lengths.== maximum(lengths))+1 # 从2开始
end

function solution15(n=20)
    dict=Dict(i=> true for i in 1:n)
    s=1
    for j in n+1:2n
        s*=j
        for i in 1:n
            if dict[i] && s%i==0
                s=round(Int,s/i)
                dict[i]=false
            end
        end
    end
    s
end

solution16(n=1000)=sum(digits((BigInt(2))^n))
function solution17(n=1000)
    number2word =Dict(
    0=>"",
    1=>"one",
    2=>"two",
    3=>"three",
    4=>"four",
    5=>"five",
    6=>"six",
    7=>"seven",
    8=>"eight",
    9=>"nine",
    10=>"ten",
    11=>"eleven",
    12=>"twelve",
    13=>"thirteen",
    14=>"fourteen",
    15=>"fifteen",
    16=>"sixteen",
    17=>"seventeen",
    18=>"eighteen",
    19=>"nineteen",
    20=>"twenty",
    30=>"thirty",
    40=>"forty",
    50=>"fifty",
    60=>"sixty",
    70=>"seventy",
    80=>"eighty",
    90=>"ninety",
    100=>"hundred",
    1000=>"one thousand",
    )
    function num2text(m)
        @assert m <=1000
        m==100 && return "one hundred"
        if m in keys(number2word)
            return number2word[m]
        elseif m <100
            d=m -m %10
            m=m-d
            return num2text(d)*" "* num2text(m)
        else
            h=m-m%100
            m=m-h
            h=div(h,100)

            return num2text(h)*" hundred"*(m!=0 ? " and " *num2text(m) : "")
        end
    end
    lettercount=0
    for i in 1:n
        lettercount+=length(replace(num2text(i)," "=>""))
    end
    lettercount
end

function solution18()
    nums = [
        [75],
        [95 64],
        [17 47 82],
        [18 35 87 10],
        [20 04 82 47 65],
        [19 01 23 75 03 34],
        [88 02 77 73 07 63 67],
        [99 65 04 28 06 16 70 92],
        [41 41 26 56 83 40 80 70 33],
        [41 48 72 33 47 32 37 16 94 29],
        [53 71 44 65 25 43 91 52 97 51 14],
        [70 11 33 28 77 73 17 78 39 68 17 57],
        [91 71 52 38 17 14 91 43 58 50 27 29 48],
        [63 66 04 68 89 53 67 30 73 16 69 87 40 31],
        [04 62 98 27 23 09 70 98 73 93 38 53 60 04 23]]
    rows=length(nums)
    rows> 1 || return maximum(nums[1])
    cols=maximum(length.(nums))
    dist=zeros(Int,(rows,cols))
    dist[1,1]=nums[1][1]
    for i in 2: rows
        dist[i,1]=nums[i][1]+dist[i-1][1]
    end
    for i in 2:rows
        for j in 2:length(nums[i])
            dist[i,j]=nums[i][j]+max(dist[i-1,j],dist[i-1,j-1])
        end
    end
    maximum(dist[rows,:])
end

function solution19()
    Y=1901
    M=1
    D=1
    W=1
    cnt=0
    leap_year(Y)= (Y%4==0 && Y%100 !=0) || Y%400==0
    function days_in_month(M,Y)
        dict=Dict{Int,Int}(1=>31, 2=> 28,3=>31,4=>30,5=>31,6=>30,7=>31,8=>31,9=>30,10=>31,11=>30,12=>31)
        if leap_year(Y)
            dict[2]+=1
        end
        dict[M]
    end
    function date_add_one(Y,M,D)
        days=days_in_month(M,Y)
        if D==days
            D=1
            if M<12
                M+=1
            else
                Y+=1
                M=1
            end
        else
            D+=1
        end
        Y,M,D
    end
    weekday_add_one(W)= W<7 ? W+1 : 1
    while Y<=2000 && M<=12 && D<=31
        if D==1 && W==7
            cnt+=1
        end
        Y,M,D=date_add_one(Y,M,D)
        W=weekday_add_one(W)
    end
    cnt
end

solution20(n=100)=sum(digits(prod(1:BigInt(n))))

function solution21(n=10000)
    d(n)::Int=sum(x for x in 1:n-1 if n%x==0) # n>=2
    s=0
    for i in 2:n
        t=d(i)
        if t >1 && t <=n && d(t)==i && t!=i # 一个数不能构成亲和数
            s+=i
        end
    end
    s
end

function solution22()
    char2num(s)=Int(lowercase(s))-Int('a')+1
    name_score(name,idx)=sum(char2num(c) for c in name)*idx
    names = sort(split(replace(read("names.txt",String),"\""=>""),","))
    sum(name_score(names[i],i) for i in 1:length(names))
end

function solution23()
    abundants=(1:28123)[is_abundant.(1:28123)]
    function is_sum_of_two_abundant(n)
        n>=24 || return false
        for i in abundants
            if   i <n && (n-i) in abundants
                return true
            end
        end
        return false
    end
    sum(x for x in 1:28123 if !is_sum_of_two_abundant(x))
end

function solution24(n=1000000)
    nums=collect(0:9)
    P=permutations(nums)
    i=0
    for p in P
        i+=1
        if i ==n
            return sum(collect(p) .* [10^(10-i) for i in 1:10])
        end
    end
end

function solution25(dig=1000)
    a = BigInt(1)
    b = BigInt(1)
    n = 2
    digits_n=length(digits(b))
    while digits_n <dig
        a,b=b,a+b
        n+=1
        digits_n=length(digits(b))
    end
    return n
end

function  solution26(n=1000)
    lens=cyclelen.(1:1000)
    findmax(lens)[2]
end
solution26()

function solution27(ub=1000)
    maxcnt=0
    a0=0
    b0=0
    for a in -(ub-1):(ub-1)
        for b in -(ub-1):(ub-1)
            mycnt = primecounter(n-> n*n+a*n+b)
            if mycnt >= maxcnt
                a0=a
                b0=b
                maxcnt=mycnt
            end
        end
    end
    return a0*b0
end
# @enter solution27()

solution27()

function solution28(w::Integer=1001)
    @assert isodd(w)
    a(n)=n*n # anti-clockwise
    b(n)=a(n)-n+1
    c(n)=b(n)-n+1
    d(n)=c(n)-n+1
    s=0
    for i in 1:div(w,2)+1
        s+=a(2i-1)+b(2i-1)+c(2i-1)+d(2i-1)
    end
    s-=3 # 1 has been added for 4 times
end

solution28()

function solution29(n=100)
    list=zeros(BigInt,99*99)
    i=1
    for a in 2:100
        for b in 2:100
            list[i]=BigInt(a)^b
            i+=1
        end
    end
    length(unique(list))
end

solution29()

function solution30(k=5)
    d=length(digits(9^k))+1
    sp_nums=Int[]
    for i in 2:10^d
        if i== sum(x->x^k,digits(i))
            push!(sp_nums,i)
        end
    end
    sum(sp_nums)
end

solution30()

function solution31(n::Int=200)
    currency=Int[1,2,5,10,20,50,100,200]
    cnt=0
    for c200 in 0:div(n,200)
        for c100 in 0:div(n,100)
            for c50 in 0:div(n,50)
                for c20 in 0:div(n,20)
                    for c10 in 0:div(n,10)
                        for c5 in 0:div(n,5)
                            for c2 in 0:div(n,2)
                                for c1 in 0:n
                                    t=c1+2c2+5c5+10c10+20c20+50c50+100c100+200c200
                                    if t==n
                                        cnt+=1
                                    elseif t>n
                                        break
                                    end
                                end
                            end
                        end
                    end
                end
            end
        end
    end
    cnt
end

function solution32()
    # This function is fairly slow!
    # 1000 x 1000= 1000,000,  10 digits!
    p=Set(0) # Initialization
    for a in 1:10000
        for b in a:10000
            c=a*b
            digs=vcat(digits(a),digits(b),digits(c))
            if length(digs)==9 && sort(digs)==collect(1:9)
                push!(p,c)
            end
        end
    end
    sum(p)
end

function solution33()
    # 49/98=1/2
    # 30/50=3/5
    # 2*2*5*5=100
    100
end

function solution34()
    factof1to9=collect(0:9)
    factof1to9[1]=1
    for i in 2:10
        factof1to9[i]=factof1to9[i-1]*(i-1)
    end
    s=0
    ub=sum(factof1to9)
    for k in 10:ub
        if k == sum(factof1to9[digits(k).+1])
            s+=k
        end
    end
    s
end

function solution35(n=1000_000)
    cnt=0
    for i in 2:n
        if isprime(i)
            all_prime=true
            for num in cycles(i)
                if !isprime(num)
                    all_prime=false
                    break
                end
            end
            if all_prime
                cnt+=1
            end
        end
    end
    cnt
end

function solution36(n=1000_000)
    s=0
    for i in 1:n
        if !iseven(i)
            tenary=string(i)
            binary=string(i,base=2)
            if tenary==reverse(tenary) && binary ==reverse(binary)
                s+=i
            end
        end
    end
    s
end

function solution37()
    s=0
    cnt=0
    i=11
    while true
        if istruncatableprime(i)
            cnt+=1
            s+=i
        end
        i+=2
        if cnt ==11
            return s
        end
    end
end


struct BernsteinPolynomial{N, T} <: AbstractPolynomial{N, T}
    coeffs::NTuple{N, T}
end

# Construction
BernsteinPolynomial(coeffs::Tuple) = BernsteinPolynomial(promote(coeffs...))
BernsteinPolynomial(coeffs...) = BernsteinPolynomial(coeffs)

# Conversion to monomial basis
Polynomial(b::BernsteinPolynomial{1}) = Polynomial(b.coeffs)
function Polynomial(b::BernsteinPolynomial{N, T}) where {N, T}
    # b(t) = (1 - t) * b1(t) + t * b2(t)
    #      = b1(t) + t * (b2(t) - b1(t))
    b1 = BernsteinPolynomial(reverse(tail(reverse(b.coeffs))))
    b2 = BernsteinPolynomial(tail(b.coeffs))
    Polynomial(b1) + Polynomial(0, 1) * Polynomial(b2 - b1)
end

# Utility
@inline constant(b::BernsteinPolynomial) = b.coeffs[1]
Base.zero(::Type{BernsteinPolynomial{N, T}}) where {N, T} = BernsteinPolynomial(ntuple(_ -> zero(T), Val(N)))
Base.zero(b::BernsteinPolynomial) = zero(typeof(b))

# Evaluation
@inline (b::BernsteinPolynomial{1})(t) = constant(b)
@inline function (b::BernsteinPolynomial)(t)
    b1 = BernsteinPolynomial(reverse(tail(reverse(b.coeffs))))
    b2 = BernsteinPolynomial(tail(b.coeffs))
    (oneunit(t) - t) * b1(t) + t * b2(t)
end

for op in [:+, :-]
    @eval begin
        # Two BernsteinPolynomials
        Base.$op(b1::BernsteinPolynomial{N}, b2::BernsteinPolynomial{N}) where {N} = BernsteinPolynomial(_map($op, b1.coeffs, b2.coeffs))

        # BernsteinPolynomials and constant
        Base.$op(b::BernsteinPolynomial, c) = BernsteinPolynomial(_map(p -> $op(p, c), b.coeffs))
        Base.$op(c, b::BernsteinPolynomial) = BernsteinPolynomial(_map(p -> $op(c, p), b.coeffs))
    end
end

for op in [:*, :/]
    @eval Base.$op(b::BernsteinPolynomial, c) = BernsteinPolynomial(_map(x -> $op(x, c), b.coeffs))
end
Base.:*(c, b::BernsteinPolynomial) = BernsteinPolynomial(_map(x -> c * x, b.coeffs))

@inline function derivative(b::BernsteinPolynomial{N}) where {N}
    # https://pages.mtu.edu/~shene/COURSES/cs3621/NOTES/spline/Bezier/bezier-der.html
    coeffs = b.coeffs
    n = N - 1
    BernsteinPolynomial(ntuple(i -> n * (coeffs[i + 1] - coeffs[i]), n))
end

"""
```math
\\int_{0}^{1} b\\left(\\tau\\right) e^{c \\tau} d\\tau
```
"""
@inline function exponential_integral(b::BernsteinPolynomial{N}, c; inv_c=inv(c), exp_c=exp(c)) where N
    pn = b.coeffs[N]
    if N === 1
        return inv_c * constant(b) * (exp_c - 1)
    else
        return inv_c * (pn * exp_c - constant(b) - exponential_integral(derivative(b), c, inv_c=inv_c, exp_c=exp_c))
    end
end


### A Pluto.jl notebook ###
# v0.12.18

using Markdown
using InteractiveUtils

# ╔═╡ 0b0902da-40f0-11eb-3e10-a1fe1bd1f364
begin
	using PyPlot
	using SpecialFunctions
	using HypergeometricFunctions
end

# ╔═╡ 0b507856-40f1-11eb-244f-f15c1038b96f
begin
	rc("text", usetex = true)
	ion()
end

# ╔═╡ eaadeff2-546b-11eb-1289-9bade75bad45
function J(α, β, n, z)
    T = eltype(z/2)
    tot = zero(T)
    for k in 0:n
        tot += binomial(α + n, k) * binomial(β + n, n - k) * ((z + 1)/2)^k * ((z - 1)/2)^(n - k)
    end
    tot
end;

# ╔═╡ Cell order:
# ╠═0b0902da-40f0-11eb-3e10-a1fe1bd1f364
# ╠═0b507856-40f1-11eb-244f-f15c1038b96f
# ╠═eaadeff2-546b-11eb-1289-9bade75bad45


"""
    Cylinder{N, T}

A `Cylinder` is a 2D rectangle or a 3D cylinder defined by its origin point,
its extremity and a radius. `origin`, `extremity` and `r`, must be specified.
"""
struct Cylinder{N,T} <: GeometryPrimitive{N,T}
    origin::Point{N,T}
    extremity::Point{N,T}
    r::T
end

"""
    Cylinder2{T}
    Cylinder3{T}

A `Cylinder2` or `Cylinder3` is a 2D/3D cylinder defined by its origin point,
its extremity and a radius. `origin`, `extremity` and `r`, must be specified.
"""
const Cylinder2{T} = Cylinder{2,T}
const Cylinder3{T} = Cylinder{3,T}

origin(c::Cylinder{N,T}) where {N,T} = c.origin
extremity(c::Cylinder{N,T}) where {N,T} = c.extremity
radius(c::Cylinder{N,T}) where {N,T} = c.r
height(c::Cylinder{N,T}) where {N,T} = norm(c.extremity - c.origin)
direction(c::Cylinder{N,T}) where {N,T} = (c.extremity .- c.origin) ./ height(c)

function rotation(c::Cylinder{2,T}) where {T}
    d2 = direction(c)
    u = @SVector [d2[1], d2[2], T(0)]
    v = @MVector [u[2], -u[1], T(0)]
    normalize!(v)
    return hcat(v, u, @SVector T[0, 0, 1])
end

function rotation(c::Cylinder{3,T}) where {T}
    d3 = direction(c)
    u = @SVector [d3[1], d3[2], d3[3]]
    if abs(u[1]) > 0 || abs(u[2]) > 0
        v = @MVector [u[2], -u[1], T(0)]
    else
        v = @MVector [T(0), -u[3], u[2]]
    end
    normalize!(v)
    w = @SVector [u[2] * v[3] - u[3] * v[2], -u[1] * v[3] + u[3] * v[1],
                  u[1] * v[2] - u[2] * v[1]]
    return hcat(v, w, u)
end

function coordinates(c::Cylinder{2,T}, nvertices=(2, 2)) where {T}
    r = Rect(c.origin[1] - c.r / 2, c.origin[2], c.r, height(c))
    M = rotation(c)
    points = coordinates(r, nvertices)
    vo = to_pointn(Point3{T}, origin(c))
    return (M * (to_pointn(Point3{T}, point) .- vo) .+ vo for point in points)
end

function faces(sphere::Cylinder{2}, nvertices=(2, 2))
    return faces(Rect(0, 0, 1, 1), nvertices)
end

function coordinates(c::Cylinder{3,T}, nvertices=30) where {T}
    if isodd(nvertices)
        nvertices = 2 * (nvertices ÷ 2)
    end
    nvertices = max(8, nvertices)
    nbv = nvertices ÷ 2

    M = rotation(c)
    h = height(c)
    range = 1:(2 * nbv + 2)
    function inner(i)
        return if i == length(range)
            c.extremity
        elseif i == length(range) - 1
            origin(c)
        else
            phi = T((2π * (((i + 1) ÷ 2) - 1)) / nbv)
            up = ifelse(isodd(i), 0, h)
            (M * Point(c.r * cos(phi), c.r * sin(phi), up)) .+ c.origin
        end
    end

    return (inner(i) for i in range)
end

function faces(c::Cylinder{3}, facets=30)
    isodd(facets) ? facets = 2 * div(facets, 2) : nothing
    facets < 8 ? facets = 8 : nothing
    nbv = Int(facets / 2)
    indexes = Vector{TriangleFace{Int}}(undef, facets)
    index = 1
    for j in 1:(nbv - 1)
        indexes[index] = (index + 2, index + 1, index)
        indexes[index + 1] = (index + 3, index + 1, index + 2)
        index += 2
    end
    indexes[index] = (1, index + 1, index)
    indexes[index + 1] = (2, index + 1, 1)

    for i in 1:length(indexes)
        i % 2 == 1 ? push!(indexes, (indexes[i][1], indexes[i][3], 2 * nbv + 1)) :
        push!(indexes, (indexes[i][2], indexes[i][1], 2 * nbv + 2))
    end
    return indexes
end


import Distances: Metric, SqEuclidean, Euclidean

"""
    SMeasureTest(m, τ, K)
    SMeasureTest(m::Int, τ, K; metric::Metric = SqEuclidean(),
        tree_metric::Metric = Euclidean())

S-measure test [1] for the directional dependence between data series.

## Algorithm

1. Create an `m-`dimensional embeddings of both `x` and `y`, resulting in
    `N` different `m`-dimensional embedding points
    ``X = \\{x_1, x_2, \\ldots, x_N \\}`` and ``X = \\{y_1, y_2, \\ldots, y_N \\}``.
    `τ` controls the embedding lag.
2. Let ``r_{i,j}`` and ``s_{i,j}`` be the indices of the `k`-th nearest neighbors
    of ``x_i`` and ``y_i``, respectively.
3. Compute the the mean squared Euclidean distance to the ``k`` nearest neighbors
    for each ``x_i``, using the indices ``r_{i, j}``.

```math
R_i^{(k)}(x) = \\dfrac{1}{k} \\sum_{i=1}^{k}(x_i, x_{r_{i,j}})^2
```

- Compute the y-conditioned mean squared Euclidean distance to the ``k`` nearest
    neighbors for each ``x_i``, now using the indices ``s_{i,j}``.

```math
R_i^{(k)}(x|y) = \\dfrac{1}{k} \\sum_{i=1}^{k}(x_i, x_{s_{i,j}})^2
```

- Define the following measure of independence, where ``0 \\leq S \\leq 1``, and
    low values indicate independence and values close to one occur for
    synchronized signals.

```math
S^{(k)}(x|y) = \\dfrac{1}{N} \\sum_{i=1}^{N} \\dfrac{R_i^{(k)}(x)}{R_i^{(k)}(x|y)}
```

## Examples

```julia
# Initialise test, specifying embedding dimension, embedding lag
# and number of nearest neighbors
test = SMeasureTest(m = 4, τ = 3, K = 2:10)
```

## References

1.  Quian Quiroga, R., Arnhold, J. & Grassberger, P. [2000] “Learning
    driver-response relationships from synchronization patterns,”
    Phys. Rev. E61(5), 5142–5148.
"""
Base.@kwdef struct SMeasureTest{N, M, TM} <: DistanceBasedCausalityTest{N} where {M, TM}
    """ The embedding dimension. Defaults to 2. """
    m::Int = 2

    """ The embedding lag. Defaults to 1. """
    τ::Int = 1

    """ Neighborhood size(s). Defaults to 2:10. """
    K = 2:10

    """ The distance metric used when computing distances. Defaults to `SqEuclidean()`"""
    metric::M = SqEuclidean()

    """
    The distance metric used when computing `KDTree` for nearest neighbor searches.
    Defaults to `Euclidean()`.
    """
    tree_metric::TM = Euclidean()

    function SMeasureTest(m, τ, K, metric::M, tree_metric::TM) where {M, TM}
        N = length(K)
        new{N, M, TM}(m, τ, K, metric, tree_metric)
    end
end

function SMeasureTest(m, τ, K; metric::Metric = SqEuclidean(), tree_metric::Metric = Euclidean())
    SMeasureTest(m = m, τ = τ, K = K; metric = metric, tree_metric = tree_metric)
end

"""

    causality(x, y, test::SMeasureTest)

## Examples

```julia
using CausalityTools, DynamicalSystemsBase, Plots

# Create an orbit of the built-in `henon2` map
npts, Ttr = 5000, 500
x, y = columns(trajectory(henon2(c_xy = 1.0), npts, Ttr = Ttr));

# Some random time series
xr, yr = rand(npts), rand(npts);

# Initialise test, specifying embedding dimension, emebdding lag
# and number of nearest neighbors
Ks = 2:10
test = SMeasureTest(m = 4, τ = 1, K = Ks)

# Compute causality statistic on random time series, and on Henon map
# time series.
Ss_r_xy = causality(xr, yr, test)
Ss_r_yx = causality(yr, xr, test)
Ss_henon_xy = causality(x, y, test)
Ss_henon_yx = causality(y, x, test);

plot(xlabel = "# nearest neighbors (k)", ylabel = "S", ylims = (-0.05, 1.05))
plot!(Ks, Ss_r_xy,  label = "random uncoupled system (x -> y)", marker = stroke(2), c = :black)
plot!(Ks, Ss_r_yx,  label = "random uncoupled system (y -> x)", marker = stroke(2), c = :red)
plot!(Ks, Ss_henon_xy, marker = stroke(2), label = "henon unidir (x -> y)")
plot!(Ks, Ss_henon_yx, marker = stroke(2), label = "henon unidir (y -> x)")
```
"""
function causality(x::AbstractVector, y::AbstractVector, test::SMeasureTest)
    Ss = zeros(Float64, length(test.K))

    for (i, k) in enumerate(test.K)
        Ss[i] = s_measure(x, y,test.m, test.τ, k, test.metric, test.tree_metric)
    end

    return Ss
end

export SMeasureTest, causality


using ActiveSetMethods
using GLPK
using Test

@testset "MathOptInterface" begin
    include("MOI_wrapper.jl")
end

@testset "External Solver Attributes Implementation with Toy Example" begin
    include("ext_solver.jl")
    @test isapprox(xsol, -1.0, rtol=1e-4)
    @test isapprox(ysol, -1.0, rtol=1e-4)
    @test status == MOI.LOCALLY_SOLVED
end

@testset "opf.jl" begin
    include("opf.jl")
    result = run_opf("../examples/acopf/case3.m", 100, "Line Search")
    @test isapprox(result["objective"], +5.90687949e+03, rtol=1e-3)
    result = run_opf("../examples/acopf/case3.m", 100, "Line Search")
    @test isapprox(result["objective"], +5.90687949e+03, rtol=1e-2)
end


py"""
class DiffBertsimas(torch.nn.Module):
    '''
    DiffBertsimas

Differentiable Bertsimas model.

Atributes:
- `D::Int`: dimension (i.e. number of assets)
- `w::PyObject`: differentiable weights
- `c::Vector{PyObject}`: convex constraints on weights
- `μ::Array{Float64,1}`: predicted mean vector
- `sqrtΣ::Array{Float64,2}`: predicted upper triangular Cholesky factor of covariance matrix
- ` solver_args::Dict=Dict("use_indirect" => true, "eps" => 1e-5, "max_iters" => 20000}`: solver arguments
    '''
    def __init__(
        self,
        D,
        w,
        c,
        μ,
        sqrtΣ,
        solver_args={'use_indirect': True, 'eps': 1e-5, 'max_iters': 20000}
    ):
        super(DiffBertsimas, self).__init__()
        self.D = D
        self.μ = μ
        Σ = sqrtΣ.T @ sqrtΣ
        self.Σ = Σ
        self.solver_args=solver_args

        λ = cp.Variable()
        π_neg = cp.Variable(D)
        π_pos = cp.Variable(D)
        θ = cp.Variable(D)

        params = cp.Parameter(D+1)

        constraints = [
            λ >= 0,
            π_neg >= 0,
            π_pos >= 0,
            θ >= 0,
            w == π_pos - π_neg,
            cp.multiply(params[0:D], (π_pos + π_neg)) - θ <= λ,
        ]

        objective = cp.Maximize(μ.T @ w - cp.sum(θ) - λ * params[D])
        problem = cp.Problem(objective, constraints + c)

        self.cvxpylayer = CvxpyLayer(
            problem,
            parameters=[params],
            variables=[w, λ, π_neg, π_pos, θ]
        )

    def forward(self, params):
        return self.cvxpylayer(params, solver_args=self.solver_args)[0]

    def get_variables(self, params):
        params_ = torch.tensor(params, dtype=torch.float64, requires_grad=True)
        return self.cvxpylayer(params_, solver_args=self.solver_args)

    def get_w(self, params):
        params_ = torch.tensor(params, dtype=torch.float64, requires_grad=True)
        return self.get_variables(params_)[0]

    def get_lagrangian(self, params):
        params_ = torch.tensor(params, dtype=torch.float64, requires_grad=True)
        D = self.D
        μ_ = torch.tensor(self.μ, dtype=torch.float64)
        Σ_ = torch.tensor(self.Σ, dtype=torch.float64)
        w_, λ_, π_neg_, π_pos_, θ_ = self.cvxpylayer(params_, solver_args=self.solver_args)
        return μ_.T @ w_ - torch.sum(θ_) - λ_ * params_[D]
"""

DiffBertsimas = py"DiffBertsimas"


@testset "Mixed Integer Programs: $solver" for solver in solvers
    if can_solve_mip(solver)
        mip_solver = !can_solve_mip(solver) ? GLPKSolverMIP() : solver

        @testset "lp fallback interface" begin
            x = Variable()
            p = minimize(x, x>=4.3)
            @test vexity(p) == AffineVexity()
            solve!(p, solver)
            @test p.optval ≈ 4.3 atol=TOL

            x = Variable(2)
            p = minimize(norm(x,1), x[1]>=4.3)
            @test vexity(p) == ConvexVexity()
            solve!(p, solver)
            @test p.optval ≈ 4.3 atol=TOL
        end

        @testset "integer variables" begin
            x = Variable(:Int)
            p = minimize(x, x>=4.3)
            @test vexity(p) == AffineVexity()
            solve!(p, mip_solver)
            @test p.optval ≈ 5 atol=TOL

            x = Variable(2, :Int)
            p = minimize(sum(x), x>=4.3)
            @test vexity(p) == AffineVexity()
            solve!(p, mip_solver)
            @test p.optval ≈ 10 atol=TOL

            x = Variable(:Int)
            y = Variable()
            p = minimize(sum(x + y), x>=4.3, y>=7)
            @test vexity(p) == AffineVexity()
            solve!(p, mip_solver)
            @test p.optval ≈ 12 atol=TOL

            x = Variable(2, :Int)
            p = minimize(norm(x, 1), x[1]>=4.3)
            @test vexity(p) == ConvexVexity()
            solve!(p, mip_solver)
            @test p.optval ≈ 5 atol=TOL

            x = Variable(2, :Int)
            p = minimize(sum(x), x[1]>=4.3, x>=0)
            @test vexity(p) == AffineVexity()
            solve!(p, mip_solver)
            @test p.optval ≈ 5 atol=TOL

            x = Variable(2, :Int)
            p = minimize(sum(x), x>=.5)
            @test vexity(p) == AffineVexity()
            solve!(p, mip_solver)
            @test p.optval ≈ 2 atol=TOL
        end

        @testset "binary variables" begin
            x = Variable(2, :Bin)
            p = minimize(sum(x), x>=.5)
            @test vexity(p) == AffineVexity()
            solve!(p, mip_solver)
            @test p.optval ≈ 2 atol=TOL

            x = Variable(2, :Bin)
            p = minimize(sum(x), x[1]>=.5, x>=0)
            @test vexity(p) == AffineVexity()
            solve!(p, mip_solver)
            @test p.optval ≈ 1 atol=TOL
        end
    end
end


using ReachabilityAnalysis, JLD2
using ReachabilityAnalysis: add_dimension

LazySets.set_ztol(Float64, 1e-15)
ISS_path = joinpath(@__DIR__, "ISS.jld2")

@load ISS_path C
const C3 = C[3, :] # variable y₃
const C3_ext = vcat(C3, fill(0.0, 3))

function ISSF01()
    @load ISS_path A B

    U = Hyperrectangle(low=[0.0, 0.8, 0.9], high=[0.1, 1., 1.]);  # input set
    X0 = BallInf(zeros(size(A, 1)), 0.0001)  # -0.0001 <= xi <= 0.0001 for all i
    prob_ISSF01 = @ivp(x' = A*x + B*u, x(0) ∈ X0, u ∈ U, x ∈ Universe(270))
end

function ISSC01()
    @load ISS_path A B

    Aext = add_dimension(A, 3)
    Aext[1:270, 271:273] = B

    U = Hyperrectangle(low=[0.0, 0.8, 0.9], high=[0.1, 1., 1.])  # input set
    X0 = BallInf(zeros(size(A, 1)), 0.0001)  # -0.0001 <= xi <= 0.0001 for all i
    X0 = X0 * U
    prob_ISSC01 = @ivp(x' = Aext*x, x(0) ∈ X0)
end


p = 2
function pow_array(x::Vector{Float64})
s = 0.0
for y in x
s = s + y^p
end
return s
end


module Statifiers

using StatsBase: std, skewness, kurtosis, variation, sem, mad,  entropy, summarystats, autocor, pacf, rle, quantile
using Dates
using DataFrames
using Random
using Statistics

export fit!,transform!
export Statifier

export fullstat,statifierrun

using TSML.TSMLTypes
import TSML.TSMLTypes.fit! # to overload
import TSML.TSMLTypes.transform! # to overload
using TSML.Utils

# Transforms instances with nominal features into one-hot form
# and coerces the instance matrix to be of element type Float64.
mutable struct Statifier <: Transformer
  model
  args

  function Statifier(args=Dict())
    default_args = Dict(
      :processmissing => true
    )
    new(nothing, mergedict(default_args, args))
  end
end

function fit!(st::Statifier, features::T=[], labels::Vector=[]) where {T<:Union{Vector,Matrix,DataFrame}}
  typeof(features) <: DataFrame || error("Statifier.fit!: data should be a dataframe: Date,Val ")
  ncol(features) == 2 || error("dataframe must have 2 columns: Date, Val")
  st.model = st.args
end

function transform!(st::Statifier, features::T=[]) where {T<:Union{Vector,Matrix,DataFrame}}
  features != [] || return DataFrame()
  typeof(features) <: DataFrame || error("Statifier.fit!: data should be a dataframe: Date,Val ")
  ncol(features) == 2 || error("dataframe must have 2 columns: Date, Val")
  sum(names(features) .== (:Date,:Value))  == 2 || error("wrong column names")
  fstat = fullstat(features[:Value])
  timestat = timevalstat(features)
  if st.args[:processmissing] == true
    # full namedtuple: stat1,stat2,bstat
    hcat(timestat,fstat...)
  else
    hcat(timestat,fstat.stat1,fstat.stat2)
  end
end

function timevalstat(features::DataFrame)
  ldates=features[:Date] |> skipmissing |> collect
  totalhours = sum(diff(ldates)).value/1000/3600 # to hours
  lcount = length(ldates)
  timestart=first(features[:Date])
  timeend=last(features[:Date])
  sfreq = totalhours/lcount
  dftime = DataFrame(tstart=timestart,tend=timeend,sfreq=sfreq)
end

function fullstat(dat::Vector)
  data = skipmissing(dat) |> collect
  lcount = length(data)
  lmax = maximum(data)
  lmin = minimum(data)
  lsm = summarystats(data)
  q1 = quantile(data,0.1)
  q2 = quantile(data,0.2)
  q8 = quantile(data,0.8)
  q9 = quantile(data,0.9)
  lks = kurtosis(data)
  lsk = skewness(data)
  lvar = variation(data)
  lsem = sem(data)
  lentropy = entropy(data)
  # assume 24-hour period
  _autolags = 1:minimum([24,length(data)-1])
  _pacflags = 1:minimum([24,div(length(data),2)-1])
  lautocor = autocor(data,_autolags) .|> abs2 |> sum |> sqrt
  lpacf = pacf(data,_pacflags;method=:yulewalker) .|> abs2 |> sum |> sqrt
  lbrle = rlestatmissingblocks(dat)
  df1=DataFrame(count=lcount,max=lmax,min=lmin,median=lsm.median,
                mean=lsm.mean,q1=q1,q2=q2,q25=lsm.q25,
                q75=lsm.q75,q8=q8,q9=q9)
  df2=DataFrame(kurtosis=lks,skewness=lsk,variation=lvar,entropy=lentropy,
            autocor=lautocor,pacf=lpacf)
  df3=DataFrame(bmedian=lbrle.bmedian,bmean=lbrle.bmean,bq25=lbrle.bq25,bq75=lbrle.bq75,
                bmin=lbrle.bmin,bmax=lbrle.bmax)
  return (stat1=df1,stat2=df2,bstat=df3)
end

function rlestatmissingblocks(dat::Vector{T}) where {T<:Union{AbstractFloat,Integer,Missing}}
  data = deepcopy(dat)
  # replace missing with dummy value for rle processing
  indxmissing=findall(x->ismissing(x),data)
  dummy = -99999
  if eltype(data) <: Union{AbstractFloat,Missing}
    data[indxmissing] .= dummy
  elseif eltype(data) <: Union{Integer,Missing}
    data[indxmissing] .= dummy
  else
    error("not a valid type")
  end
  _rle = rle(data)
  indx=findall(x->x == dummy,_rle[1])
  blocks = _rle[2][indx]
  _sm = summarystats(blocks)
  return (bq25=_sm.q25,bmean=_sm.mean,bmedian=_sm.median,bq75=_sm.q75,bmax=_sm.max,bmin=_sm.min,bmiss=_sm.nmiss)
end

end


function Base.cov{T <: Model}(
    ::Type{T},
    inputs::Vector,
    weights::Vector,
)
    # Construct the negative log likelihood function and its gradient.
    loss = LogLoss()
    λ = 0.0
    nll, nll_gr! = make_closures(T, loss, λ, inputs, weights)

    # Find the MLE for the model.
    res = fit(T, inputs, weights)
    Θ = Optim.minimizer(res)

    # Compute the Hessian at the MLE.
    H = ForwardDiff.hessian(nll, Θ)

    # Construct the covariance matrix.
    return inv(H)
end


include("euler/euler.jl")
using .BigIntegers: BigInteger
using .Numbers: digits_sum
using BenchmarkTools
BenchmarkTools.DEFAULT_PARAMETERS.samples = 100

compute(n::Int)::Int = digits_sum(prod(map(BigInteger, 1:n)))

compute(10)

compute(100)

@benchmark compute(100)