#
#      V(k,n) - The Stiefel-Manifold
#  A point on this manifold is a n×k matrix.
#  A point in an tangential space is a n×k matrix.
#
import LinearAlgebra: norm, dot, nullspace, det, tr, qr, triu, rank, svd, diag, sylvester, lyap, Diagonal
import Base: exp, log, show, cat, rand, Matrix, real, atan
export Stiefel, StPoint, StTVector, getValue
export dot, exp, log, manifoldDimension, norm, parallelTransport, randomTVector, randomMPoint, retractionQR, retractionPolar, inverseRetractionPolar, inverseRetractionQR, projection, retraction, inverseRetraction
export zeroTVector
#
# Type definitions
#

@doc doc"""
    Stiefel{T<:Union{U, Complex{U}} <: Manifold

The manifold $\mathcal M = \mathrm{St}(k,n)$ represented by $n\times k$
orthogonal matrices, that are either real- or complex-valued.

# Abbreviation
  `St`

# Constructor
     Stiefel(k, n[, d=Float64])

generate the manifold $\mathrm{St}(k,n)$ where the integer `n` is the
number of rows and `k` is the number of columns of the matrices and the optional
parameter `d` sets the `DataType` of the matrix entries.
"""
struct Stiefel{T<:Union{U, Complex{U}} where U<:AbstractFloat} <: Manifold
  name::String
  dimensionlines::Int
  dimensioncolumns::Int
  abbreviation::String
  function Stiefel{T}(dimensioncolumns::Int, dimensionlines::Int) where T<:Union{U, Complex{U}} where U<:AbstractFloat
    if dimensioncolumns > dimensionlines
      throw(ErrorException("dimensioncolumns can't be bigger than dimensionlines: $dimensioncolumns > $dimensionlines"))
    else
     new("Stiefel-Manifold St($dimensioncolumns,$dimensionlines) in $T", dimensionlines, dimensioncolumns,"V($dimensioncolumns,$dimensionlines)")
    end
  end
end
Stiefel(dimensionlines::Int, dimensioncolumns::Int, D::DataType=Float64) = Stiefel{D}(dimensionlines::Int, dimensioncolumns::Int)

@doc doc"""
    StPoint <: MPoint

A point $x$ on the manifold $\mathcal M = \mathrm{St}(k,n)$ is represented by an
orthogonal matrix from $\mathbb{K}^{n\times k}.$

# Constructor
    StPoint(Matrix)

where Matrix is an orthogonal matrix of dimension $n×k$.
"""
struct StPoint{T<:Union{U, Complex{U}} where U<:AbstractFloat} <: MPoint
  value::Matrix{T}
  StPoint{T}(value::Matrix{T}) where T<:Union{U, Complex{U}} where U<:AbstractFloat = new(value)
end
getValue(x::StPoint) = x.value;
StPoint(value::Matrix{T}) where T<:Union{U, Complex{U}} where U<:AbstractFloat = StPoint{T}(value)

@doc doc"""
    StTVector <: TVector

A tangent vector $\xi \in T_x\mathcal M$ on the manifold
$\mathcal M = \mathrm{St}(k,n)$. The tangent space is given by as

$T_x\mathrm{St}(k,n) = \bigl\{\xi \in \mathbb{K}^{n\times k} \big| x^{\mathrm{T}}ξ+ξ^{\mathrm{T}}x=0 \bigr\}$.

# Constructor

    StTVector(ξ)

where `ξ` is an $n\times k$ `Matrix` that satisfies the above.
"""
struct StTVector{T<:Union{U, Complex{U}} where U<:AbstractFloat} <: TVector
  value::Matrix{T}
  StTVector{T}(value::Matrix{T}) where T<:Union{U, Complex{U}} where U<:AbstractFloat = new(value)
end
getValue(ξ::StTVector) = ξ.value;
StTVector(value::Matrix{T}) where T<:Union{U, Complex{U}} where U<:AbstractFloat = StTVector{T}(value)

#
# Traits
#
@traitimpl IsMatrixM{Stiefel}
@traitimpl IsMatrixP{StPoint}
@traitimpl IsMatrixTV{StTVector}

#
# Functions
#
@doc doc"""
    dot(M,x,ξ,ν)

compute the Riemannian inner product for two [`StTVector`](@ref)s `ξ` and `ν`
from $T_x\mathcal M$ of the [`Stiefel`](@ref) manifold `M` given by

$\langle \xi, \nu \rangle_x = \operatorname{trace}({\bar \xi}^{\mathrm{T}}\nu).$
"""
function dot(M::Stiefel{T}, x::StPoint{T}, ξ::StTVector{T}, ν::StTVector{T})::Float64 where T<:Union{U, Complex{U}} where U<:AbstractFloat
  return real(tr(getValue(ξ)'*getValue(ν)))
end

@doc doc"""
    exp(M,x,ξ [,t=1.0])

compute the exponential map on the [`Stiefel`](@ref) manifold `M`$=\mathrm{SO}(n)$ with
respect to the [`StPoint`](@ref) `x` and the [`StTVector`](@ref) `ξ`, which can
be shortened with `t` to `tξ`. The formula reads

$\operatorname{exp}_{x} tξ = \begin{pmatrix}
   x\\t\xi
 \end{pmatrix}
 \operatorname{Exp}
 \left(
 \begin{pmatrix} {\bar x}^{\mathrm{T}}\xi & -{\bar \xi}^{\mathrm{T}}\xi\\
 I_{k×k} & {\bar x}^{\mathrm{T}}\xi\end{pmatrix}
 \right)
\begin{pmatrix}  \operatorname{Exp}( -{\bar x}^{\mathrm{T}}\xi) \\ 0_{k×k}\end{pmatrix}$

where $\operatorname{Exp}$ denotes matrix exponential, and $I_{k×k}$ and
$0_{k×k} are the identity matrix and the zero matrix of dimension $k×k$,
respectively.
"""
function exp(M::Stiefel{T},x::StPoint{T},ξ::StTVector{T},t::Float64=1.0) where T<:Union{U, Complex{U}} where U<:AbstractFloat
  Z = zeros(M.dimensioncolumns,M.dimensioncolumns)
  I = one(Z)
  Ξ = t*getValue(ξ)
  Y = [getValue(x) Ξ] * exp([getValue(x)'*Ξ -Ξ'*Ξ; I getValue(x)'*Ξ]) * [exp(-getValue(x)'*Ξ); Z]
  StPoint{T}(Y)
end

@doc doc"""
    inverseRetractionPolar(M,x,y)

return a [`StTVector`](@ref) `ξ` of the tagent space $T_x\mathrm{SO}(n)$
of the [`StPoint`](@ref) `x` on the [`Stiefel`](@ref) manifold `M`
with which the [`StPoint`](@ref) `y` can be reached by the
[`retractionPolar`](@ref)  after time 1.
The formula reads

$ξ = ys-x$

where $s$ is the solution to the Lyapunov equation

$x^{\mathrm{T}}ys + s(x^{\mathrm{T}}y)^{\mathrm{T}} + 2\mathrm{I}_k = 0.$

This function is implemented only for the case $\mathbb{K}=\mathbb{R}$.
"""
function inverseRetractionPolar(M::Stiefel{T},x::StPoint{T},y::StPoint{T}) where T<:AbstractFloat
  A = transpose(getValue(x)) * getValue(y)
  H = -2 * one(transpose(getValue(x))*getValue(x))
  B = lyap(A, H)
  C = getValue(y)*B - getValue(x)
  StTVector{T}( C )
end

@doc doc"""
    inverseRetractionQR(M,x,y)

return a [`StTVector`](@ref) `ξ` of the tagent space $T_x\mathrm{SO}(n)$
of the [`StPoint`](@ref) `x` on the [`Stiefel`](@ref) manifold `M`
with which the [`StPoint`](@ref) `y` can be reached by the
[`retractionQR`](@ref) after time 1.
This function is implemented only for the case $\mathbb{K}=\mathbb{R}$.
This is also the standard retraction.
"""
function inverseRetractionQR(M::Stiefel{T},x::StPoint{T},y::StPoint{T}) where T<:AbstractFloat
  A = transpose(getValue(x)) * getValue(y)
  R = zeros(M.dimensioncolumns, M.dimensioncolumns)
  for i = 1:M.dimensioncolumns
    b = zeros(i)
    b[end] = 1
    b[1:(end-1)] = - transpose(R[1:(i-1), 1:(i-1)]) * A[i, 1:(i-1)]
    R[1:i, i] = A[1:i, 1:i] \ b
  end
  C =  getValue(y)*R-getValue(x)
  StTVector{T}( C )
end

function inverseRetraction(M::Stiefel{T},x::StPoint{T},y::StPoint{T}) where T<:AbstractFloat
  return inverseRetractionQR(M,x,y)
end

@doc doc"""
    manifoldDimension(x)

return the dimension of the [`Stiefel`](@ref) manifold `M`$= \mathrm{St}(k,n)$, the
[`StPoint`](@ref) `x`, itself embedded in $\mathbb R^{n\times k}$, belongs to.
The dimension for $\mathbb{K}=\mathbb{R}$ is given by

$nk - \frac{1}{2}k(k+1)$

and for $\mathbb{K}=\mathbb{C}$

$2nk - k^2.$
"""
function manifoldDimension(x::StPoint{T}) where T<:AbstractFloat
  return Int(size(getValue(x),1) * size(getValue(x),2) - 0.5 * size(getValue(x),2) * (size(getValue(x),2) + 1))
end

function manifoldDimension(x::StPoint{T}) where T<:Complex{U} where U<:AbstractFloat
  return Int(2 * size(getValue(x),1) * size(getValue(x),2) - size(getValue(x),2)^2)
end

@doc doc"""
    manifoldDimension(M)

return the dimension of the [`Stiefel`](@ref) manifold `M`.
The dimension for $\mathbb{K}=\mathbb{R}$ is given by

$nk - \frac{1}{2}k(k+1)$

and for $\mathbb{K}=\mathbb{C}$

$2nk - k^2.$
"""
function manifoldDimension(M::Stiefel{T}) where T<:AbstractFloat
  return Int(M.dimensionlines * M.dimensioncolumns - 0.5 * M.dimensioncolumns * (M.dimensioncolumns + 1))
end
function manifoldDimension(M::Stiefel{T}) where T<:Complex{U} where U<:AbstractFloat
  return Int(2 * M.dimensionlines * M.dimensioncolumns - (M.dimensioncolumns)^2)
end

@doc doc"""
    norm(M,x,ξ)

compute the norm of the [`StTVector`](@ref) `ξ` in the tangent space
$T_x\mathcal M$ at [`StPoint`](@ref) `x` of the [`Stiefel`](@ref) manifold `M`.

$\lVert \xi \rVert_x = \sqrt{\sum_{i,j=0}^n \xi_{ij}^2}$

where $\xi_{ij}$ are the entries of the matrix `ξ`, i.e. the norm is given by
the Frobenius norm of `ξ`.
"""
function norm(M::Stiefel{T}, x::StPoint{T}, ξ::StTVector{T})::Float64 where T<:Union{U, Complex{U}} where U<:AbstractFloat
  norm(getValue(ξ))
end

@doc doc"""
    parallelTransport(M,x,y,ξ)

compute the paralllel transport of the [`StTVector`](@ref) `ξ` from
the tangent space $T_x\mathcal M$ at [`StPoint`](@ref) `x` to
$T_y\mathcal M$ at [`StPoint`](@ref) `y` on the [`Stiefel`](@ref) manifold `M`
provided that the corresponding [`geodesic`](@ref) $g(\cdot;x,y)$ is unique.
The formula reads

$P_{x\to y}(\xi) = \operatorname{proj}_{\mathcal M}(y,\xi).$

where $\operatorname{proj}_{\mathcal M}$ is the projection onto the
tangent space $T_y\mathcal M$.
"""
function parallelTransport(M::Stiefel{T}, x::StPoint{T}, y::StPoint{T}, ξ::StTVector{T}) where T<:Union{U, Complex{U}} where U<:AbstractFloat
    projection(M,y,getValue(ξ))
end

@doc doc"""
    projection(M,x,q)

project a `Matrix q` orthogonally on the tangent space of the
[`StPoint`](@ref) `x` on the [`Stiefel`](@ref) manifold `M`. The formula reads

$\operatorname{proj}_{\mathcal M}(x,q) = q-xB$

where

$B=\frac{1}{2} (x^{\mathrm{T}}{\bar q})^{\mathrm{T}} {\bar x}^{\mathrm{T}}q.$

# see also
[`parallelTransport`](@ref), [`randomTVector`](@ref)
"""
function projection(M::Stiefel{T}, x::StPoint{T}, q::Matrix{T}) where T<:Union{U, Complex{U}} where U<:AbstractFloat
    A = getValue(x)'*q
    B = 0.5 * (A' + A)
    return StTVector{T}(q - getValue(x) * B)
end

@doc doc"""
    randomMPoint(M [,:Gaussian, σ=1.0])

return a random (Gaussian) [`StPoint`](@ref) `x` on the manifold
[`Stiefel`](@ref) manifold `M` by generating a (Gaussian) matrix with standard deviation
`σ` and return the orthogonalized version, i.e. return the Q
component of the QR decomposition of the random matrix of size $n×k$.
"""
function randomMPoint(M::Stiefel{T}, ::Val{:Gaussian}, σ::Float64=1.0) where T<:Union{U, Complex{U}} where U<:AbstractFloat
    A = σ*randn(T, (M.dimensionlines, M.dimensioncolumns))
    return StPoint{T}( Matrix(qr(A).Q) )
end

@doc doc"""
    randomTVector(M,x [,:Gaussian, σ=1.0])

return a random vector [`StTVector`](@ref) in the tangential space
$T_x\mathrm{St}(k,n)$ by generating a random matrix of size $n×k$ and projecting
it onto [`StPoint`](@ref) `x` with [`projection`](@ref).
"""
function randomTVector(M::Stiefel{T}, x::StPoint{T}, ::Val{:Gaussian}, σ::Float64=1.0) where T<:Union{U, Complex{U}} where U<:AbstractFloat
  y = σ * randn(T, (M.dimensionlines,M.dimensioncolumns))
  Y = projection(M, x, y)
  return 1/(norm(getValue(Y))) * Y
end

@doc doc"""
    retractionPolar(M,x,ξ,[t=1.0])

move the [`StPoint`](@ref) `x` in the direction of the [`StTVector`](@ref) `ξ`
on the  [`Stiefel`](@ref) manifold `M`. This SVD-based retraction is a second-order
approximation of the exponential map [`exp`](@ref). Let

$USV = x + tξ$

be the singular value decomposition, then the formula reads

$\operatorname{retr}_x\xi = UV^\mathrm{T}.$

This function is implemented only for the case $\mathbb{K}=\mathbb{R}$.
"""
function retractionPolar(M::Stiefel{T},x::StPoint{T},ξ::StTVector{T},t::Float64=1.0) where T<:AbstractFloat
  y = getValue(x) + t * getValue(ξ)
  S = svd(y)
  A = S.U * transpose(S.V)
  StPoint{T}(A)
end

@doc doc"""
    retractionQR(M,x,ξ,[t=1.0])

move the [`StPoint`](@ref) `x` in the direction of the [`StTVector`](@ref) `ξ`
on the  [`Stiefel`](@ref) manifold `M`. This QR-based retraction is a
first-order approximation of the exponential map [`exp`](@ref). Let

$QR = x + tξ$

be the QR decomposition, then the formula reads

$\operatorname{retr}_x\xi = QD$

where D is a $n×k$ matrix with the signs of the diagonal entries of $R$ plus
$0.5$ on the upper diagonal, i.e.

$D_{ij}=\begin{cases} \operatorname{sgn}(R_{ij}+0,5) & \text{if} \; i=j \\ 0 & \, \text{otherwise} \end{cases}.$

This is also the standard retraction.
"""
function retractionQR(M::Stiefel{T},x::StPoint{T},ξ::StTVector{T},t::Float64=1.0) where T<:Union{U, Complex{U}} where U<:AbstractFloat
  y = getValue(x) + t * getValue(ξ)
  QRdecomp = qr(y)
  d = diag(QRdecomp.R)
  D = Diagonal( sign.( sign.(d .+ 0.5) ) )
  B = zeros(M.dimensionlines,M.dimensioncolumns)
  B[1:M.dimensioncolumns,1:M.dimensioncolumns] = D
  A = QRdecomp.Q * B
  StPoint{T}(A)
end

function retraction(M::Stiefel{T},x::StPoint{T},ξ::StTVector{T},t::Float64=1.0) where T<:Union{U, Complex{U}} where U<:AbstractFloat
  return retractionQR(M,x,ξ,t)
end

@doc doc"""
    ξ = zeroTVector(M,x)

return a zero vector in the tangent space $T_x\mathcal M$ of the
[`StPoint`](@ref) $x\in\mathrm{St}(k,n)$ on the [`Stiefel`](@ref) manifold `M`.
"""
function zeroTVector(M::Stiefel{T},x::StPoint{T}) where T<:Union{U, Complex{U}} where U<:AbstractFloat
  return StTVector{T}( zero(getValue(x)) )
end

@doc doc"""
    validateMPoint(M,x)

validate that the [`StPoint`](@ref) `x` is a valid point on the
[`Stiefel`](@ref) manifold `M`, i.e. that both dimensions and the rank are correct
as well as that all columns are orthonormal.
"""
function validateMPoint(M::Stiefel{T}, x::StPoint{T}) where T<:Union{U, Complex{U}} where U<:AbstractFloat
  if size(getValue(x), 1) ≠ M.dimensionlines
    throw( ErrorException("The dimension of $x must be $(M.dimensionlines) × $(M.dimensioncolumns) but it is $(size(getValue(x), 1)) × $(size(getValue(x), 2)).") )
  end
  if size(getValue(x), 2) ≠ M.dimensioncolumns
    throw( ErrorException("The dimension of $x must be $(M.dimensionlines) × $(M.dimensioncolumns) but it is $(size(getValue(x), 1)) × $(size(getValue(x), 2)).") )
  end
if rank(getValue(x)) ≠ M.dimensioncolumns #is this necessary?
    throw( ErrorException("$x must have  full rank.") )
  end

  z=getValue(x)'*getValue(x)
  if z ≉ one(z)
    throw( ErrorException("$x has to be orthonormal but it's not.") )
  end
end
@doc doc"""
    validateMPoint(M,x,ξ)

validate that the [`StTVector`](@ref) `ξ` is a valid tangent vector to
[`StPoint`](@ref) `x` on the [`Stiefel`](@ref) manifold `M`, i.e. that both
dimensions agree and ${\bar x}^{\mathrm{T}}$ is skew symmetric.
"""
function validateTVector(M::Stiefel{T}, x::StPoint{T}, ξ::StTVector{T}) where T<:Union{U, Complex{U}} where U<:AbstractFloat
  if size(getValue(ξ), 1) ≠ M.dimensionlines
    throw( ErrorException("The dimension of $ξ must be $(M.dimensionlines) × $(M.dimensioncolumns) but it is $(size(getValue(ξ), 1)) × $(size(getValue(ξ), 2))") )
  end
  if size(getValue(ξ), 2) ≠ M.dimensioncolumns
    throw( ErrorException("The dimension of $ξ must be $(M.dimensionlines) × $(M.dimensioncolumns) but it is $(size(getValue(ξ), 1)) × $(size(getValue(ξ), 2))") )
  end
  if norm(getValue(x)'*getValue(ξ) + getValue(ξ)'*getValue(x)) > 10^(-15)
    throw( ErrorException("The matrix $x'$ξ must be skew-symmetric!") )
  end
end


"""
    simulate_random_snparray(n::Integer, p::Integer, s::String; mafs::Vector{Float64}, min_ma::Integer)

Creates a random SnpArray in the current directory without missing value,
where each SNP has ⫺5 (default) minor alleles.

Note: if supplied minor allele frequency is extremely small, it could take a long time for
the simulation to generate samples where at least `min_ma` (defaults to 5) are present.

# Arguments:
- `n`: number of samples
- `p`: number of SNPs
- `s`: name of SnpArray that will be created (memory mapped) in the current directory. To not memory map, use `undef`.

# Optional Arguments:
- `mafs`: vector of desired minor allele freuqencies (uniform(0, 0.5) by default)
- `min_ma`: the minimum number of minor alleles that must be present for each SNP (defaults to 5)
"""
function simulate_random_snparray(n::Int64, p::Int64, s::Union{String, UndefInitializer};
                                  mafs::Vector{Float64}=zeros(Float64, p), min_ma::Int = 5)

    @assert all(0.0 .<= mafs .<= 0.5) "Minor allele frequencies must be all in the range (0, 0.5)"

    if mafs != zeros(Float64, p)
        return _random_snparray(n, p, s, mafs, min_ma=min_ma)
    end

    #first simulate a random {0, 1, 2} matrix with each SNP drawn from Binomial(2, r[i])
    A1 = BitArray(undef, n, p)
    A2 = BitArray(undef, n, p)
    for j in 1:p
        minor_alleles = 0
        maf = 0
        while minor_alleles <= min_ma
            maf = 0.5rand()
            for i in 1:n
                A1[i, j] = rand(Bernoulli(maf))
                A2[i, j] = rand(Bernoulli(maf))
            end
            minor_alleles = sum(view(A1, :, j)) + sum(view(A2, :, j))
        end
        mafs[j] = maf
    end

    #fill the SnpArray with the corresponding x_tmp entry
    return _make_snparray(A1, A2, s)
end

"""
This function requires a vector of minor alleles frequencies to perform simulation.
It will create a random SnpArray in the current directory without missing value,
where each SNP has at least 5 minor alleles.
"""
function _random_snparray(n::Int64, p::Int64, s::Union{String, UndefInitializer},
                          mafs::Vector{Float64}; min_ma::Int = 5)
    all(0.0 .<= mafs .<= 0.5) || throw(ArgumentError("vector of minor allele frequencies must be in (0, 0.5)"))
    any(mafs .<= 0.0005) && @warn("Provided minor allele frequencies contain entries smaller than 0.0005, simulation may take long if sample size is small and min_ma = $min_ma is large")

    #first simulate a random {0, 1, 2} matrix with each SNP drawn from Binomial(2, r[i])
    A1 = BitArray(undef, n, p)
    A2 = BitArray(undef, n, p)
    for j in 1:p
        minor_alleles = 0
        maf = mafs[j]
        while minor_alleles <= min_ma
            for i in 1:n
                A1[i, j] = rand(Bernoulli(maf))
                A2[i, j] = rand(Bernoulli(maf))
            end
            minor_alleles = sum(view(A1, :, j)) + sum(view(A2, :, j))
        end
    end

    #fill the SnpArray with the corresponding x_tmp entry
    return _make_snparray(A1, A2, s)
end

"""
Make a SnpArray from 2 BitArrays.
"""
function _make_snparray(A1::BitArray, A2::BitArray, s::Union{String, UndefInitializer})
    n, p = size(A1)
    x = SnpArray(s, n, p)
    for i in 1:(n*p)
        c = A1[i] + A2[i]
        if c == 0
            x[i] = 0x00
        elseif c == 1
            x[i] = 0x02
        elseif c == 2
            x[i] = 0x03
        else
            throw(error("matrix shouldn't have missing values!"))
        end
    end
    return x
end

"""
    simulate_correlated_snparray(n, p, s; block_length, hap, prob)

Simulates a SnpArray with correlation. SNPs are divided into blocks where each
adjacent SNP is the same with probability prob. There are no correlation between blocks.

# Arguments:
- `n`: number of samples
- `p`: number of SNPs
- `s`: name of SnpArray that will be created (memory mapped) in the current directory. To not memory map, use `undef`.

# Optional arguments:
- `block_length`: length of each LD block
- `hap`: number of haplotypes to simulate for each block
- `prob`: with probability `prob` an adjacent SNP would be the same.
"""
function simulate_correlated_snparray(n::Int64, p::Int64, s::Union{String, UndefInitializer};
            block_length::Int64=20, hap::Int=20, prob::Float64=0.75)

    @assert mod(p, block_length) == 0 "block_length ($block_length) is not divible by p ($p)"
    @assert 0 < prob < 1 "transition probably should be between 0 and 1, got $prob"

    x = SnpArray(s, n, p)
    haplotypes = zeros(hap, block_length)
    snps = zeros(block_length)
    blocks = Int(p / block_length)

    @inbounds for b in 1:blocks

        #create pool of haplotypes for each block
        _sample_haptotypes!(haplotypes, prob)

        for i in 1:n
            #sample 2 haplotypes with replacement from the pool of haplotypes
            row1 = rand(1:hap)
            row2 = rand(1:hap)
            for j in 1:block_length
                snps[j] = haplotypes[row1, j] + haplotypes[row2, j]
            end

            #copy haplotypes into x
            _copy_blocks!(x, i, snps, b, block_length)
        end
    end

    return x
end

function _sample_haptotypes!(haplotypes::Matrix, prob::Float64)
    n, p = size(haplotypes)
    fill!(haplotypes, 0)

    @inbounds for i in 1:n
        cur_row_sum = 0
        while cur_row_sum == 0
            curr = rand(0:1)
            haplotypes[i, 1] = curr
            cur_row_sum += curr
            for j in 2:p
                stay = rand(Bernoulli(prob)) #stay = 1 means retain the current value
                curr = (stay == 1 ? curr : 1 - curr)
                haplotypes[i, j] = curr
                cur_row_sum += curr
            end
        end
    end
end

function _copy_blocks!(x::SnpArray, row, snps, cur_block, block_length)
    #copy sampled snps into SnpArray
    @inbounds for k in 1:length(snps)
        c = snps[k]
        col = (cur_block - 1) * block_length + k
        if c == 0
            x[row, col] = 0x00
        elseif c == 1
            x[row, col] = 0x02
        elseif c == 2
            x[row, col] = 0x03
        else
            throw(error("SNP values should be 0, 1, or 2 but was $c"))
        end
    end
end

"""
    simulate_random_response(x::SnpArray, xbm::SnpBitMatrix, k::Int, d::UnionAll, l::Link)

This function simulates a random response (trait) vector `y` based on provided x, β, distirbution,
and link function. When the distribution is from Poisson, Gamma, or Negative Binomial, we simulate `β ∼ N(0, 0.3)`
to roughly ensure the mean of response `y` doesn't become too large. For other distributions,
we choose `β ∼ N(0, 1)`.

# Arguments
- `x`: The SnpArray
- `xbm`: SnpBitMatrix type of your SnpArray
- `k`: the true number of predictors.
- `d`: The distribution of the simulated trait (note `typeof(d) = UnionAll` but `typeof(d())` is an actual distribution: e.g. Normal)
- `l`: The link function. Input `canonicallink(d())` if you want to use the canonical link of `d`.

# Optional arguments
- `nn`: The number of success until stopping in negative binomial regression, defaults to 10
- `α`: Shape parameter of the gamma distribution, defaults to 1
"""
function simulate_random_response(x::SnpArray, xbm::SnpBitMatrix, k::Int,
                                  d::UnionAll, l::Link; nn = 10,
                                  α = 1) where {T <: Float}
    n, p = size(x)
    if (typeof(d) <: NegativeBinomial) || (typeof(d) <: Gamma)
        l == LogLink() || throw(ArgumentError("Distribution $d must use LogLink!"))
    end

    #simulate a random model β from a normal distribution
    true_b = zeros(p)
    if d == Poisson || d == Gamma || d == NegativeBinomial
        true_b[1:k] = rand(Normal(0, 0.3), k)
    else
        true_b[1:k] = randn(k)
    end
    shuffle!(true_b)
    correct_position = findall(x -> x != 0, true_b)

    #simulate phenotypes (e.g. vector y)
    if d == Normal || d == Poisson || d == Bernoulli
        prob = linkinv.(l, xbm * true_b)
        clamp!(prob, -20, 20)
        y = [rand(d(i)) for i in prob]
    elseif d == NegativeBinomial
        μ = linkinv.(l, xbm * true_b)
        clamp!(μ, -20, 20)
        prob = 1 ./ (1 .+ μ ./ nn)
        y = [rand(d(nn, i)) for i in prob] #number of failtures before nn success occurs
    elseif d == Gamma
        μ = linkinv.(l, xbm * true_b)
        β = 1 ./ μ # here β is the rate parameter for gamma distribution
        y = [rand(d(α, i)) for i in β] # α is the shape parameter for gamma
    end
    y = Float64.(y)

    return y, true_b, correct_position
end

"""
    adhoc_add_correlation!(x::SnpArray, ρ::Float64, pos::Int64, location::Vector{Int})

Makes 1 SNP (a column of `x`) correlate with SNPs in `location` with correlation coefficient roughly `ρ`.

# Arguments:
- `x`: the snparray
- `ρ`: correlation coefficient
- `pos`: the position of the target SNP that everything would be correlated to
- `location`: All the SNPs that shall be correlated with the SNP at position `pos` with correlation `ρ`.
"""
function adhoc_add_correlation!(x::SnpArray, ρ::Float64, pos::Int64, location::Vector{Int})
    @assert 0 <= ρ <= 1 "correlation coefficient must be in (0, 1) but was $ρ"

    for loc in location
        for i in 1:size(x, 1)
            prob = rand(Bernoulli(ρ))
            prob == 1 && (x[i, loc] = x[i, pos]) #make 2nd column the same as 1st ~ρ% of the time
        end
        corr = cor(x[:, loc], x[:, pos])
    end
end

"""
    simulate_rare_variants(mafs::Vector{Float64}, penetrance::Vector{Float64}, cases::Int, controls::Int)

TODO: implement Algrithm 1 of https://www.ncbi.nlm.nih.gov/pmc/articles/PMC3025646/
"""
function simulate_rare_variants(mafs::Vector{Float64}, penetrance::Vector{Float64},
                                cases::Int, controls::Int)
    return nothing # not implemented yet
end

"""
    make_bim_fam_files(x::SnpArray, y, name::String)

Creates .bim and .bed files from a SnpArray.

# Arguments:
- `x`: A SnpArray (i.e. `.bed` file on the disk) for which you wish to create corresponding `.bim` and `.fam` files.
- `name`: string that should match the `.bed` file (Do not include `.bim` or `.fam` extensions in `name`).
- `y`: Trait vector that will go in to the 6th column of `.fam` file.
"""
function make_bim_fam_files(x::SnpArray, y, name::String)
    ly = length(y)
    n, p = size(x)
    @assert n == ly "dimension mismatch: phenotype data has length $ly but SnpArray has $n samples"

    #create .bim file structure: https://www.cog-genomics.org/plink2/formats#bim
    open(name * ".bim", "w") do f
        for i in 1:p
            write(f, "1 \t $i \t 0 \t 1 \t 1 \t 2 \n")
        end
    end

    #create .fam file structure: https://www.cog-genomics.org/plink2/formats#bim
    open(name * ".fam", "w") do f
        for i in 1:n
            yi = y[i]
            write(f, "$i \t 1 \t 0 \t 0 \t 1 \t $yi \n")
        end
    end
end


@testset "fc" begin
    config = [2, 20,20,20,3]
    x = constant(rand(10,2))
    θ = ae_init(config)
    u, du = fc(x, config[2:end], θ)
    y = ae(x, config[2:end], θ)

    @test run(sess, u) ≈ run(sess, y)

    DU = run(sess, du)
    @test run(sess, tf.gradients(y[:,1], x)[1]) ≈ DU[:,1,:]
    @test run(sess, tf.gradients(y[:,2], x)[1]) ≈ DU[:,2,:]
    @test run(sess, tf.gradients(y[:,3], x)[1]) ≈ DU[:,3,:]
end


asparams(::Affine, ::Val{:μ}) = asℝ
asparams(::Affine, ::Val{:σ}) = asℝ₊
asparams(::Type{A}, ::Val{:μ}) where {A<:Affine} = asℝ
asparams(::Type{A}, ::Val{:σ}) where {A<:Affine} = asℝ₊

asparams(::Affine{N,M,T}, ::Val{:μ}) where {N,M,T<:AbstractArray} = as(Array, asℝ, size(d.μ))
asparams(::Affine{N,M,T}, ::Val{:σ}) where {N,M,T<:AbstractArray} = as(Array, asℝ, size(d.σ))


function givenscheck(T,n,tol)
    e = eps(real(T))
    A = rand(T,n,n)
    Ab = big.(A)
    G,_ = givens(A[4,5],A[5,5],4,5)
    Gb,_ = givens(Ab[4,5],Ab[5,5],4,5)
    B = copy(A)
    Bb = copy(Ab)
    lmul!(G,B)
    lmul!(Gb,Bb)
    @test norm(B-Bb) < tol * n * e
    rmul!(A,G)
    rmul!(Ab,Gb)
    @test norm(A-Ab) < tol * n * e
end

@testset "givens $T" for T in (Double64, Double32, Complex{Double64})
    n = 67
    givenscheck(T,n,4)
end


function is_equivalent_by_translation(cryst::Crystal, b1::BondRaw, b2::BondRaw)
    # Displacements between the two bonds
    D1 = b2.ri - b1.ri
    D2 = b2.rj - b1.rj
    # Round components of D1 to nearest integers
    n = round.(D1, RoundNearest)
    # If both n ≈ D1 and n ≈ D2, then the bonds are equivalent by translation
    return norm(n - D1) < cryst.symprec && norm(n - D2) < cryst.symprec
end

# Generate list of all symmetries that transform b2 into b1, along with parity
function symmetries_between_bonds(cryst::Crystal, b1::BondRaw, b2::BondRaw)
    # Fail early if two bonds describe different real-space distances
    # (dimensionless error tolerance is measured relative to the minimum lattice
    # constant ℓ)
    if b1 != b2
        ℓ = minimum(norm, eachcol(cryst.lat_vecs))
        d1 = distance(cryst, b1) / ℓ
        d2 = distance(cryst, b2) / ℓ
        if abs(d1-d2) > cryst.symprec
            return Tuple{SymOp, Bool}[]
        end
    end

    ret = Tuple{SymOp, Bool}[]
    for s in cryst.symops
        b2′ = transform(s, b2)
        if is_equivalent_by_translation(cryst, b1, b2′)
            push!(ret, (s, true))
        elseif is_equivalent_by_translation(cryst, b1, reverse(b2′))
            push!(ret, (s, false))
        end
    end
    return ret
end

"Is there a symmetry operation that transforms `b1` into either `b2` or its reverse?"
function is_related_by_symmetry(cryst::Crystal, b1::BondRaw, b2::BondRaw)
    return !isempty(symmetries_between_bonds(cryst::Crystal, b1::BondRaw, b2::BondRaw))
end

function is_related_by_symmetry(cryst::Crystal, b1::Bond, b2::Bond)
    return is_related_by_symmetry(cryst, BondRaw(cryst, b1), BondRaw(cryst, b2))
end

"Returns all bonds in `cryst` for which `bond.i == i`"
function all_bonds_for_atom(cryst::Crystal, i::Int, max_dist; min_dist=0.0)
    # be a little generous with the minimum and maximum distances
    ℓ = minimum(norm, eachcol(cryst.lat_vecs))
    max_dist += 4 * cryst.symprec * ℓ
    min_dist -= 4 * cryst.symprec * ℓ

    # columns are the reciprocal vectors
    recip_vecs = 2π * inv(cryst.lat_vecs)'

    # box_lengths[i] represents the perpendicular distance between two parallel
    # boundary planes spanned by lattice vectors a_j and a_k (where indices j
    # and k differ from i)
    box_lengths = [a⋅b/norm(b) for (a,b) = zip(eachcol(cryst.lat_vecs), eachcol(recip_vecs))]
    n_max = round.(Int, max_dist ./ box_lengths, RoundUp)

    bonds = Bond[]

    # loop over neighboring cells
    for n1 in -n_max[1]:n_max[1]
        for n2 in -n_max[2]:n_max[2]
            for n3 in -n_max[3]:n_max[3]
                n = SVector(n1, n2, n3)

                # loop over all atoms within neighboring cell
                for j in eachindex(cryst.positions)
                    b = Bond(i, j, n)
                    if min_dist <= distance(cryst, b) <= max_dist
                        push!(bonds, b)
                    end
                end
            end
        end
    end

    return bonds
end

# Calculate score for a bond. Lower would be preferred.
function _score_bond(cryst::Crystal, b::Bond)
    # Favor bonds with fewer nonzero elements in basis matrices J
    Js = basis_for_symmetry_allowed_couplings(cryst, b)
    nnz = [count(abs.(J) .> 1e-12) for J in Js]
    score = Float64(sum(nnz))

    # Favor bonds with smaller unit cell displacements. Positive
    # displacements are slightly favored over negative displacements.
    # Displacements in x are slightly favored over y, etc.
    score += norm((b.n .- 0.1) .* [0.07, 0.08, 0.09])

    # Favor smaller indices and indices where i < j
    score += 1e-2 * (b.i + b.j) + 1e-2 * (b.i < b.j ? -1 : +1)

    return score
end

"Indices of the unique elements in `a`, ordered by their first appearance."
function unique_indices(a)
    map(x->x[1], unique(x->x[2], enumerate(a)))
end

"""    reference_bonds(cryst::Crystal, max_dist)

Returns a full list of bonds, one for each symmetry equivalence class, up to
distance `max_dist`. The reference bond `b` for each equivalence class is
selected according to a scoring system that prioritizes simplification of the
elements in `basis_for_symmetry_allowed_couplings(cryst, b)`."""
reference_bonds(cryst::Crystal, max_dist) = reference_bonds(cryst, convert(Float64, max_dist))

function reference_bonds(cryst::Crystal, max_dist::Float64)
    # Bonds, one for each equivalence class
    ref_bonds = Bond[]
    for i in unique_indices(cryst.classes)
        for b in all_bonds_for_atom(cryst, i, max_dist)
            if !any(is_related_by_symmetry(cryst, b, b′) for b′ in ref_bonds)
                push!(ref_bonds, b)
            end
        end
    end

    # Sort by distance
    sort!(ref_bonds, by=b->distance(cryst, b))

    # Replace each canonical bond by the "best" equivalent bond
    return map(ref_bonds) do rb
        # Find full set of symmetry equivalent bonds
        equiv_bonds = unique([transform(cryst, s, rb) for s in cryst.symops])
        # Take the bond with lowest score
        scores = [_score_bond(cryst, b) for b in equiv_bonds]
        return equiv_bonds[findmin(scores)[2]::Int]
        # TODO: In Julia 1.7 the above two lines become
        #     return argmin(b -> _score_bond(cryst, b), equiv_bonds)
    end
end

"""
    all_symmetry_related_bonds_for_atom(cryst::Crystal, i::Int, b::Bond)

Returns a list of all bonds that start at atom `i`, and that are symmetry
equivalent to bond `b` or its reverse.
"""
function all_symmetry_related_bonds_for_atom(cryst::Crystal, i::Int, b_ref::Bond)
    bs = Bond[]
    dist = distance(cryst, b_ref)
    for b in all_bonds_for_atom(cryst, i, dist; min_dist=dist)
        if is_related_by_symmetry(cryst, b_ref, b)
            push!(bs, b)
        end
    end
    return bs
end

"""
    all_symmetry_related_bonds(cryst::Crystal, b::Bond)

Returns a list of all bonds that are symmetry-equivalent to bond `b` or its
reverse.
"""
function all_symmetry_related_bonds(cryst::Crystal, b_ref::Bond)
    bs = Bond[]
    for i in eachindex(cryst.positions)
        append!(bs, all_symmetry_related_bonds_for_atom(cryst, i, b_ref))
    end
    bs
end

"""    coordination_number(cryst::Crystal, i::Int, b::Bond)

Returns the number times that atom `i` participates in a bond equivalent to `b`.
In other words, the count of bonds that begin at atom `i` and that are
symmetry-equivalent to `b` or its reverse.

Defined as `length(all_symmetry_related_bonds_for_atom(cryst, i, b))`.
"""
function coordination_number(cryst::Crystal, i::Int, b::Bond)
    return length(all_symmetry_related_bonds_for_atom(cryst, i, b))
end


module MLE

__precompile__(true)

using Distributions
using Optim
using LinearAlgebra

const EULER_MASCHERONI = big"0.57721566490153286060651209008240243104215933593992"

const REAL_SUPPORT = [-Inf, Inf]
const NON_NEGATIVE_SUPPORT = [0, Inf]
const UNIT_SUPPORT = Float64[0, 1]

"""
    parameters(D)

Extract the names of parameters from the distribution D to be estimated.

# Examples:

```jldoctest
julia> parameters(Normal)
(:μ, :σ)
```

"""
parameters(::Type{D}) where {D <: UnivariateDistribution} = fieldnames(D)

"""
    mle_starting_point(D, data)

Compute a starting point for minimizing log likelihood of the distribution D
for the dataset.

# Examples:

```jldoctest
julia> mle_starting_point(Normal, [1,2,3])
(μ = 2.0, σ = 1.0)
```

"""
function mle_starting_point(::Type{D}, data) where {D<:UnivariateDistribution}
    # ones(Float64, fieldcount(D))
    params = parameters(D)
    (; zip(parameters(D), ones(Float64, shape(params)))...)
end

function mle_starting_point(::Type{Gumbel}, data)
    μ = mean(data)
    θ = var(data)
    (;μ, θ)
end

parameter_support(::Type{Gumbel}) = (μ=REAL_SUPPORT, θ=NON_NEGATIVE_SUPPORT)

parameter_support(::Type{Weibull}) = (α=NON_NEGATIVE_SUPPORT, θ=NON_NEGATIVE_SUPPORT)
mle_starting_point(::Type{Weibull}, data) = (;α = mean(data), θ = var(data))

mle_starting_point(::Type{Normal}, data) = (; μ = mean(data), σ = var(data))
parameter_support(::Type{Normal}) = (μ=REAL_SUPPORT, σ=NON_NEGATIVE_SUPPORT)

mle_starting_point(::Type{Poisson}, data) = (; λ = mean(data))
parameter_support(::Type{Poisson}) = (; λ=NON_NEGATIVE_SUPPORT)

mle_starting_point(::Type{NegativeBinomial}, data) = (; p = mean(data) / maximum(data), r=maximum(data))
parameter_support(::Type{NegativeBinomial}) = (; p=UNIT_SUPPORT, r=NON_NEGATIVE_SUPPORT)

mle_starting_point(::Type{Gamma}, data) = (α=mean(data), θ=var(data))
parameter_support(::Type{Gamma}) = (α=NON_NEGATIVE_SUPPORT, θ=NON_NEGATIVE_SUPPORT)

mle_starting_point(::Type{Binomial}, data) = (;p=mean(data) / maximum(data), n=maximum(data))
parameter_support(::Type{Binomial}) = (;p=UNIT_SUPPORT, n=NON_NEGATIVE_SUPPORT)
Binomial(n::Float64, p::Float64; check_args) = Binomial(Int(n), p; check_args)

mle_starting_point(::Type{Geometric}, data) = (; p=1.0/mean(data))
parameter_support(::Type{Geometric}) = (; p=UNIT_SUPPORT)

mle_starting_point(::Type{Beta}, data) = (α=1.0, β=1.0)
parameter_support(::Type{Beta}) = (α=NON_NEGATIVE_SUPPORT, β=NON_NEGATIVE_SUPPORT)

mle_starting_point(::Type{Exponential}, data) = (;θ = mean(data))
parameter_support(::Type{Exponential}) = (;θ=NON_NEGATIVE_SUPPORT)

struct MLEConfig{D}
    distribution::D
    data::AbstractArray
    params::Tuple
    params_to_solve::Vector{Bool}
    starting_point::Vector
    constraints::Optim.AbstractConstraints
    custom_bounds::Bool
end

function make_constraints(d::Type{<:UnivariateDistribution}, start, holding, bounds)
    params = parameters(d)

    lower = [ bounds[p][1] for p in params if p ∈ keys(start) && p ∉ keys(holding) ]
    upper = [ bounds[p][2] for p in params if p ∈ keys(start) && p ∉ keys(holding) ]

    TwiceDifferentiableConstraints(lower, upper)
end

function _get_bounds(d::Type{<:UnivariateDistribution}, bounds::Nothing)
    parameter_support(d)
end
function _get_bounds(d::Type{<:UnivariateDistribution}, bounds)
    params = parameters(d)

    foreach(params) do p
        if p ∉ keys(bounds)
            bounds = (; p=> REAL_SUPPORT, bounds...)
        end
    end
    bounds
end

function prepare_mle(d::Type{<:UnivariateDistribution},
                     data::AbstractArray;
                     start=nothing,
                     bounds=nothing,
                     holding=nothing,
                     kwargs...)

    params = parameters(d)

    start = isnothing(start) ? mle_starting_point(d, data) : start
    holding = isnothing(holding) ? (;) : holding

    to_solve = Bool[(p ∈ keys(start)) && (p ∉ keys(holding)) for p in params]

    starting_point = Float64[
        get(holding, p, get(start, p, missing)) |> Float64
        for p in params
    ]
    @assert !any(ismissing(v) for v in starting_point)

    custom_bounds = !isnothing(bounds)
    bounds = _get_bounds(d, bounds)
    constraints = make_constraints(d, start, holding, bounds)

    MLEConfig(d,
              data,
              params,
              to_solve,
              starting_point,
              constraints,
              custom_bounds)
end

function fit_numerical_mle(d::Type{<:UnivariateDistribution},
                 data::AbstractArray,
                 start,
                 bounds=nothing;
                 kwargs...)
    fit_numerical_mle(d, data; start=start, bounds=bounds, kwargs...)
end

"""
    fit_numerical_mle(D, data[, start, bounds, holding])

Use numerical methods to fit a MLE from a distribution with twice
differentiable density. MLEs are fit using Optim with the IPNewton constrained
solver.

"""
function fit_numerical_mle(d::Type{<:UnivariateDistribution},
                 data::AbstractArray;
                 start=nothing,
                 bounds=nothing,
                 holding=nothing,
                 kwargs...)

    config = prepare_mle(d, data; start, bounds, holding, kwargs...)
    mle_solve(config; kwargs...)
end

function mle_solve(config::MLEConfig; kwargs...)
    parameters = copy(config.starting_point)
    to_solve = config.params_to_solve

    function objective!(x)
        parameters[to_solve] = x
        dist = config.distribution(parameters...)

        -loglikelihood(dist, config.data)
    end

    vars = parameters[to_solve]
    func = TwiceDifferentiable(objective!, vars)

    opt = optimize(func,
                   config.constraints,
                   vars,
                   IPNewton(),
                   Optim.Options(),
                  )

    converged = Optim.converged(opt)
    loglike = -Optim.minimum(opt)
    result = Optim.minimizer(opt)
    fit_parameters = (; zip(config.params, parameters)...)
    n = Optim.iterations(opt)

    summary = (;
        converged,
        type=(all(to_solve) && !config.custom_bounds ? :MLE : :Restricted),
        fit_parameters,
        fit=config.distribution(parameters...; check_args=false),
        result,
        loglikelihood=loglike,
        n
    )

    if converged
        objective!(result)
        H = Optim.hessian!(func, result)

        vcov = inv(H)

        D = diag(vcov)
        if any(D .< 0)
            D = Complex.(D)
        end
        sd = sqrt.(D)

        (; summary..., H, vcov, sd)
    else
        summary
    end
end

export parameter_support, mle_starting_point, fit_numerical_mle

end # module


import Voronoi

using Random
using Printf
include("Draw.jl")

function randf(start, finish, n)
   v = rand(n)
   return map(x -> start + x*(finish-start), v)
end

function gradientDescent()
   WIDTH = 100.0
   HEIGHT = 100.0

   seed = 100
   println("Random seed: ", seed)
   Random.seed!(seed)

   command = nothing
   for n in 2:100
      points = convert(Array{Tuple{Real, Real}}, collect(zip(randf(1, WIDTH-1, n), randf(1, HEIGHT-1, n))))

      iterations = 1
      fortunes = 1

      time = @elapsed V, f, ∇f, ξ = Voronoi.Optimization.init(points, WIDTH, HEIGHT)

      μ = 0.5
      α = 0.01
      ϵ = 10^-1
      λ₀ = 0.01

      times = time

      while ξ > ϵ
         @printf("f = %.10g\t\t|∇f| = %.10g\t\ttime: %s seconds\n", f, ξ, time)
         Draw.init(WIDTH, HEIGHT)
         Draw.voronoiDiagram(V)

         d = ∇f

         time = @elapsed V, f, ∇f, ξ, forts = Voronoi.Optimization.lineSearch(points, d, α, λ₀, μ, ϵ, V, f, ∇f, ξ)

         times += time
         fortunes += forts
         iterations += 1

         if command != "a"
            println("Press Return for a step or enter \"a\" to animate until the end.")
            command = readline(stdin)
         else
            sleep(0.0000001) # make sure stuff is drawn
         end
      end

      println("n: ", n, "\t\titer: ", iterations, "\t\tfortunes: ", fortunes, "\t\ttime: ", times)
      #Draw.init(WIDTH, HEIGHT)
      #Draw.voronoiDiagram(V)
      #command = readline(stdin)
   end
end

gradientDescent()


function loadpuzzle(s::String)
    s = replace(s, r"\s" => "")
    s = map(c -> c in "-._" ? 0 : parse(Int, c), s |> collect)
    reshape(s, 9, 9) |> permutedims
end
loadpuzzle(mat::Matrix{<:Integer}) = copy(mat)
function possinit!(P::Matrix{UInt16}, grid::Matrix)
    P .= 0
    for r in 1:9
        for c in 1:9
            p = (r - 1) ÷ 3 * 3 + (c - 1) ÷ 3 + 1
            v = grid[r, c]
            @assert 0 <= v <= 9
            if v > 0
                bm = 1 << (v-1)
                P[1, r] |= bm
                P[2, c] |= bm
                P[3, p] |= bm
            end
        end
    end
end
function intlog2(v)
    if v == 1 << 4
        return 4
    elseif v < 1 << 4
        if v == 1 << 2
            return 2
        elseif v > 1 << 2
            return 3
        else
            return v - 1
        end
    else
        if v == 1 << 7
            return 7
        elseif v > 1 << 7
            return 8
        elseif v < 1 << 6
            return 5
        else
            return 6
        end
    end
end
# 朴素深搜（顺序搜索，逐列）
function sudokudfs!(grid::Matrix, P::Matrix{UInt16}=zeros(UInt16, 3, 9), id=1)
    if id == 82
        return true
    end
    r = (id - 1) % 9 + 1
    c = (id - 1) ÷ 9 + 1
    p = (r - 1) ÷ 3 * 3 + (c - 1) ÷ 3 + 1
    if grid[r, c] == 0
        @inbounds mask = P[1, r] | P[2, c] | P[3, p]
        while mask < 1 << 9 - 1
            mask2 = mask | (mask+1)
            bm = mask2 - mask
            @inbounds P[1, r] |= bm
            @inbounds P[2, c] |= bm
            @inbounds P[3, p] |= bm
            if sudokudfs!(grid, P, id + 1)
                grid[r, c] = intlog2(bm) + 1
                return true
            end
            @inbounds P[1, r] &= ~bm
            @inbounds P[2, c] &= ~bm
            @inbounds P[3, p] &= ~bm
            mask = mask2
        end
        return false
    else
        return sudokudfs!(grid, P, id + 1)
    end
end
function naivesolver!(grid, P::Matrix{UInt16}=zeros(UInt16, 3, 9); check=true)
    possinit!(P, grid)
    valid = sudokudfs!(grid, P)
    return check ? (valid ? grid : nothing) : grid
end

# 重排搜索顺序
function sudokudfs!(grid::Matrix{<:Integer}, P::Matrix{UInt16}, order::Vector, id=1)
    if id == 82
        return true
    end
    id2 = order[id]
    r = (id2 - 1) % 9 + 1
    c = (id2 - 1) ÷ 9 + 1
    p = (r - 1) ÷ 3 * 3 + (c - 1) ÷ 3 + 1
    if grid[r, c] == 0
        @inbounds mask = P[1, r] | P[2, c] | P[3, p]
        while mask < 1 << 9 - 1
            mask2 = mask | (mask+1)
            bm = mask2 - mask
            @inbounds P[1, r] |= bm
            @inbounds P[2, c] |= bm
            @inbounds P[3, p] |= bm
            if sudokudfs!(grid, P, id + 1)
                grid[r, c] = intlog2(bm) + 1
                return true
            end
            @inbounds P[1, r] &= ~bm
            @inbounds P[2, c] &= ~bm
            @inbounds P[3, p] &= ~bm
            mask = mask2
        end
        return false
    else
        return sudokudfs!(grid, P, id + 1)
    end

end
function presort()
    ids = [[1:9...], [9:-1:1...],[1:9...], [9:-1:1...],
    [1:9...], [9:-1:1...],[1:9...], [9:-1:1...],[1:9...]]# 一条龙序
    ids = vcat(broadcast((a, b) -> a .+ b, ids, 0:9:72)...)
end
const ORDER = presort()
function reordersolver!(grid, P::Matrix{UInt16}=zeros(UInt16, 3, 9), order=ORDER; check=true)
    possinit!(P, grid)
    valid = sudokudfs!(grid, P, order)
    return check ? (valid ? grid : nothing) : grid
end

# 优先深搜
function possinit!(P::AbstractArray{T,3}, b::AbstractArray{W,2}) where {T,W}
    P .= 0
    for r in 1:9
        for c in 1:9
            if b[r, c] > 0
                hold!(P, r, c, b[r, c])
                P[r, c, :] .+= 1
                P[r, c, b[r, c]] -= 1
            end
        end
    end
end
function hold!(P::AbstractArray{T,3}, r, c, v) where T
    for i in 1:c - 1
        P[r, i, v] += 1
    end
    for i in c + 1:9
        P[r, i, v] += 1
    end
    for i in 1:r - 1
        P[i, c, v] += 1
    end
    for i in r + 1:9
        P[i, c, v] += 1
    end
    rr = (r - 1) % 3 + 1
    rc = (c - 1) % 3 + 1
    s = ((1, 2), (-1, 1), (-2, -1))
    for i in s[rr]
        for j in s[rc]
            P[r + i, c + j, v] += 1
        end
    end
end

function release!(P::AbstractArray{T,3}, r, c, v) where T
    for i in 1:c - 1
        P[r, i, v] -= 1
    end
    for i in c + 1:9
        P[r, i, v] -= 1
    end
    for i in 1:r - 1
        P[i, c, v] -= 1
    end
    for i in r + 1:9
        P[i, c, v] -= 1
    end
    rr = (r - 1) % 3 + 1
    rc = (c - 1) % 3 + 1
    s = ((1, 2), (-1, 1), (-2, -1))
    for i in s[rr]
        for j in s[rc]
            P[r + i, c + j, v] -= 1
        end
    end
end
mutable struct KVpair{KT,VT}
    key::KT
    value::VT
end
Base.first(p::KVpair) = p.key
Base.last(p::KVpair) = p.value

function heapinit!(P::AbstractArray{T,3}) where T
    pocount = sum(P .== 0, dims=3)[:]
    heap = tracHeap([KVpair(v, id) for (id, v) in enumerate(pocount)])
    buildminheap!(heap)
    return heap
end
function increasehold(P::AbstractArray{T,3}, H::tracHeap, r, c, v) where T
    if P[r, c, v] == 0
        id = (c - 1) * 9 + r
        lc = H.locL[id]
#         @assert last(H[lc])==id
        if lc <= length(H)
            oldK = first(H[lc])
            H[lc].key -= 1
            heapchangekey!(H, lc, H[lc], oldK)
        end

    end
P[r, c, v] += 1
end

function decreasehold(P::AbstractArray{T,3}, H::tracHeap, r, c, v) where T # for release!
    if P[r, c, v] == 1
        id = (c - 1) * 9 + r
        lc = H.locL[id]
#         @assert last(H[lc])==id
        if lc <= length(H)
#             oldK = first(H[lc])
            H[lc].key += 1
#             heapchangekey!(H, lc, H[lc]) #交给undo!去完成结构调整
        end
    end
P[r, c, v] -= 1
end

function hold!(P::AbstractArray{T,3}, r, c, v, H::tracHeap) where T
    for i in 1:c - 1
        increasehold(P, H, r, i, v)
    end
    for i in c + 1:9
        increasehold(P, H, r, i, v)
    end
    for i in 1:r - 1
        increasehold(P, H, i, c, v)
    end
    for i in r + 1:9
        increasehold(P, H, i, c, v)
    end
    rr = (r - 1) % 3 + 1
    rc = (c - 1) % 3 + 1
    s = ((1, 2), (-1, 1), (-2, -1))
    for i in s[rr]
        for j in s[rc]
            increasehold(P, H, r + i, c + j, v)
        end
    end
end

function release!(P::AbstractArray{T,3}, r, c, v, H::tracHeap) where T
    for i in 1:c - 1
        decreasehold(P, H, r, i, v)
    end
    for i in c + 1:9
        decreasehold(P, H, r, i, v)
    end
    for i in 1:r - 1
        decreasehold(P, H, i, c, v)
    end
    for i in r + 1:9
        decreasehold(P, H, i, c, v)
    end
    rr = (r - 1) % 3 + 1
    rc = (c - 1) % 3 + 1
    s = ((1, 2), (-1, 1), (-2, -1))
    for i in s[rr]
        for j in s[rc]
            decreasehold(P, H, r + i, c + j, v)
        end
    end
end
function sudokudfs!(P::AbstractArray{T,3}, H::tracHeap, id=last(heapextractmin!(H))) where T
    r = (id - 1) % 9 + 1
    c = (id - 1) ÷ 9 + 1
    if length(H) == 0
        for v in 1:9
            if P[r, c, v] == 0
                return true
            end
        end
        return false
    end
    for v in 1:9
        if P[r, c, v] == 0
            snapshot1 = historylen(H)
            hold!(P, r, c, v, H)
            snapshot2 = historylen(H)
            id = last(heapextractmin!(H))
            if sudokudfs!(P, H, id)
#                 @show id
                return true
            end
            undo!(H, tostep=snapshot2)
            release!(P, r, c, v, H)
            undo!(H, tostep=snapshot1)
        end
    end
    return false
end
function prioritysolver!(b, P::AbstractArray{T,3}=zeros(Int, 9, 9, 9); check=true) where T
    possinit!(P, b)
    H = heapinit!(P)
    valid = sudokudfs!(P, H)
    ans = reshape(mapslices(argmin, P, dims=3), 9, 9)
    return check ? (valid ? ans : nothing) : ans
end
function solvesudoku!(sudoku::Matrix{<:Integer}, args...; solver=naivesolver!, kargs...)
    solver(sudoku, args...; kargs...)
end
solvesudoku(sudoku, args...; kargs...) = solvesudoku!(loadpuzzle(sudoku), args...; kargs...)


function horner(coeff, x0)
	degree = length(coeff) - 1
	y = coeff[1]
	z = 0
	for j ∈ 2:(degree + 1)
		println(coeff[j])
		z = y + x0 * z
		y = x0 * y + coeff[j]
	end
	return (y, z)
end

println(horner([5, 2, 1], 5)[1])


#
###
# AbstractSpace interface
###

### getters
"""
args:
    space::AbstractSpace{T,D}
ret:
    (x1, ..., xD,) # incl end points
"""
function get_grid end

"""
args:
    space::AbstractSpace{T,D}
ret:
    AbstractField{Integer, D}
"""
function get_global_numbering end

"""
get domain

args:
    space::AbstractSpace
ret:
    AbstractDomain
"""
function get_domain end

"""
args:
    space::AbstractSpace{T,D}
    i::Integer
ret:
    ith basis function that can be evaluated
    anywhere in Space.domain
"""
function get_basis end

### interpolation
"""
Interpolate function values to to points.

args:
    points::vector of coordinates
    u::AbstractField
    space::AbstractSpace{T,D}
ret:
    u interpolated to points
"""
function interp end

"""
Point-to-point interpolant between
spaces on same domain. used for
dealiasing

args:
    space1::AbstractSpace{T,D}
    space2::AbstractSpace{T,D}
ret:
    interpolation operator from
    space1 to space2
"""
function interpOp end

### vector calculus ops
"""
Gradient Operator
Compute gradient of u∈H¹(Ω).

Continuity isn't enforced across
element boundaries for gradients

args:
    space::AbstractSpace{T,D}
ret:
    gradOp: u -> [dudx1, ..., dudxD]
"""
function gradOp end

"""
Mass Operator

args:
    space::AbstractSpace{T,D}
    space_dealias
ret:
    massOp: AbstractField -> AbstractField
"""
function massOp end

"""
Laplace Operator

for v,u in H¹₀(Ω)

(v,-∇² u) = (vx,ux) + (vy,uy)\n
         := a(v,u)\n

args:
    space::AbstractSpace{T,D}
    space_dealias
ret:
    laplaceOp: AbstractField -> AbstractField
"""
function laplaceOp(space::AbstractSpace)
    D = gradOp(space)
    M = massOp(space)

    lapl = D' * [M] * D

    first(lapl)
end

"""
args:
    space::AbstractSpace{<:Number,D}
    vel...::AbstractField{<:Number,D}
    space_dealias
ret:
    advectionOp: AbstractField -> AbstractField

for v,u,T in H¹₀(Ω)

(v,(u⃗⋅∇)T) = (v,ux*∂xT + uy*∂yT)\n
           = v' *B*(ux*∂xT + uy*∂yT)\n

implemented as

R'R * QQ' * B * (ux*∂xT + uy*∂yT)

(u⃗⋅∇)T = ux*∂xT + uy*∂yT

       = [ux uy] * [Dx] T
                   [Dx]
"""
function advectionOp(space::AbstractSpace{<:Number,D},
                     vel::AbstractField{<:Number,D}...
                    ) where{D}
    V = [DiagonalOp.(vel)...]

    Dx = gradOp(space)
    M  = massOp(space)

    advectOp = V' * [M] * Dx

    first(advectOp)
end

"""
Divergence
"""
function divergenceOp end # TODO write fallback using gradOp

#function divergence(vec::Vector{AbstractSpace{<:Number,D}}, space::AbstractSpace{<:Number,D}) where{D}
#    Dx = grad(space)
#    div = NullOp{D}()
#    for i=1:D
#        Dxi = Dx[i]
#        Vi  = V[i]
#        div += Dxi(V[i])
#    end
#
#    div
#end

### dealiased operators

function massOp(space1::AbstractSpace{<:Number,D},
                space2::AbstractSpace{<:Number,D};
                J = nothing,
               ) where{D}
    J12 = J !== nothing ? J : interpOp(space1, space2)

    M2 = massOp(space2)

    J12' ∘ M2 ∘ J12
end

function laplaceOp(space1::AbstractSpace{<:Number,D},
                   space2::AbstractSpace{<:Number,D};
                   J = nothing,
                  ) where{D}
    J12 = J !== nothing ? J : interpOp(space1, space2)

    M2 = massOp(space2)
    D1 = gradOp(space1)
    JD = [J12] .∘ D1

    laplOp = JD' * [M2] * JD

    first(laplOp)
end

"""
ux,uy, ∇T are interpolated to
a grid with higher polynomial order
for dealiasing (over-integration)
so we don't commit any
"variational crimes"
"""
function advectionOp(space1::AbstractSpace{<:Number,D},
                     space2::AbstractSpace{<:Number,D},
                     vel::AbstractField{<:Number,D}...;
                     J = nothing,
                    ) where{D}
    J12 = J !== nothing ? J : interpOp(space1, space2)

    V1 = [DiagonalOp.(vel)...]
    V2 = J12 .* V1

    M2 = massOp(space2)
    D1 = gradOp(space1)

    advectOp = [J12]' * V2' * [M2] * [J12] * D1

    first(advectOp)
end

###
# AbstractSpectralSpace
###

"""
get number of points
"""
function get_numpoints end

###
# Tensor Product Spaces
###

#struct TensorProductSpace{T,D1+D2,
#                          Tspace1<:AbstractSpace{<:Number, D1},
#                          Tspace2<:AbstractSpace{<:Number, D2},
#                         } <: AbstractTensorProductSpace{T,D1+D2}
#    space1::Tspace1
#    space2::Tspace2
#end
#


import ProgressMeter
#######################################################################################
# Original correlation sum
#######################################################################################
using Distances, Roots
export correlationsum, grassberger_dim, boxed_correlationsum,
estimate_r0_buenoorovio, data_boxing, autoprismdim, estimate_r0_theiler

"""
    correlationsum(X, ε::Real; w = 0, norm = Euclidean(), q = 2) → C_q(ε)
Calculate the `q`-order correlation sum of `X` (`Dataset` or timeseries)
for a given radius `ε` and `norm`. They keyword `show_progress = false` can be used
to display a progress bar for large `X`.

The function [`boxed_correlationsum`](@ref) is faster and should be preferred over this one.

## Description
The correlation sum is done using the formula:
```math
C_2(\\epsilon) = \\frac{2}{(N-w)(N-w-1)}\\sum_{i=1}^{N}\\sum_{j=1+w+i}^{N} B(||X_i - X_j|| < \\epsilon)
```
for `q=2` and
```math
C_q(\\epsilon) = \\left[\\frac{1}{\\alpha} \\sum_{i=w+1}^{N-w}\\left[\\sum_{j:|i-j| > w} B(||X_i - X_j|| < \\epsilon)\\right]^{q-1}\\right]^{1/(q-1)}
```
where
```math
\\alpha = (N-2w)(N-2w-1)^{(q-1)}
```
for `q≠2`, where ``N`` is its length and ``B`` gives 1 if the argument is
`true`. `w` is the [Theiler window](@ref). If `ε` is a vector its values have to be
ordered. See the article of Grassberger for the general definition [^Grassberger2007] and the book "Nonlinear Time Series Analysis" [^Kantz2003], Ch. 6, for
a discussion around `w` and choosing best values and Ch. 11.3 for the
explicit definition of the q-order correlationsum.

    correlationsum(X, εs::AbstractVector; w, norm, q) → C_q(ε)

If `εs` is a vector, `C_q` is calculated for each `ε ∈ εs`.
If also `q=2`, we attempt to do further optimizations are done, if the allocation
a matrix of size `N×N` is possible

See [`grassberger`](@ref) for more. See also [`takens_best_estimate`](@ref).

[^Grassberger]: Peter Grassberger (2007) [Grassberger-Procaccia algorithm. Scholarpedia, 2(5):3043.](http://dx.doi.org/10.4249/scholarpedia.3043)

[^Kantz]: Kantz, H., & Schreiber, T. (2003). [More about invariant quantities. In Nonlinear Time Series Analysis (pp. 197-233). Cambridge: Cambridge University Press.](https://doi:10.1017/CBO9780511755798.013)
"""
function correlationsum(X, ε; q = 2, norm = Euclidean(), w = 0, show_progress = false)
    if q == 2
        correlationsum_2(X, ε, norm, w, show_progress)
    else
        correlationsum_q(X, ε, q, norm, w, show_progress)
    end
end

function correlationsum_2(X, ε::Real, norm, w, show_progress)
    N = length(X)
    if show_progress
        progress = ProgressMeter.Progress(N; desc = "Correlation sum: ", dt = 1.0)
    end
    C = zero(eltype(X))
    @inbounds for (i, x) in enumerate(X)
        for j in i+1+w:N
            C += evaluate(norm, x, X[j]) < ε
        end
        show_progress && ProgressMeter.next!(progress)
    end
    return C * 2 / ((N-w-1)*(N-w))
end

function correlationsum_q(X, ε::Real, q, norm, w, show_progress)
    N, C = length(X), zero(eltype(X))
    normalisation = (N-2w)*(N-2w-one(eltype(X)))^(q-1)
    if show_progress
        progress = ProgressMeter.Progress(length(1+w:N-w); desc = "Correlation sum: ", dt = 1.0)
    end
    for i in 1+w:N-w
        x = X[i]
        C_current = zero(eltype(X))
        # computes all distances from 0 up to i-w
        for j in 1:i-w-1
            C_current += evaluate(norm, x, X[j]) < ε
        end
        # computes all distances after i+w till the end
        for j in i+w+1:N
            C_current += evaluate(norm, x, X[j]) < ε
        end
        C += C_current^(q - 1)
        show_progress && ProgressMeter.next!(progress)
    end
    return (C / normalisation) ^ (1 / (q-1))
end

# Optimized version
function correlationsum_2(X, εs::AbstractVector, norm, w, show_progress)
    @assert issorted(εs) "Sorted εs required for optimized version."
    d = try
        distancematrix(X, norm)
    catch err
        @warn "Couldn't create distance matrix ($(typeof(err))). Using slower algorithm..."
        return [correlationsum_2(X, ε, norm, w, show_progress) for ε in εs]
    end
    return correlationsum_2_fb(X, εs, d, w, show_progress) # function barrier
end
function correlationsum_2_fb(X, εs, d, w, show_progress)
    Cs = zeros(eltype(X), length(εs))
    N = length(X)
    factor = 2/((N-w)*(N-1-w))
    if show_progress
        K = length(length(εs)÷2:-1:1)
        M = K + length((length(εs)÷2 + 1):length(εs))
        progress = ProgressMeter.Progress(M; desc = "Correlation sum: ", dt = 1.0)
    end

    # First loop: mid-way ε until lower saturation point (C=0)
    for (ki, k) in enumerate(length(εs)÷2:-1:1)
        ε = εs[k]
        for i in 1:N
            @inbounds Cs[k] += count(d[j, i] < ε for j in i+1+w:N)
        end
        show_progress && ProgressMeter.update!(progress, ki)
        Cs[k] == 0 && break
    end
    # Second loop: mid-way ε until higher saturation point (C=max)
    for (ki, k) in enumerate((length(εs)÷2 + 1):length(εs))
        ε = εs[k]
        for i in 1:N
            @inbounds Cs[k] += count(d[j, i] < ε for j in i+1+w:N)
        end
        show_progress && ProgressMeter.update!(progress, ki+K)
        if Cs[k] ≈ 1/factor
            Cs[k:end] .= 1/factor
            break
        end
    end
    show_progress && ProgressMeter.finish!(progress)
    return Cs .* factor
end

function correlationsum_q(X, εs::AbstractVector, q, norm, w, show_progress)
    @assert issorted(εs) "Sorted εs required for optimized version."
    Nε, T, N = length(εs), eltype(X), length(X)
    Cs = zeros(T, Nε)
    normalisation = (N-2w)*(N-2w-one(T))^(q-1)
    if show_progress
        progress = ProgressMeter.Progress(length(1+w:N-w); desc = "Correlation sum: ", dt = 1.0)
    end
    for i in 1+w:N-w
        x = X[i]
        C_current = zeros(T, Nε)
        # Compute distances for j outside the Theiler window
        for j in Iterators.flatten((1:i-w-1, i+w+1:N))
            dist = evaluate(norm, x, X[j])
            for k in Nε:-1:1
                if dist < εs[k]
                    C_current[k] += 1
                else
                    break
                end
            end
        end
        Cs .+= C_current .^ (q-1)
        show_progress && ProgressMeter.next!(progress)
    end
    return (Cs ./ normalisation) .^ (1/(q-1))
end

function distancematrix(X, norm = Euclidean())
    N = length(X)
    d = zeros(eltype(X), N, N)
    @inbounds for i in 1:N
        for j in i+1:N
            d[j, i] = evaluate(norm, X[i], X[j])
        end
    end
    return d
end

"""
    grassberger_dim(data, εs = estimate_boxsizes(data); kwargs...) → D_C
Use the method of Grassberger and Proccacia[^Grassberger1983], and the correction by
Theiler[^Theiler1986], to estimate the correlation dimension `D_C` of the given `data`.

This function does something extremely simple:
```julia
cm = correlationsum(data, εs; kwargs...)
return linear_region(log.(sizes), log(cm))[2]
```
i.e. it calculates [`correlationsum`](@ref) for various radii and then tries to find
a linear region in the plot of the log of the correlation sum versus log(ε).
See [`generalized_dim`](@ref) for a more thorough explanation.

See also [`takens_best_estimate`](@ref).

[^Grassberger1983]: Grassberger and Proccacia, [Characterization of strange attractors, PRL 50 (1983)](https://journals-aps-org.e-bis.mpimet.mpg.de/prl/abstract/10.1103/PhysRevLett.50.346)

[^Theiler1986]: Theiler, [Spurious dimension from correlation algorithms applied to limited time-series data. Physical Review A, 34](https://doi.org/10.1103/PhysRevA.34.2427)
"""
function grassberger_dim(data::AbstractDataset, εs = estimate_boxsizes(data); kwargs...)
    cm = correlationsum(data, εs; kwargs...)
    return linear_region(log.(εs), log.(cm))[2]
end

################################################################################
# Boxed Correlation sum (we distribute data to boxes beforehand)
################################################################################
"""
    boxed_correlationsum(X::Dataset, εs, r0 = maximum(εs); kwargs...) → Cs

Estimate the box assisted q-order correlation sum[^Kantz2003] `Cs` out of a
dataset `X` for each radius in `εs`, by splitting the data into boxes of size `r0`
beforehand. This method is much faster than [`correlationsum`](@ref), **provided that** the
box size `r0` is significantly smaller than then the attractor length.
A good estimate for `r0` is [`estimate_r0_buenoorovio`](@ref).

    boxed_correlationsum(X; kwargs...) → εs, Cs

In this method the minimum inter-point distance and [`estimate_r0_buenoorovio`](@ref)
are used to estimate good `εs` for the calculation, which are also returned.

## Keywords
* `q = 2`
* `P = autoprismdim(data)` : The prism dimension.
* `w = 0` : The [Theiler window](@ref).
* `show_progress = false` : Whether to display a progress bar for the calculation.

## Description
`C_q(ε)` is calculated for every `ε ∈ εs` and each of the boxes to then be
summed up afterwards. The method of splitting the data into boxes was
implemented according to Theiler[^Theiler1987]. `w` is the [Theiler window](@ref).
`P` is the prism dimension. If `P` is unequal to the dimension of the data, only the
first `P` dimensions are considered for the box distribution (this is called the
prism-assisted version). By default `P` is choosen automatically.

The function is explicitly optimized for `q = 2` but becomes quite slow for `q ≠ 2`.

See [`correlationsum`](@ref) for the definition of `C_q`
and also [`data_boxing`](@ref) to use the algorithm that splits data into boxes.

[^Kantz]: Kantz, H., & Schreiber, T. (2003). [More about invariant quantities. In Nonlinear Time Series Analysis (pp. 197-233). Cambridge: Cambridge University Press.](https://doi:10.1017/CBO9780511755798.013)

[^Theiler1987]: Theiler, [Efficient algorithm for estimating the correlation dimension from a set of discrete points. Physical Review A, 36](https://doi.org/10.1103/PhysRevA.36.4456)
"""
function boxed_correlationsum(data; q = 2, P = autoprismdim(data), w = 0, show_progress=false)
    r0, ε0 = estimate_r0_buenoorovio(data, P)
    @assert  r0 < ε0 "The calculated box size was smaller than the minimum interpoint " *
    "distance. Please choose manually."
    εs = 10.0 .^ range(log10(ε0), log10(r0), length = 16)
    boxed_correlationsum(data, εs, r0; q, P, w, show_progress)
end

function boxed_correlationsum(
        data, εs, r0 = maximum(εs);
        q = 2, P = autoprismdim(data), w = 0,
        show_progress = false,
    )
    @assert P ≤ size(data, 2) "Prism dimension has to be lower or equal than " *
    "data dimension."
    boxes, contents = data_boxing(data, r0, P)
    if q == 2
        boxed_correlationsum_2(boxes, contents, data, εs; w, show_progress)
    else
        boxed_correlationsum_q(boxes, contents, data, εs, q; w, show_progress)
    end
end

"""
    autoprismdim(data, version = :bueno)

An algorithm to find the ideal choice of a prism dimension for [`boxed_correlationsum`](@ref).
`version = :bueno` uses `P=2`, while `version = :theiler` uses Theiler's original suggestion.
"""
function autoprismdim(data, version = :bueno)
    D = dimension(data)
    N = length(data)
    if version == :bueno
        return min(D, 2)
    elseif version == :theiler
        if D > 0.75 * log2(N)
            return max(2, ceil(0.5 * log2(N)))
        else
            return D
        end
    end
end

"""
    data_boxing(data, r0, P) → boxes, contents
Distribute the `data` points into boxes of size `r0`. Return box positions
and the contents of each box as two separate vectors. Implemented according to
the paper by Theiler[^Theiler1987] improving the algorithm by Grassberger and
Procaccia[^Grassberger1983]. If `P` is smaller than the dimension of the data,
only the first `P` dimensions are considered for the distribution into boxes.

See also: [`boxed_correlationsum`](@ref).

[^Theiler1987]: Theiler, [Efficient algorithm for estimating the correlation dimension from a set of discrete points. Physical Review A, 36](https://doi.org/10.1103/PhysRevA.36.4456)

[^Grassberger1983]: Grassberger and Proccacia, [Characterization of strange attractors, PRL 50 (1983)](https://journals-aps-org.e-bis.mpimet.mpg.de/prl/abstract/10.1103/PhysRevLett.50.346)
"""
function data_boxing(data, r0, P)
    @assert P ≤ size(data, 2) "Prism dimension has to be ≤ than data dimension."
    mini = minima(data)[1:P]

    # Map each datapoint to its bin edge and sort the resulting list:
    bins = map(point -> floor.(Int, (point[1:P] - mini)/r0), data)
    permutations = sortperm(bins, alg=QuickSort)

    boxes = unique(bins[permutations])
    contents = Vector{Vector{Int}}()
    sizehint!(contents, length(boxes))

    prior, prior_perm = 1, permutations[1]
    # distributes all permutation indices into boxes
    for (index, perm) in enumerate(permutations)
        if bins[perm] ≠ bins[prior_perm]
            push!(contents, permutations[prior:index-1])
            prior, prior_perm = index, perm
        end
    end
    push!(contents, permutations[prior:end])

    Dataset(boxes), contents
end

"""
    boxed_correlationsum_2(boxes, contents, data, εs; w = 0)
For a vector of `boxes` and the indices of their `contents` inside of `data`,
calculate the classic correlationsum of a radius or multiple radii `εs`.
`w` is the Theiler window, for explanation see [`boxed_correlationsum`](@ref).
"""
function boxed_correlationsum_2(boxes, contents, data, εs; w = 0, show_progress = false)
    Cs = zeros(eltype(data), length(εs))
    N = length(data)
    M = length(boxes)
    if show_progress
        progress = ProgressMeter.Progress(M; desc = "Boxed correlation sum: ", dt = 1.0)
    end
    for index in 1:M
        indices_neighbors = find_neighborboxes_2(index, boxes, contents)
        indices_box = contents[index]
        Cs .+= inner_correlationsum_2(indices_box, indices_neighbors, data, εs; w)
        show_progress && ProgressMeter.update!(progress, index)
    end
    Cs .* (2 / ((N - w) * (N - w - 1)))
end

"""
    find_neighborboxes_2(index, boxes, contents) → indices
For an `index` into `boxes` all neighbouring boxes beginning from the current
one are searched. If the found box is indeed a neighbour, the `contents` of
that box are added to `indices`.
"""
function find_neighborboxes_2(index, boxes, contents)
    indices = Int[]
    box = boxes[index]
    N_box = length(boxes)
    for index2 in index:N_box
        if evaluate(Chebyshev(), box, boxes[index2]) < 2
            append!(indices, contents[index2])
        end
    end
    indices
end

"""
    inner_correlationsum_2(indices_X, indices_Y, data, εs; norm = Euclidean(), w = 0)
Calculates the classic correlation sum for values `X` inside a box, considering
`Y` consisting of all values in that box and the ones in neighbouring boxes for
all distances `ε ∈ εs` calculated by `norm`. To obtain the position of the
values in the original time series `data`, they are passed as `indices_X` and
`indices_Y`.

`w` is the Theiler window. Each index to the original array is checked for the
distance of the compared index. If this absolute value is not higher than `w`
its element is not used in the calculation of the correlationsum.

See also: [`correlationsum`](@ref)
"""
function inner_correlationsum_2(indices_X, indices_Y, data, εs; norm = Euclidean(), w = 0)
    @assert issorted(εs) "Sorted εs required for optimized version."
    Cs, Ny, Nε = zeros(length(εs)), length(indices_Y), length(εs)
    for (i, index_X) in enumerate(indices_X)
    	x = data[index_X]
        for j in i+1:Ny
            index_Y = indices_Y[j]
            # Check for Theiler window.
            if abs(index_Y - index_X) > w
                # Calculate distance.
		        dist = evaluate(norm, data[index_Y], x)
		        for k in Nε:-1:1
		            if dist < εs[k]
		                Cs[k] += 1
		            else
		                break
		            end
		        end
		    end
        end
    end
    return Cs
end

"""
    boxed_correlationsum_q(boxes, contents, data, εs, q; w = 0)
For a vector of `boxes` and the indices of their `contents` inside of `data`,
calculate the `q`-order correlationsum of a radius or radii `εs`.
`w` is the Theiler window, for explanation see [`boxed_correlationsum`](@ref).
"""
function boxed_correlationsum_q(boxes, contents, data, εs, q; w = 0, show_progress = false)
    q <= 1 && @warn "This function is currently not specialized for q <= 1" *
    " and may show unexpected behaviour for these values."
    Cs = zeros(eltype(data), length(εs))
    N = length(data)
    M = length(boxes)
    if show_progress
        progress = ProgressMeter.Progress(M; desc = "Boxed correlation sum: ", dt = 1.0)
    end
    for index in 1:M
        indices_neighbors = find_neighborboxes_q(index, boxes, contents, q)
        indices_box = contents[index]
        Cs .+= inner_correlationsum_q(indices_box, indices_neighbors, data, εs, q; w)
        show_progress && ProgressMeter.update!(progress, index)
    end
    clamp.((Cs ./ ((N - 2w) * (N - 2w - 1) ^ (q-1))), 0, Inf) .^ (1 / (q-1))
end

"""
    find_neighborboxes_q(index, boxes, contents, q) → indices
For an `index` into `boxes` all neighbouring boxes are searched. If the found
box is indeed a neighbour, the `contents` of that box are added to `indices`.
"""
function find_neighborboxes_q(index, boxes, contents, q)
    indices = Int[]
    box = boxes[index]
    for (index2, box2) in enumerate(boxes)
        if evaluate(Chebyshev(), box, box2) < 2
            append!(indices, contents[index2])
        end
    end
    indices
end

"""
    inner_correlationsum_q(indices_X, indices_Y, data, εs, q::Real; norm = Euclidean(), w = 0)
Calculates the `q`-order correlation sum for values `X` inside a box,
considering `Y` consisting of all values in that box and the ones in
neighbouring boxes for all distances `ε ∈ εs` calculated by `norm`. To obtain
the position of the values in the original time series `data`, they are passed
as `indices_X` and `indices_Y`.

`w` is the Theiler window. The first and last `w` points of this data set are
not used by themselves to calculate the correlationsum.

See also: [`correlationsum`](@ref)
"""
function inner_correlationsum_q(indices_X, indices_Y, data, εs, q::Real; norm = Euclidean(), w = 0)
    @assert issorted(εs) "Sorted εs required for optimized version."
    Cs = zeros(length(εs))
    N, Ny, Nε = length(data), length(indices_Y), length(εs)
    for i in indices_X
        # Check that this index is not within Theiler window of the boundary
        # This step is neccessary for easy normalisation.
        (i < w + 1 || i > N - w) && continue
        C_current = zeros(Nε)
        x = data[i]
        for j in indices_Y
            # Check that this index is not whithin the Theiler window
        	if abs(i - j) > w
                # Calculate the distance for the correlationsum
		        dist = evaluate(norm, x, data[j])
		        for k in Nε:-1:1
		            if dist < εs[k]
		                C_current[k] += 1
		            else
		                break
		            end
		        end
		    end
        end
        Cs .+= C_current .^ (q-1)
    end
    return Cs
end

#######################################################################################
# Good boxsize estimates for boxed correlation sum
#######################################################################################
"""
    estimate_r0_theiler(X::Dataset) → r0, ε0
Estimate a reasonable size for boxing the data `X` before calculating the
[`boxed_correlationsum`](@ref) proposed by Theiler[^Theiler1987].
Return the boxing size `r0` and minimum inter-point distance in `X`, `ε0`.

To do so the dimension is estimated by running the algorithm by Grassberger and
Procaccia[^Grassberger1983] with `√N` points where `N` is the number of total
data points. Then the optimal boxsize ``r_0`` computes as
```math
r_0 = R (2/N)^{1/\\nu}
```
where ``R`` is the size of the chaotic attractor and ``\\nu`` is the estimated dimension.

[^Theiler1987]: Theiler, [Efficient algorithm for estimating the correlation dimension from a set of discrete points. Physical Review A, 36](https://doi.org/10.1103/PhysRevA.36.4456)

[^Grassberger1983]: Grassberger and Proccacia, [Characterization of strange attractors, PRL 50 (1983)](https://journals-aps-org.e-bis.mpimet.mpg.de/prl/abstract/10.1103/PhysRevLett.50.346)
"""
function estimate_r0_theiler(data)
    N = length(data)
    mini, maxi = minmaxima(data)
    R = mean(maxi .- mini)
    # Sample √N datapoints for a rough estimate of the dimension.
    data_sample = data[unique(rand(1:N, ceil(Int, sqrt(N))))] |> Dataset
    # Define radii for the rough dimension estimate
    min_d, _ = minimum_pairwise_distance(data)
    if min_d == 0
        @warn(
        "Minimum distance in the dataset is zero! Probably because of having data "*
        "with low resolution, or duplicate data points. Setting to `d₊/1000` for now.")
        min_d = R/(10^3)
    end
    lower = log10(min_d)
    εs = 10 .^ range(lower, stop = log10(R), length = 12)
    # Actually estimate the dimension.
    cm = correlationsum(data_sample, εs)
    ν = linear_region(log.(εs), log.(cm), tol = 0.5, warning = false)[2]
    # The combination yields the optimal box size
    r0 = R * (2/N)^(1/ν)
    return r0, min_d
end

"""
    estimate_r0_buenoorovio(X::Dataset, P = autoprismdim(X))
Estimates a reasonable size for boxing the time series `X` proposed by
Bueno-Orovio and Pérez-García[^Bueno2007] before calculating the correlation
dimension as presented by Theiler[^Theiler1983]. If instead of boxes, prisms
are chosen everything stays the same but `P` is the dimension of the prism.
To do so the dimension `ν` is estimated by running the algorithm by Grassberger
and Procaccia[^Grassberger1983] with `√N` points where `N` is the number of
total data points.
An effective size `ℓ` of the attractor is calculated by boxing a small subset
of size `N/10` into boxes of sidelength `r_ℓ` and counting the number of filled
boxes `η_ℓ`.
```math
\\ell = r_\\ell \\eta_\\ell ^{1/\\nu}
```
The optimal number of filled boxes `η_opt` is calculated by minimising the number
of calculations.
```math
\\eta_\\textrm{opt} = N^{2/3}\\cdot \\frac{3^\\nu - 1}{3^P - 1}^{1/2}.
```
`P` is the dimension of the data or the number of edges on the prism that don't
span the whole dataset.

Then the optimal boxsize ``r_0`` computes as
```math
r_0 = \\ell / \\eta_\\textrm{opt}^{1/\\nu}.
```

[^Bueno2007]: Bueno-Orovio and Pérez-García, [Enhanced box and prism assisted algorithms for computing the correlation dimension. Chaos Solitons & Fractrals, 34(5)](https://doi.org/10.1016/j.chaos.2006.03.043)

[^Theiler1987]: Theiler, [Efficient algorithm for estimating the correlation dimension from a set of discrete points. Physical Review A, 36](https://doi.org/10.1103/PhysRevA.36.4456)

[^Grassberger1983]: Grassberger and Proccacia, [Characterization of strange attractors, PRL 50 (1983)](https://journals-aps-org.e-bis.mpimet.mpg.de/prl/abstract/10.1103/PhysRevLett.50.346)
"""
function estimate_r0_buenoorovio(X, P = autoprismdim(X))
    mini, maxi = minmaxima(X)
    N = length(X)
    R = mean(maxi .- mini)
    # The possibility of a bad pick exists, if so, the calculation is repeated.
    ν = zero(eltype(X))
    min_d, _ = minimum_pairwise_distance(X)
    if min_d == 0
        @warn(
        "Minimum distance in the dataset is zero! Probably because of having data "*
        "with low resolution, or duplicate data points. Setting to `d₊/1000` for now.")
        min_d = R/(10^3)
    end

    # Sample N/10 datapoints out of data for rough estimate of effective size.
    sample1 = X[unique(rand(1:N, N÷10))] |> Dataset
    r_ℓ = R / 10
    η_ℓ = length(data_boxing(sample1, r_ℓ, P)[1])
    r0 = zero(eltype(X))
    while true
        # Sample √N datapoints for rough dimension estimate
        sample2 = X[unique(rand(1:N, ceil(Int, sqrt(N))))] |> Dataset
        # Define logarithmic series of radii.
        εs = 10.0 .^ range(log10(min_d), log10(R); length = 16)
        # Estimate ν from a sample using the Grassberger Procaccia algorithm.
        cm = correlationsum(sample2, εs)
        ν = linear_region(log.(εs), log.(cm); tol = 0.5, warning = false)[2]
        # Estimate the effictive size of the chaotic attractor.
        ℓ = r_ℓ * η_ℓ^(1/ν)
        # Calculate the optimal number of filled boxes according to Bueno-Orovio
        η_opt = N^(2/3) * ((3^ν - 1/2) / (3^P - 1))^(1/2)
        # The optimal box size is the effictive size divided by the box number # to the power of the inverse dimension.
        r0 = ℓ / η_opt^(1/ν)
        !isnan(r0) && break
    end
    return r0, min_d
end


#----------------------------------------------------------------------------# Weight
abstract Weight
abstract BatchWeight <: Weight  # only BatchWeight types work for fitting by batch
abstract StochasticWeight <: BatchWeight

#-----------------------------------------------------------------------# EqualWeight
"""
One of the `Weight` types.  Observations are weighted equally.  For analytical
updates, the online algorithm will give results equal to the offline version.

- `EqualWeight()`
"""
type EqualWeight <: BatchWeight
    nobs::Int
    EqualWeight() = new(0)
    EqualWeight(n::Int) = new(n)
end

#-----------------------------------------------------------------# ExponentialWeight
"""
One of the `Weight` types.  Updates are performed with a constant weight
`λ = 2 / (1 + lookback)`.

- `ExponentialWeight(λ::Float64)`
- `ExponentialWeight(lookback::Int)`
"""
type ExponentialWeight <: Weight
    nobs::Int
    λ::Float64
    function ExponentialWeight(λ::Real = 1.0)
            @assert 0 <= λ <= 1
            new(0, λ)
    end
    ExponentialWeight(lookback::Integer) = ExponentialWeight(2.0 / (lookback + 1))
end

#----------------------------------------------------------------# BoundedEqualWeight
"""
One of the `Weight` types.  Uses `EqualWeight` until reaching `λ = 2 / (1 + lookback)`,
then weights are held constant.

- `BoundedEqualWeight(λ::Float64)`
- `BoundedEqualWeight(lookback::Int)`
"""
type BoundedEqualWeight <: Weight
    nobs::Int
    λ::Float64
    function BoundedEqualWeight(λ::Real = 1.0)
            @assert 0 <= λ <= 1
            new(0, λ)
    end
    BoundedEqualWeight(lookback::Integer) = BoundedEqualWeight(2.0 / (lookback + 1))
end

#----------------------------------------------------------------------# LearningRate
"""
One of the `Weight` types.  It's primary use is for the OnlineStats that use stochastic
approximation (`StatLearn`, `QuantReg`, `QuantileMM`, `QuantileSGD`, `NormalMix`, and
`KMeans`).  The weight at update `t` is `1 / t ^ r`.  When weights reach `λ`, they are
held consant.  Compare to `LearningRate2`.

- `LearningRate(r = 0.5, λ = 0.0)`
"""
type LearningRate <: StochasticWeight
    nobs::Int
    nups::Int
    r::Float64
    λ::Float64
    function LearningRate(r::Real = 0.5, λ::Real = 0.0)
        @assert 0 < r <= 1
        @assert λ >= 0
        new(0, 0, r, λ)
    end
end

#---------------------------------------------------------------------# LearningRate2
"""
One of the `Weight` types.  It's primary use is for the OnlineStats that use stochastic
approximation (`StatLearn`, `QuantReg`, `QuantileMM`, `QuantileSGD`, `NormalMix`, and
`KMeans`).  The weight at update `t` is `1 / (1 + c * (t - 1))`.  When weights reach
`λ`, they are held consant.  Compare to `LearningRate`.

- `LearningRate2(c = 0.5, λ = 0.0)`
"""
type LearningRate2 <: StochasticWeight
    nobs::Int
    nups::Int
    c::Float64
    λ::Float64
    function LearningRate2(c::Real = 0.5, λ = 0.0)
        @assert λ >= 0
        @assert c > 0
        new(0, 0, c, λ)
    end
end

#---------------------------------------------------------------------------# methods
StatsBase.nobs(w::Weight) = w.nobs
nups(w::StochasticWeight) = w.nups

# increase nobs by the number of new observations
updatecounter!(w::Weight, n2::Int = 1)              = (w.nobs += n2)
updatecounter!(w::StochasticWeight, n2::Int = 1)    = (w.nobs += n2; w.nups += 1)
updatecounter!(o::OnlineStat, n2::Int = 1) = updatecounter!(o.weight, n2)

# Get weight for update of size n2, typically immediately after updatecounter!
weight(w::EqualWeight, n2::Int = 1)         = n2 / w.nobs
weight(w::BoundedEqualWeight, n2::Int = 1)  = nobs(w) == 0 ? 1.0 : max(w.λ, n2 / w.nobs)
weight(w::ExponentialWeight, n2::Int = 1)   = w.λ
weight(w::LearningRate, n2::Int = 1)        = max(w.λ, exp(-w.r * log(w.nups)))
weight(w::LearningRate2, n2::Int = 1)       = max(w.λ, 1.0 / (1.0 + w.c * (w.nups - 1)))
weight(o::OnlineStat, n2::Int = 1) = weight(o.weight, n2)


using PkgSRA
using Plots, Random
using Flux, DiffEqFlux, OrdinaryDiffEq, Statistics
using DSP
pyplot()
Random.seed!(11)

#####
##### Define the controlled dynamical system
#####
# Load example problem
include("paper_settings.jl")
include(EXAMPLE_FOLDERNAME*"example_mass_spring.jl")

# Define the time dependent forcing function
# num_ctr = 2;
# U_starts = rand(num_ctr,1) .* tspan[2]
# U_width = 0.05;
# U_func(t) = U_func_time(t, u, U_width, U_starts)
U_func = spring_forcing_example(tspan)

# Solve the system
dat = Array(solve_msd_system(U_func))
dat_unforced = Array(solve_msd_system())
# Add a constant to the data
n, m = size(dat)
dat_and_const = vcat(dat, ones(1,m))

# Get true control signal
U_true = zeros(size(dat))
for (i, t) in enumerate(collect(ts))
    U_true[:,i] = U_func(t, dat[:,1])
end

#####
##### SRA: initialize the separation for dynamics vs. control
#####
dat_grad = numerical_derivative(dat, ts)

U_guess = sra_dmd(dat_and_const, dat_grad, quantile_threshold=0.8)[1]
A, B = dmdc(dat_and_const, dat_grad, U_guess)
# Loop a couple of times
for i in 1:20
    # Update guess for control signal
    global U_guess = calculate_residual(dat_and_const,
                                        dat_grad,
                                        A)
    sparsify_signal!(U_guess, quantile_threshold=0.8)
    # Update guess for dynamics
    global A, B = dmdc(dat_and_const, dat_grad, U_guess)
end

#####
##### Try to learn the dynamics directly
#####
nn_dim = 64
    U_dim = 1 # Expect that the true dimensionality is 1
core_dyn_naive = Dense(n, n, identity)
# Part 2: Neural network for control signal
ctr_dyn_naive = Chain(Dense(1,nn_dim, initb=(x)-> tspan[2].*rand(x)),
                        Dense(nn_dim, nn_dim, σ),
                        Dense(nn_dim, nn_dim, σ),
                        Dense(nn_dim, U_dim),
                        Dense(U_dim,2)) # Equivalent to the B matrix
p = param(Float32[1.0; 1.0])

function dudt_naive(u::AbstractArray,p,t)
    core_dyn_naive(u) + p.*ctr_dyn_naive([t])
end
prob_naive = ODEProblem(dudt_naive,u0,tspan,p)

_x = param(u0) # The initial condition can be learned too

function predict_rd_naive()
    Flux.Tracker.collect(diffeq_rd(p,prob_naive,Tsit5(),u0=_x, saveat=ts))
end
function loss_naive(λ=0.01)
    return sum(abs2,predict_rd_naive() .- dat) +      # Good fit
            λ*sum(abs,ctr_dyn_naive(collect(ts)'))   # L1 norm on ctr
end

# SOLVE the mixed NODE
ps = params(core_dyn_naive, ctr_dyn_naive, p);

# Initial quick learning
opt = ADAM(0.01)
dat_dummy = Iterators.repeated((), 300)
Flux.train!(loss_naive, ps,
            dat_dummy, opt, cb = ()->@show loss_naive())
# Run for a long time
opt = ADAM(0.001)
tol = 25
dat_dummy = Iterators.repeated((), 20)
num_iter = 1
max_iter = 10
while loss_naive() > tol
    Flux.train!(loss_naive, ps,
                dat_dummy, opt, cb = ()->@show loss_naive())
    global num_iter += 1
    num_iter > max_iter && break
end

#####
##### Plot and save: uninitialized
#####
# Full fit
sol = Flux.data(predict_rd_naive())
p_naive_full = plot(sol[2,:], label="Simulation",
                    color=COLOR_DICT["intrinsic"], lw=2);
    plot!(dat_and_const[2,:], label="Data",
                    color=COLOR_DICT["data"], lw=3);
    title!("Unitilialized Fit");

# Only "dynamics"
function dudt_naive_only_dyn(u::AbstractArray,p,t)
    core_dyn_naive(u)
end
prob_naive_only_dyn = ODEProblem(dudt_naive_only_dyn,u0,tspan,p)
function predict_rd_naive_dyn()
    Flux.Tracker.collect(diffeq_rd(p,prob_naive_only_dyn,Tsit5(),u0=_x, saveat=ts))
end
p_naive_dyn = plot(Flux.data(predict_rd_naive_dyn())[2,:], legend=false,
                    color=COLOR_DICT["intrinsic"], lw=2);
    title!("Intrinsic Dynamics");

# Only "control"
dat_naive_ctr = Flux.data(ctr_dyn_naive(collect(ts)'))
p_naive_ctr = plot(dat_naive_ctr[2,:], legend=false,
                    color=COLOR_DICT["control_time"], lw=2);
    title!("Control Signal");

# Save everything
my_layout = @layout [p_naive_full  [p_naive_dyn; p_naive_ctr]];
    p_all = plot(p_naive_full, p_naive_dyn, p_naive_ctr, layout = my_layout)
fname = FIGURE_FOLDERNAME * "fig_uninitialized.png";
savefig(p_all, fname)

#####
##### Instead: initialize the controller and then learn
#####
# SRA: Learn a neural network to approximate the guess
# Note: need to initialize the first layer so that the whole space is used
ctr_dyn = Chain(Dense(1,nn_dim, initb=(x)-> tspan[2].*rand(x)),
                Dense(nn_dim, nn_dim, σ),
                Dense(nn_dim, nn_dim, σ),
                Dense(nn_dim, U_dim),
                Dense(U_dim,2)) # Equivalent to the B matrix
ps = Flux.params(ctr_dyn)
    loss_U() = sum(abs2,ctr_dyn(collect(ts)') .- U_guess)
    loss_to_use = loss_U

# Initial fast learning
tol = 500
    dat_dummy = Iterators.repeated((), 100)
    opt = ADAM(0.01)
while loss_to_use() > 5*tol
    Flux.train!(loss_to_use, ps, dat_dummy, opt, cb = ()->@show loss_to_use())
end
opt = ADAM(0.001)
while loss_to_use() > tol
    Flux.train!(loss_to_use, ps, dat_dummy, opt, cb = ()->@show loss_to_use())
end
println("Finished learning control signal")
# plot(Flux.data(ctr_dyn(collect(ts)')[2,:]))

## Now learn entire system
core_dyn = Dense(param(Float32.(A[:, 1:n])),
                 param(Float32.(A[:, end])),
                 identity)
# Part 2: Neural network for control signal; initialized above
p = param(Float32[1.0; 1.0])

function dudt_(u::AbstractArray,p,t)
    core_dyn(u) + p.*ctr_dyn([t])
end
prob = ODEProblem(dudt_,u0,tspan,p)

_x = param(u0) # The initial condition can be learned too

function predict_rd()
    Flux.Tracker.collect(diffeq_rd(p,prob,Tsit5(),u0=_x, saveat=ts))
end
function loss_rd(λ1=0.01, λ2=0.1)
    return sum(abs2,predict_rd() .- dat) +      # Good fit
            λ1*loss_U() +                # Control signal close to guess
            λ2*sum(abs,ctr_dyn(collect(ts)'))   # L1 norm on ctr
end

## SOLVE the initialized mixed NODE
ps = params(core_dyn,ctr_dyn,p);

# Initial quick learning
opt = ADAM(0.0001)
dat_dummy = Iterators.repeated((), 10)
Flux.train!(loss_rd, ps,
            dat_dummy, opt, cb = ()->@show loss_rd())
# Run for a long time
opt = ADAM(0.00001)
tol = 50
dat_dummy = Iterators.repeated((), 10)
while loss_rd() > tol
    Flux.train!(loss_rd, ps,
                dat_dummy, opt, cb = ()->@show loss_rd())
    break
end

#####
##### Plot and save: Initialized
#####
# Full fit
sol = Flux.data(predict_rd())
p_full = plot(dat_and_const[2,:], label="Data",
                            color=COLOR_DICT["data"], lw=3);
    plot!(sol[2,:], label="Simulation",
                            color=COLOR_DICT["intrinsic"], lw=2);
    title!("Initialized Fit");

# Only "dynamics"
function dudt_only_dyn(u::AbstractArray,p,t)
    core_dyn(u)
end
prob_only_dyn = ODEProblem(dudt_only_dyn,u0,tspan,p)
function predict_rd_dyn()
    Flux.Tracker.collect(diffeq_rd(p,prob_only_dyn,Tsit5(),u0=_x, saveat=ts))
end
p_dyn = plot(dat_unforced[2,:], label="True",
                    color=COLOR_DICT["intrinsic_true"], lw=3)
        plot!(Flux.data(predict_rd_dyn())[2,:], label="Simulation",
                    color=COLOR_DICT["intrinsic"], lw=2);
    title!("Intrinsic Dynamics");

# Only "control"
dat_full_ctr = Flux.data(p.*ctr_dyn(collect(ts)'))
p_ctr = plot(U_true[2,:], label="True", lw=2,
            color=COLOR_DICT["control_true"])
        plot!(dat_full_ctr[2,:], label="Simulation",
                color=COLOR_DICT["control_time"], lw=2);
    title!("Control Signal");

# Save everything
my_layout = @layout [p_full  [p_dyn; p_ctr]];
    p_all = plot(p_full, p_dyn, p_ctr, layout = my_layout)
fname = FIGURE_FOLDERNAME * "fig_initialized.png";
savefig(p_all, fname)


include("../src/casingcalculations.jl")

@testset "Casing pressure segment drop" begin
    sg_gas = 0.7
    P_pc = HankinsonWithWichertPseudoCriticalPressure(sg_gas, 0, 0)
    _, T_pc, _ = HankinsonWithWichertPseudoCriticalTemp(sg_gas, 0, 0)
    ΔP = calculate_casing_pressuresegment(300, 10, (100+180)/2, 4500, sg_gas, KareemEtAlZFactor, P_pc, T_pc)
    @test ΔP ≈ 332.7-300 atol = 1
end

@testset "Casing pressure traverse" begin
    md = [0, 2250, 4500]
    tvd = md
    inc = [0, 0, 0]
    id = 2.441

    temps = [100., 140., 180.]

    testwell = Wellbore(md, inc, tvd, id)

    pressures = casing_traverse_topdown(wellbore = testwell, temperatureprofile = temps,
                                        CHP = 300 - pressure_atmospheric, sg_gas = 0.7, dp_est = 10)

    @test pressures[end] ≈ (332.7 - pressure_atmospheric) atol = 5

    #model = WellModel(wellbore = testwell, temperatureprofile = temps, CHP = 300 - pressure_atmospheric, sg_gas_inj = 0.7, dp_est_inj = 10)
    #pressures2 = casing_traverse_topdown(model)
    #@test pressures2[end] ≈ (332.7 - pressure_atmospheric) atol = 5
end


using DynamicalSystems, PyPlot

cd()
mkpath("hhanim")

ds = Systems.henonheiles()

# Grid of initial conditions at given energy:

energy(x,y,px,py) = 0.5(px^2 + py^2) + potential(x,y)
potential(x, y) = 0.5(x^2 + y^2) + (x^2*y - (y^3)/3)
function generate_ics(E, n)
    ys = range(-0.4, stop = 1.0, length = n)
    pys = range(-0.5, stop = 0.5, length = n)
    ics = Vector{Vector{Float64}}()
    for y in ys
        V = potential(0.0, y)
        V ≥ E && continue
        for py in pys
            Ky = 0.5*(py^2)
            Ky + V ≥ E && continue
            px = sqrt(2(E - V - Ky))
            ic = [0.0, y, px, py]
            push!(ics, [0.0, y, px, py])
        end
    end
    return ics
end

Emin = 0.1
Emax = 1/6
Es = range(Emin, stop = Emax, length = 100)

density = 20
tfinal = 2000.0
tgali = 1000.0
tinteg = tangent_integrator(ds, 4)

# Initialize required stuff for plotting
mpl = PyPlot.matplotlib
norm = mpl[:colors][:Normalize](0.0, 1.0)
c_m = mpl[:cm][:Spectral]
s_m = mpl[:cm][:ScalarMappable](cmap=c_m, norm=norm)
s_m[:set_array]([])

# %%
ioff()
for j in 1:length(Es)
    j==1 && (tim = time())
    E = Es[j]
    println("Energy $(j)/$(length(Es))")
    figure(figsize = (12, 8))
    ax = gca()

    ics = generate_ics(E, density)

    for u in ics

        # compute section:
        psos = poincaresos(ds, (1, 0.0), tfinal; u0 = u)

        # compute gali (using advanced usage)
        reinit!(tinteg, u, orthonormal(4,4))
        regularity = gali(tinteg, tgali, 1, 1e-12)[2][end]/tgali

        # Plot PSOS with color corresponding to regularity
        y = psos[:, 2]; py = psos[:, 4]
        plot(y, py, ls = "None", ms = 0.8, marker = "o",
        color = s_m[:to_rgba](regularity), alpha = 0.75)
    end
    cb = colorbar(s_m, ticks=[])
    # cb = colorbar(s_m, ticks=[lowerend, 1.0])
    # cb[:ax][:set_yticklabels](["chaotic", "periodic"], rotation="vertical")
    cb[:set_label]("regularity")
    xlim(-0.55, 1.05)
    ylim(-0.65, 0.65)
    xlabel("\$y\$")
    ylabel("\$p_y\$")
    text(0.7, 0.9, "\$E = $(round(E, digits=5))\$", transform = ax[:transAxes])
    tight_layout()
    savefig("hhanim/galipsos_$j.png", dpi = 200)
    if j == 1
        ttim = time()
        println("Average time per energy: $(round(ttim - tim, digits=3)) secs")
    end
end
close("all")
ion()

# %% produce animation

framerate = 10
savename = "hhanim/galipsos"
anim = `ffmpeg -y -framerate $(framerate) -start_number 1 -i $(savename)_%d.png
-c:v libx264 -pix_fmt yuv420p -preset veryslow -profile:v baseline
-filter:v scale=1200:-1 -b:v 2048k -level 3.0 $(savename).mp4`
run(anim)

if true
    for i in 1:length(Es)
        rm(savename*"_$(i).png")
    end
end


using GeometryTypes, Quaternions, GLVisualize, Reactive, GLAbstraction

function Quaternions.qrotation{T<:Real}(axis::Quaternion{T}, theta::T)
    ax = Vec{3,T}(axis.v1, axis.v2, axis.v3)
    qrotation(ax, theta)
end

function Quaternions.qrotation{T<:Real}(axis::Quaternion{T}, z::Quaternion{T}, theta::T)
    q = qrotation(axis, theta)
    q*z*conj(q)
end

function quaternion_belt_trick(timesignal)

    rad1 = 1f0 # outer circle
    rad2 = 0.2f0 # inner circle

    center1 = Point2f0(rad1,rad1) # left
    center2 = Point2f0(rad1+2f0,rad1) # right

    u1234 = Quaternion(0f0,0f0,1f0,0f0) # gets rotated
    axis  = Quaternion(0f0,0f0,0f0,1f0) # axis of qrotationdd
    mi  = 80
    mj  = 10
    rxs = zeros(Float32, mi+1,mj+1)
    rys = zeros(Float32, mi+1,mj+1)
    rzs = zeros(Float32, mi+1,mj+1)

    max_frames = 96
    dphi = 2*pi/(max_frames-1) # include both ends

    u1234_s = foldp(u1234, timesignal) do v0, _
        qrotation(axis, v0, Float32(dphi))
    end
    xyz = const_lift(plot_belts, rxs, rys, rzs, u1234_s, center1, center2, rad1, rad2)
    x = const_lift(getindex, xyz, 1)
    y = const_lift(getindex, xyz, 2)
    z = const_lift(getindex, xyz, 3)
    (x,y,z)
end

function plot_belts(rxs, rys, rzs, u1234, center1, center2, rad1, rad2)
    mi  = 80
    mj  = 10
    iis = 1f0:(mi+1f0)
    rs  = rad1-((rad1-rad2)*(iis-1f0)/mi)
    for ii=1:(mi+1)
        jjs   = 1f0:(mj+1f0)
        xs    = (rad2/mj)*((jjs-1f0)-(mj/2f0))
        ys    = sqrt(rs[ii]*rs[ii]-xs.*xs)
        theta = Float32(2*pi*(ii-1)/mi)

        for jj=1:(mj+1)
            efgh    = Quaternion(0f0, xs[jj], ys[jj], 0f0)
            wxyz    = qrotation(efgh, u1234, theta)
            rxs[ii,jj] = center1[1]+wxyz.v1
            rys[ii,jj] = center1[2]+wxyz.v2
            rzs[ii,jj] = wxyz.v3
        end
    end
    rxs, rys, rzs
end

if !isdefined(:runtests)
    window = glscreen()
    timesignal = loop(linspace(0f0, 1f0, 360))
end

xyz = quaternion_belt_trick(timesignal)
_view(visualize(xyz, :surface), window)

if !isdefined(:runtests)
    renderloop(window)
end


function powSpec(real::Number, imag::Number)
  return real^2+imag^2
end

function RMSpowSpec(real::Number, imag::Number)
  return sqrt(real^2+imag^2)
end

function logPowSpec(real::Number, imag::Number)
  return log(real^2+imag^2)
end


module TestHetero
using GaussianProcesses, ScikitLearnBase
using GaussianProcesses: set_params!, get_params, num_params,
                         update_cK!, update_mll!, update_target_and_dtarget!
using Test, Random

Random.seed!(1)
@testset "GPE heteroscedastic" begin
    d, n = 3, 10

    X = 2π * rand(d, n)
    y = [sum(sin, view(X, :, i)) / d for i in 1:n]
    mZero = MeanZero()
    kern = SE(0.0, 0.0)
    logNoise = range(0.0, stop=1.0, length=n)

    ntest = 5
    Xtest = randn(d, ntest)

    @testset "VectorParam" begin
        vp = GaussianProcesses.wrap_param(logNoise)
        @test all(get_params(vp) .== logNoise)
        @test num_params(vp) == n
        @test length(vp) == n
    end

    @testset "GPE constructors" begin
        gp1 = GPE(X, y, mZero, kern, logNoise)
        gp2 = GPE(X, y, mZero, kern, GaussianProcesses.wrap_param(logNoise))
    end

    gp = GP(X, y, mZero, kern, logNoise)

    @testset "update_cK!" begin
        update_cK!(gp)
    end
    @testset "update_mll" begin
        update_mll!(gp)
        @test isfinite(gp.mll)
    end

    @testset "update_dmll!" begin
        update_target_and_dtarget!(gp; noise=false)
        @test isfinite(gp.mll)
        @test isfinite(gp.target)
        @test all(isfinite.(gp.dmll))
        @test all(isfinite.(gp.dtarget))
    end

    @testset "num_params" begin
        @test num_params(gp) == 0+2+n
    end

    @testset "show" begin
        show(devnull, gp) # doesn't crash
    end

    @testset "Predictive mean at test locations" begin
        f_pred, sig = predict_f(gp, Xtest)
        f_pred, sig = predict_f(gp, Xtest; full_cov=true)
    end

    # ScikitLearn interface test
    @testset "ScikitLearn interface" begin
        gp_sk = ScikitLearnBase.fit!(GPE(), X', y)
        f_pred = ScikitLearnBase.predict(gp_sk, X')
    end

    # Modify kernel and update
    @testset "Update" begin
        gp.kernel.ℓ2 = 4.0
        X_pred = 2π * rand(d, n)

        GaussianProcesses.update_target!(gp)
        f_pred, sig = predict_f(gp, X_pred)
    end

    #Check that the rand function works
    @testset "Random GP sampling" begin
        X_test = 2π * rand(d, n)
        samples = rand(gp, X_test)
    end

    @testset "params round trip" begin
        params_1 = deepcopy(get_params(gp))
        set_params!(gp, params_1)
        params_2 = get_params(gp)
        @test params_1 ≈ params_2
    end

end
end