###### Weight vector #####
abstract type AbstractWeights{S<:Real, T<:Real, V<:AbstractVector{T}} <: AbstractVector{T} end

"""
    @weights name

Generates a new generic weight type with specified `name`, which subtypes `AbstractWeights`
and stores the `values` (`V<:RealVector`) and `sum` (`S<:Real`).
"""
macro weights(name)
    return quote
        mutable struct $name{S<:Real, T<:Real, V<:AbstractVector{T}} <: AbstractWeights{S, T, V}
            values::V
            sum::S
        end
        $(esc(name))(vs) = $(esc(name))(vs, sum(vs))
    end
end

eltype(wv::AbstractWeights) = eltype(wv.values)
length(wv::AbstractWeights) = length(wv.values)
values(wv::AbstractWeights) = wv.values
sum(wv::AbstractWeights) = wv.sum
isempty(wv::AbstractWeights) = isempty(wv.values)

Base.getindex(wv::AbstractWeights, i) = getindex(wv.values, i)
Base.size(wv::AbstractWeights) = size(wv.values)

@propagate_inbounds function Base.setindex!(wv::AbstractWeights, v::Real, i::Int)
    s = v - wv[i]
    wv.values[i] = v
    wv.sum += s
    v
end

"""
    varcorrection(n::Integer, corrected=false)

Compute a bias correction factor for calculating `var`, `std` and `cov` with
`n` observations. Returns ``\\frac{1}{n - 1}`` when `corrected=true`
(i.e. [Bessel's correction](https://en.wikipedia.org/wiki/Bessel's_correction)),
otherwise returns ``\\frac{1}{n}`` (i.e. no correction).
"""
@inline varcorrection(n::Integer, corrected::Bool=false) = 1 / (n - Int(corrected))

@weights Weights

@doc """
    Weights(vs, wsum=sum(vs))

Construct a `Weights` vector with weight values `vs`.
A precomputed sum may be provided as `wsum`.

The `Weights` type describes a generic weights vector which does not support
all operations possible for [`FrequencyWeights`](@ref), [`AnalyticWeights`](@ref)
and [`ProbabilityWeights`](@ref).
""" Weights

"""
    weights(vs)

Construct a `Weights` vector from array `vs`.
See the documentation for [`Weights`](@ref) for more details.
"""
weights(vs::RealVector) = Weights(vs)
weights(vs::RealArray) = Weights(vec(vs))

"""
    varcorrection(w::Weights, corrected=false)

Returns ``\\frac{1}{\\sum w}`` when `corrected=false` and throws an `ArgumentError`
if `corrected=true`.
"""
@inline function varcorrection(w::Weights, corrected::Bool=false)
    corrected && throw(ArgumentError("Weights type does not support bias correction: " *
                                     "use FrequencyWeights, AnalyticWeights or ProbabilityWeights if applicable."))
    1 / w.sum
end

@weights AnalyticWeights

@doc """
    AnalyticWeights(vs, wsum=sum(vs))

Construct an `AnalyticWeights` vector with weight values `vs`.
A precomputed sum may be provided as `wsum`.

Analytic weights describe a non-random relative importance (usually between 0 and 1)
for each observation. These weights may also be referred to as reliability weights,
precision weights or inverse variance weights. These are typically used when the observations
being weighted are aggregate values (e.g., averages) with differing variances.
""" AnalyticWeights

"""
    aweights(vs)

Construct an `AnalyticWeights` vector from array `vs`.
See the documentation for [`AnalyticWeights`](@ref) for more details.
"""
aweights(vs::RealVector) = AnalyticWeights(vs)
aweights(vs::RealArray) = AnalyticWeights(vec(vs))

"""
    varcorrection(w::AnalyticWeights, corrected=false)

* `corrected=true`: ``\\frac{1}{\\sum w - \\sum {w^2} / \\sum w}``
* `corrected=false`: ``\\frac{1}{\\sum w}``
"""
@inline function varcorrection(w::AnalyticWeights, corrected::Bool=false)
    s = w.sum

    if corrected
        sum_sn = sum(x -> (x / s) ^ 2, w)
        1 / (s * (1 - sum_sn))
    else
        1 / s
    end
end

@weights FrequencyWeights

@doc """
    FrequencyWeights(vs, wsum=sum(vs))

Construct a `FrequencyWeights` vector with weight values `vs`.
A precomputed sum may be provided as `wsum`.

Frequency weights describe the number of times (or frequency) each observation
was observed. These weights may also be referred to as case weights or repeat weights.
""" FrequencyWeights

"""
    fweights(vs)

Construct a `FrequencyWeights` vector from a given array.
See the documentation for [`FrequencyWeights`](@ref) for more details.
"""
fweights(vs::RealVector) = FrequencyWeights(vs)
fweights(vs::RealArray) = FrequencyWeights(vec(vs))

"""
    varcorrection(w::FrequencyWeights, corrected=false)

* `corrected=true`: ``\\frac{1}{\\sum{w} - 1}``
* `corrected=false`: ``\\frac{1}{\\sum w}``
"""
@inline function varcorrection(w::FrequencyWeights, corrected::Bool=false)
    s = w.sum

    if corrected
        1 / (s - 1)
    else
        1 / s
    end
end

@weights ProbabilityWeights

@doc """
    ProbabilityWeights(vs, wsum=sum(vs))

Construct a `ProbabilityWeights` vector with weight values `vs`.
A precomputed sum may be provided as `wsum`.

Probability weights represent the inverse of the sampling probability for each observation,
providing a correction mechanism for under- or over-sampling certain population groups.
These weights may also be referred to as sampling weights.
""" ProbabilityWeights

"""
    pweights(vs)

Construct a `ProbabilityWeights` vector from a given array.
See the documentation for [`ProbabilityWeights`](@ref) for more details.
"""
pweights(vs::RealVector) = ProbabilityWeights(vs)
pweights(vs::RealArray) = ProbabilityWeights(vec(vs))

"""
    varcorrection(w::ProbabilityWeights, corrected=false)

* `corrected=true`: ``\\frac{n}{(n - 1) \\sum w}`` where ``n`` equals `count(!iszero, w)`
* `corrected=false`: ``\\frac{1}{\\sum w}``
"""
@inline function varcorrection(w::ProbabilityWeights, corrected::Bool=false)
    s = w.sum

    if corrected
        n = count(!iszero, w)
        n / (s * (n - 1))
    else
        1 / s
    end
end

##### Equality tests #####

for w in (AnalyticWeights, FrequencyWeights, ProbabilityWeights, Weights)
    @eval begin
        Base.isequal(x::$w, y::$w) = isequal(x.sum, y.sum) && isequal(x.values, y.values)
        Base.:(==)(x::$w, y::$w)   = (x.sum == y.sum) && (x.values == y.values)
    end
end

Base.isequal(x::AbstractWeights, y::AbstractWeights) = false
Base.:(==)(x::AbstractWeights, y::AbstractWeights)   = false

##### Weighted sum #####

## weighted sum over vectors

"""
    wsum(v, w::AbstractVector, [dim])

Compute the weighted sum of an array `v` with weights `w`, optionally over the dimension `dim`.
"""
wsum(v::AbstractVector, w::AbstractVector) = dot(v, w)
wsum(v::AbstractArray, w::AbstractVector) = dot(vec(v), w)

# Note: the methods for BitArray and SparseMatrixCSC are to avoid ambiguities
Base.sum(v::BitArray, w::AbstractWeights) = wsum(v, values(w))
Base.sum(v::SparseArrays.SparseMatrixCSC, w::AbstractWeights) = wsum(v, values(w))
Base.sum(v::AbstractArray, w::AbstractWeights) = dot(v, values(w))

## wsum along dimension
#
#  Brief explanation of the algorithm:
#  ------------------------------------
#
#  1. _wsum! provides the core implementation, which assumes that
#     the dimensions of all input arguments are consistent, and no
#     dimension checking is performed therein.
#
#     wsum and wsum! perform argument checking and call _wsum!
#     internally.
#
#  2. _wsum! adopt a Cartesian based implementation for general
#     sub types of AbstractArray. Particularly, a faster routine
#     that keeps a local accumulator will be used when dim = 1.
#
#     The internal function that implements this is _wsum_general!
#
#  3. _wsum! is specialized for following cases:
#     (a) A is a vector: we invoke the vector version wsum above.
#         The internal function that implements this is _wsum1!
#
#     (b) A is a dense matrix with eltype <: BlasReal: we call gemv!
#         The internal function that implements this is _wsum2_blas!
#
#     (c) A is a contiguous array with eltype <: BlasReal:
#         dim == 1: treat A like a matrix of size (d1, d2 x ... x dN)
#         dim == N: treat A like a matrix of size (d1 x ... x d(N-1), dN)
#         otherwise: decompose A into multiple pages, and apply _wsum2!
#         for each
#
#     (d) A is a general dense array with eltype <: BlasReal:
#         dim <= 2: delegate to (a) and (b)
#         otherwise, decompose A into multiple pages
#

function _wsum1!(R::AbstractArray, A::AbstractVector, w::AbstractVector, init::Bool)
    r = wsum(A, w)
    if init
        R[1] = r
    else
        R[1] += r
    end
    return R
end

function _wsum2_blas!(R::StridedVector{T}, A::StridedMatrix{T}, w::StridedVector{T}, dim::Int, init::Bool) where T<:BlasReal
    beta = ifelse(init, zero(T), one(T))
    trans = dim == 1 ? 'T' : 'N'
    BLAS.gemv!(trans, one(T), A, w, beta, R)
    return R
end

function _wsumN!(R::StridedArray{T}, A::StridedArray{T,N}, w::StridedVector{T}, dim::Int, init::Bool) where {T<:BlasReal,N}
    if dim == 1
        m = size(A, 1)
        n = div(length(A), m)
        _wsum2_blas!(view(R,:), reshape(A, (m, n)), w, 1, init)
    elseif dim == N
        n = size(A, N)
        m = div(length(A), n)
        _wsum2_blas!(view(R,:), reshape(A, (m, n)), w, 2, init)
    else # 1 < dim < N
        m = 1
        for i = 1:dim-1; m *= size(A, i); end
        n = size(A, dim)
        k = 1
        for i = dim+1:N; k *= size(A, i); end
        Av = reshape(A, (m, n, k))
        Rv = reshape(R, (m, k))
        for i = 1:k
            _wsum2_blas!(view(Rv,:,i), view(Av,:,:,i), w, 2, init)
        end
    end
    return R
end

function _wsumN!(R::StridedArray{T}, A::DenseArray{T,N}, w::StridedVector{T}, dim::Int, init::Bool) where {T<:BlasReal,N}
    @assert N >= 3
    if dim <= 2
        m = size(A, 1)
        n = size(A, 2)
        npages = 1
        for i = 3:N
            npages *= size(A, i)
        end
        rlen = ifelse(dim == 1, n, m)
        Rv = reshape(R, (rlen, npages))
        for i = 1:npages
            _wsum2_blas!(view(Rv,:,i), view(A,:,:,i), w, dim, init)
        end
    else
        _wsum_general!(R, identity, A, w, dim, init)
    end
    return R
end

# General Cartesian-based weighted sum across dimensions
@generated function _wsum_general!(R::AbstractArray{RT}, f::supertype(typeof(abs)),
                                   A::AbstractArray{T,N}, w::AbstractVector{WT}, dim::Int, init::Bool) where {T,RT,WT,N}
    quote
        init && fill!(R, zero(RT))
        wi = zero(WT)
        if dim == 1
            @nextract $N sizeR d->size(R,d)
            sizA1 = size(A, 1)
            @nloops $N i d->(d>1 ? (1:size(A,d)) : (1:1)) d->(j_d = sizeR_d==1 ? 1 : i_d) begin
                @inbounds r = (@nref $N R j)
                for i_1 = 1:sizA1
                    @inbounds r += f(@nref $N A i) * w[i_1]
                end
                @inbounds (@nref $N R j) = r
            end
        else
            @nloops $N i A d->(if d == dim
                                   wi = w[i_d]
                                   j_d = 1
                               else
                                   j_d = i_d
                               end) @inbounds (@nref $N R j) += f(@nref $N A i) * wi
        end
        return R
    end
end

@generated function _wsum_centralize!(R::AbstractArray{RT}, f::supertype(typeof(abs)),
                                      A::AbstractArray{T,N}, w::AbstractVector{WT}, means,
                                      dim::Int, init::Bool) where {T,RT,WT,N}
    quote
        init && fill!(R, zero(RT))
        wi = zero(WT)
        if dim == 1
            @nextract $N sizeR d->size(R,d)
            sizA1 = size(A, 1)
            @nloops $N i d->(d>1 ? (1:size(A,d)) : (1:1)) d->(j_d = sizeR_d==1 ? 1 : i_d) begin
                @inbounds r = (@nref $N R j)
                @inbounds m = (@nref $N means j)
                for i_1 = 1:sizA1
                    @inbounds r += f((@nref $N A i) - m) * w[i_1]
                end
                @inbounds (@nref $N R j) = r
            end
        else
            @nloops $N i A d->(if d == dim
                                   wi = w[i_d]
                                   j_d = 1
                               else
                                   j_d = i_d
                               end) @inbounds (@nref $N R j) += f((@nref $N A i) - (@nref $N means j)) * wi
        end
        return R
    end
end

# N = 1
_wsum!(R::StridedArray{T}, A::DenseArray{T,1}, w::StridedVector{T}, dim::Int, init::Bool) where {T<:BlasReal} =
    _wsum1!(R, A, w, init)

# N = 2
_wsum!(R::StridedArray{T}, A::DenseArray{T,2}, w::StridedVector{T}, dim::Int, init::Bool) where {T<:BlasReal} =
    (_wsum2_blas!(view(R,:), A, w, dim, init); R)

# N >= 3
_wsum!(R::StridedArray{T}, A::DenseArray{T,N}, w::StridedVector{T}, dim::Int, init::Bool) where {T<:BlasReal,N} =
    _wsumN!(R, A, w, dim, init)

_wsum!(R::AbstractArray, A::AbstractArray, w::AbstractVector, dim::Int, init::Bool) =
    _wsum_general!(R, identity, A, w, dim, init)

## wsum! and wsum

wsumtype(::Type{T}, ::Type{W}) where {T,W} = typeof(zero(T) * zero(W) + zero(T) * zero(W))
wsumtype(::Type{T}, ::Type{T}) where {T<:BlasReal} = T

"""
    wsum!(R, A, w, dim; init=true)

Compute the weighted sum of `A` with weights `w` over the dimension `dim` and store
the result in `R`. If `init=false`, the sum is added to `R` rather than starting
from zero.
"""
function wsum!(R::AbstractArray, A::AbstractArray{T,N}, w::AbstractVector, dim::Int; init::Bool=true) where {T,N}
    1 <= dim <= N || error("dim should be within [1, $N]")
    ndims(R) <= N || error("ndims(R) should not exceed $N")
    length(w) == size(A,dim) || throw(DimensionMismatch("Inconsistent array dimension."))
    # TODO: more careful examination of R's size
    _wsum!(R, A, w, dim, init)
end

function wsum(A::AbstractArray{T}, w::AbstractVector{W}, dim::Int) where {T<:Number,W<:Real}
    length(w) == size(A,dim) || throw(DimensionMismatch("Inconsistent array dimension."))
    _wsum!(similar(A, wsumtype(T,W), Base.reduced_indices(axes(A), dim)), A, w, dim, true)
end

# extended sum! and wsum

Base.sum!(R::AbstractArray, A::AbstractArray, w::AbstractWeights{<:Real}, dim::Int; init::Bool=true) =
    wsum!(R, A, values(w), dim; init=init)

Base.sum(A::AbstractArray{<:Number}, w::AbstractWeights{<:Real}, dim::Int) = wsum(A, values(w), dim)

###### Weighted means #####

"""
    wmean(v, w::AbstractVector)

Compute the weighted mean of an array `v` with weights `w`.
"""
function wmean(v::AbstractArray{<:Number}, w::AbstractVector)
    Base.depwarn("wmean is deprecated, use mean(v, weights(w)) instead.", :wmean)
    mean(v, weights(w))
end

"""
    mean(A::AbstractArray, w::AbstractWeights[, dim::Int])

Compute the weighted mean of array `A` with weight vector `w`
(of type `AbstractWeights`). If `dim` is provided, compute the
weighted mean along dimension `dim`.

# Examples
```julia
w = rand(n)
mean(x, weights(w))
```
"""
mean(A::AbstractArray, w::AbstractWeights) = sum(A, w) / sum(w)

"""
    mean(R::AbstractArray, , A::AbstractArray, w::AbstractWeights[, dim::Int])

Compute the weighted mean of array `A` with weight vector `w`
(of type `AbstractWeights`) along dimension `dim`, and write results to `R`.
"""
mean!(R::AbstractArray, A::AbstractArray, w::AbstractWeights, dim::Int) =
    rmul!(Base.sum!(R, A, w, dim), inv(sum(w)))

wmeantype(::Type{T}, ::Type{W}) where {T,W} = typeof((zero(T)*zero(W) + zero(T)*zero(W)) / one(W))
wmeantype(::Type{T}, ::Type{T}) where {T<:BlasReal} = T

mean(A::AbstractArray{T}, w::AbstractWeights{W}, dim::Int) where {T<:Number,W<:Real} =
    mean!(similar(A, wmeantype(T, W), Base.reduced_indices(axes(A), dim)), A, w, dim)

###### Weighted median #####
function median(v::AbstractArray, w::AbstractWeights)
    throw(MethodError(median, (v, w)))
end

"""
    median(v::RealVector, w::AbstractWeights)

Compute the weighted median of `x`, using weights given by a weight vector `w`
(of type `AbstractWeights`). The weight and data vectors must have the same length.

The weighted median ``x_k`` is the element of `x` that satisfies
``\\sum_{x_i < x_k} w_i \\le \\frac{1}{2} \\sum_{j} w_j`` and
``\\sum_{x_i > x_k} w_i \\le \\frac{1}{2} \\sum_{j} w_j``.

If a weight has value zero, then its associated data point is ignored.
If none of the weights are positive, an error is thrown.
`NaN` is returned if `x` contains any `NaN` values.
An error is raised if `w` contains any `NaN` values.
"""
function median(v::RealVector, w::AbstractWeights{<:Real})
    isempty(v) && error("median of an empty array is undefined")
    if length(v) != length(w)
        error("data and weight vectors must be the same size")
    end
    @inbounds for x in w.values
        isnan(x) && error("weight vector cannot contain NaN entries")
    end
    @inbounds for x in v
        isnan(x) && return x
    end
    mask = w.values .!= 0
    if !any(mask)
        error("all weights are zero")
    end
    if all(w.values .<= 0)
        error("no positive weights found")
    end
    v = v[mask]
    wt = w[mask]
    midpoint = w.sum / 2
    maxval, maxind = findmax(wt)
    if maxval > midpoint
        v[maxind]
    else
        permute = sortperm(v)
        cumulative_weight = zero(eltype(wt))
        i = 0
        for (_i, p) in enumerate(permute)
            i = _i
            if cumulative_weight == midpoint
                i += 1
                break
            elseif cumulative_weight > midpoint
                cumulative_weight -= wt[p]
                break
            end
            cumulative_weight += wt[p]
        end
        if cumulative_weight == midpoint
            middle(v[permute[i-2]], v[permute[i-1]])
        else
            middle(v[permute[i-1]])
        end
    end
end

"""
    wmedian(v, w)

Compute the weighted median of an array `v` with weights `w`, given as either a
vector or an `AbstractWeights` vector.
"""
wmedian(v::RealVector, w::RealVector) = median(v, weights(w))
wmedian(v::RealVector, w::AbstractWeights{<:Real}) = median(v, w)

###### Weighted quantile #####

"""
    quantile(v, w::AbstractWeights, p)

Compute the weighted quantiles of a vector `v` at a specified set of probability
values `p`, using weights given by a weight vector `w` (of type `AbstractWeights`).
Weights must not be negative. The weights and data vectors must have the same length.

With [`FrequencyWeights`](@ref), the function returns the same result as
`quantile` for a vector with repeated values.
With non `FrequencyWeights`,  denote ``N`` the length of the vector, ``w`` the vector of weights,
``h = p (\\sum_{i<= N}w_i - w_1) + w_1`` the cumulative weight corresponding to the
probability ``p`` and ``S_k = \\sum_{i<=k}w_i`` the cumulative weight for each
observation, define ``v_{k+1}`` the smallest element of `v` such that ``S_{k+1}``
is strictly superior to ``h``. The weighted ``p`` quantile is given by ``v_k + \\gamma (v_{k+1} -v_k)``
with  ``\\gamma = (h - S_k)/(S_{k+1}-S_k)``. In particular, when `w` is a vector
of ones, the function returns the same result as `quantile`.
"""
function quantile(v::RealVector{V}, w::AbstractWeights{W}, p::RealVector) where {V,W<:Real}
    # checks
    isempty(v) && error("quantile of an empty array is undefined")
    isempty(p) && throw(ArgumentError("empty quantile array"))

    w.sum == 0 && error("weight vector cannot sum to zero")
    length(v) == length(w) || error("data and weight vectors must be the same size, got $(length(v)) and $(length(w))")
    for x in w.values
        isnan(x) && error("weight vector cannot contain NaN entries")
        x < 0 && error("weight vector cannot contain negative entries")
    end

    # remove zeros weights and sort
    wsum = sum(w)
    nz = .!iszero.(w)
    vw = sort!(collect(zip(view(v, nz), view(w, nz))))
    N = length(vw)

    # prepare percentiles
    ppermute = sortperm(p)
    p = p[ppermute]
    p = bound_quantiles(p)

    # prepare out vector
    out = Vector{typeof(zero(V)/1)}(undef, length(p))
    fill!(out, vw[end][1])

    # start looping on quantiles
    Sk, Skold = zero(W), zero(W)
    vk, vkold = zero(V), zero(V)
    k = 0

    for i in 1:length(p)
        if isa(w, FrequencyWeights)
            h = p[i] * (wsum - 1) + 1
        else
            h = p[i] * (wsum - vw[1][2]) + vw[1][2]
        end
        while Sk <= h
            k += 1
            if k > N
               # out was initialized with maximum v
               return out
            end
            Skold, vkold = Sk, vk
            vk, wk = vw[k]
            Sk += wk
        end
        if isa(w, FrequencyWeights)
            out[ppermute[i]] = vkold + min(h - Skold, 1) * (vk - vkold)
        else
            out[ppermute[i]] = vkold + (h - Skold) / (Sk - Skold) * (vk - vkold)
        end
    end
    return out
end

# similarly to statistics.jl in Base
function bound_quantiles(qs::AbstractVector{T}) where T<:Real
    epsilon = 100 * eps()
    if (any(qs .< -epsilon) || any(qs .> 1+epsilon))
        throw(ArgumentError("quantiles out of [0,1] range"))
    end
    T[min(one(T), max(zero(T), q)) for q = qs]
end

quantile(v::RealVector, w::AbstractWeights{<:Real}, p::Number) = quantile(v, w, [p])[1]

"""
    wquantile(v, w, p)

Compute the `p`th quantile(s) of `v` with weights `w`, given as either a vector
or an `AbstractWeights` vector.
"""
wquantile(v::RealVector, w::AbstractWeights{<:Real}, p::RealVector) = quantile(v, w, p)
wquantile(v::RealVector, w::AbstractWeights{<:Real}, p::Number) = quantile(v, w, [p])[1]
wquantile(v::RealVector, w::RealVector, p::RealVector) = quantile(v, weights(w), p)
wquantile(v::RealVector, w::RealVector, p::Number) = quantile(v, weights(w), [p])[1]


module effective_interaction_simple

using ..CollectiveSpins

const e_z = [0,0,1]
const gamma = 1

function effective_interactions(S::SpinCollection{T,P,G}) where {T,P,G}
    omega_eff::float(G) = 0.
    gamma_eff::float(G) = 0.
    origin = zeros(3)
    for spin = S.spins
        omega_eff += interaction.Omega(origin, spin.position, S.polarizations[1], S.polarizations[1], gamma, gamma)
        gamma_eff += interaction.Gamma(origin, spin.position, S.polarizations[1], S.polarizations[1], gamma, gamma)
    end
    return omega_eff, gamma_eff
end

# Finite symmetric systems

function triangle_orthogonal(a::T) where T<:Real
    positions = Vector{float(T)}[[a,0,0], [a/2, sqrt(3. /4)*a,0]]
    S = SpinCollection(positions, e_z; gammas=gamma)
    return effective_interactions(S)
end

function square_orthogonal(a::T) where T<:Real
    positions = Vector{float(T)}[[a,0,0], [0,a,0], [a,a,0]]
    S = SpinCollection(positions, e_z; gammas=gamma)
    return effective_interactions(S)
end

function polygon_orthogonal(N::Int, a::T) where T<:Real
    @assert N>2
    dα = 2*pi/N
    R = a/(2*sin(dα/2))
    positions = Vector{float(T)}[]
    for i=1:(N-1)
        x = R*cos(i*dα)
        y = R*sin(i*dα)
        push!(positions, float(T)[x-R,y,0])
    end
    S = SpinCollection(positions, e_z; gammas=gamma)
    return effective_interactions(S)
end

function rectangle_orthogonal(a::U, b::T) where {U<:Real,T<:Real}
    positions = Vector{float(T)}[[a,0,0], [0,b,0], [a,b,0]]
    S = SpinCollection(positions, e_z; gammas=gamma)
    return effective_interactions(S)
end

function cube_orthogonal(a::T) where T<:Real
    positions = Vector{float(T)}[]
    for ix=0:1, iy=0:1, iz=0:1
        if ix==0 && iy==0 && iz==0
            continue
        end
        push!(positions, [ix*a, iy*a, iz*a])
    end
    S = SpinCollection(positions, e_z; gammas=gamma)
    return effective_interactions(S)
end

function box_orthogonal(a::T, b::U, c::V) where {T<:Real,U<:Real,V<:Real}
   positions = Vector{float(T)}[]
    for ix=0:1, iy=0:1, iz=0:1
        if ix==0 && iy==0 && iz==0
            continue
        end
        push!(positions, [ix*a, iy*b, iz*c])
    end
    S = SpinCollection(positions, e_z; gammas=gamma)
    return effective_interactions(S)
end

# Infinite 1D symmetric systems

function chain(a::T, Θ, N::Int) where T<:Real
    positions = Vector{float(T)}[]
    for ix=-N:N
        if ix==0
            continue
        end
        push!(positions, [ix*a, 0., 0.])
    end
    S = SpinCollection(positions, Float64[cos(Θ), 0., sin(Θ)]; gammas=gamma)
    return effective_interactions(S)
end

function chain_orthogonal(a::T, N::Int) where T<:Real
    positions = Vector{float(T)}[]
    for ix=-N:N
        if ix==0
            continue
        end
        push!(positions, [ix*a, 0., 0.])
    end
    S = SpinCollection(positions, e_z; gammas=gamma)
    return effective_interactions(S)
end

# Infinite 2D symmetric systems

function squarelattice_orthogonal(a::T, N::Int) where {T<:Real}
    positions = Vector{float(T)}[]
    for ix=-N:N, iy=-N:N
        if ix==0 && iy==0
            continue
        end
        push!(positions, [ix*a, iy*a, 0.])
    end
    S = SpinCollection(positions, e_z; gammas=gamma)
    return effective_interactions(S)
end

function hexagonallattice_orthogonal(a::T, N::Int) where {T<:Real}
    positions = Vector{float(T)}[]
    ax = sqrt(3.0/4)*a
    for iy=1:N
        push!(positions, [0, iy*a, 0])
        push!(positions, [0, -iy*a, 0])
    end
    for ix=[-1:-2:-N; 1:2:N]
        Ny = div(2*N+1-abs(ix),2)
        for iy=0:Ny-1
            push!(positions, [ax*ix, (0.5+iy)*a, 0])
            push!(positions, [ax*ix, -(0.5+iy)*a, 0])
        end
    end
    for ix=[-2:-2:-N; 2:2:N]
        Ny = div(2*N-abs(ix),2)
        push!(positions, [ax*ix, 0, 0])
        for iy=1:Ny
            push!(positions, [ax*ix, iy*a, 0])
            push!(positions, [ax*ix, -iy*a, 0])
        end
    end
    S = SpinCollection(positions, e_z; gammas=gamma)
    return effective_interactions(S)
end

# Infinite 3D symmetric systems

function cubiclattice_orthogonal(a::T, N::Int) where {T<:Real}
    positions = Vector{float(T)}[]
    for ix=-N:N, iy=-N:N, iz=-N:N
        if ix==0 && iy==0 && iz==0
            continue
        end
        push!(positions, [ix*a, iy*a, iz*a])
    end
    S = SpinCollection(positions, e_z; gammas=gamma)
    return effective_interactions(S)
end

function tetragonallattice_orthogonal(a::T, b::U, N::Int) where {T<:Real,U<:Real}
    positions = Vector{float(T)}[]
    for ix=-N:N, iy=-N:N, iz=-N:N
        if ix==0 && iy==0 && iz==0
            continue
        end
        push!(positions, [ix*a, iy*a, iz*b])
    end
    S = SpinCollection(positions, e_z; gammas=gamma)
    return effective_interactions(S)
end

function hexagonallattice3d_orthogonal(a::T, b::U, N::Int) where {T<:Real,U<:Real}
    positions = Vector{float(T)}[]
    ax = sqrt(3.0/4)*a
    for iz=-N:N
        for iy=1:N
            push!(positions, [0, iy*a, iz*b])
            push!(positions, [0, -iy*a, iz*b])
        end
        for ix=[-1:-2:-N; 1:2:N]
            Ny = div(2*N+1-abs(ix),2)
            for iy=0:Ny-1
                push!(positions, [ax*ix, (0.5+iy)*a, iz*b])
                push!(positions, [ax*ix, -(0.5+iy)*a, iz*b])
            end
        end
        for ix=[-2:-2:-N; 2:2:N]
            Ny = div(2*N-abs(ix),2)
            push!(positions, [ax*ix, 0, iz*b])
            for iy=1:Ny
                push!(positions, [ax*ix, iy*a, iz*b])
                push!(positions, [ax*ix, -iy*a, iz*b])
            end
        end
    end
    S = SpinCollection(positions, e_z; gammas=gamma)
    return effective_interactions(S)
end

end # module


function binarysearch(lst::Vector{T}, value::T, low=1, high=length(lst)) where T
    if isempty(lst) return 0 end
    if low ≥ high
        if low > high || lst[low] != value
            return 0
        else
            return low
        end
    end
    mid = (low + high) ÷ 2
    if lst[mid] > value
        return binarysearch(lst, value, low, mid-1)
    elseif lst[mid] < value
        return binarysearch(lst, value, mid+1, high)
    else
        return mid
    end
end


# STFT/ISTFT

function blackman(n::Integer)
    const a0, a1, a2 = 0.42, 0.5, 0.08
    t = 2*pi/(n-1)
    [a0 - a1*cos(t*k) + a2*cos(t*k*2) for k=0:n-1]
end

function hanning(n::Integer)
    [0.5*(1-cos(2*pi*k/(n-1))) for k=0:n-1]
end

# countframes returns the number of frames that will be processed.
function countframes{T<:Number}(x::Vector{T}, framelen::Int, hopsize::Int)
    div(length(x) - framelen, hopsize) + 1
end

# splitframes performs overlapping frame splitting.
function splitframes{T<:Number}(x::Vector{T},
                                framelen::Int=1024,
                                hopsize::Int=framelen/2)
    const N = countframes(x, framelen, hopsize)
    frames = Array(eltype(x), framelen, N)

    for i=1:N
        frames[:,i] = x[(i-1)*hopsize+1:(i-1)*hopsize+framelen]
    end

    return frames
end

# stft performs the Short-Time Fourier Transform (STFT) for real signals.
function stft{T<:Real}(x::Vector{T},
                       framelen::Int=1024,
                       hopsize::Int=div(framelen,2),
                       window=hanning(framelen))
    frames = splitframes(x, framelen, hopsize)

    const freqbins = div(framelen, 2) + 1
    spectrogram = Array(Complex64, freqbins, size(frames,2))
    for i=1:size(frames,2)
        spectrogram[:,i] = rfft(frames[:,i] .* window)
    end

    return spectrogram
end

# istft peforms the Inverse STFT to recover the original signal from STFT
# coefficients.
function istft{T<:Complex}(spectrogram::Matrix{T},
                           framelen::Int=1024,
                           hopsize::Int=div(framelen,2),
                           window=hanning(framelen))
    const numframes = size(spectrogram, 2)

    expectedlen = framelen + (numframes-1)*hopsize
    reconstructed = zeros(expectedlen)
    windowsum = zeros(expectedlen)
    const windowsquare = window .* window

    # Overlapping addition
    for i=1:numframes
        s, e = (i-1)*hopsize+1, (i-1)*hopsize+framelen
        r = irfft(spectrogram[:,i], framelen)
        reconstructed[s:e] += r .* window
        windowsum[s:e] += windowsquare
    end

    # Normalized by window
    for i=1:endof(reconstructed)
        # avoid zero division
        if windowsum[i] > 1.0e-7
            reconstructed[i] /= windowsum[i]
        end
    end

    return reconstructed
end


function Binary_Search(x,key)     #declared function
mid = Int(floor((1 + length(x))/2))  #calculating the mid of array
           try
           if x[mid] == key
           println("Key found")        # if key is found
           elseif x[mid]< key           #comparision to see if right side of key is to iterated
           Binary_Search(x[mid+1:end],key) #recursive call
           else x[mid]>key            #comparision to see if left side of key is to iterated
           Binary_Search(x[1:mid-1],key)  #recursive call
           end
           catch
           println("Key: $key not found") # else key is not found
           end
end

x = [1,4,6,8,9,19,56]
key = 10
Binary_Search(x,key)


# # Bessel function
# An Bessel function of the first kind of order ``\nu`` can be computed using Taylor expansion

# ```math
#     J_\nu(z) = \sum\limits_{n=0}^{\infty} \frac{(z/2)^\nu}{\Gamma(k+1)\Gamma(k+\nu+1)} (-z^2/4)^{n}
# ```

# where ``\Gamma(n) = (n-1)!`` is the Gamma function. One can compute the accumulated item iteratively as ``s_n = -\frac{z^2}{4} s_{n-1}``.

using NiLang, NiLang.AD
using ForwardDiff: Dual

# Since we need to use logarithmic numbers to handle the sequential mutiplication.
# Let's first add patch about the conversion between `ULogarithmic` and `Dual` number.
function Base.convert(::Type{Dual{T,V,N}}, x::ULogarithmic) where {T,V,N}
	Dual{T,V,N}(exp(x.log))
end

@i function ibesselj(y!::T, ν, z::T; atol=1e-8) where T
	if z == 0
		if v == 0
			out! += 1
		end
	else
		@routine @invcheckoff begin
			k ← 0
			@ones ULogarithmic{T} lz halfz halfz_power_2 s
			@zeros T out_anc
			lz *= convert(z)
			halfz *= lz / 2
			halfz_power_2 *= halfz ^ 2
			## s *= (z/2)^ν/ factorial(ν)
			s *= halfz ^ ν
			for i=1:ν
				s /= i
			end
			out_anc += convert(s)
			@from k==0 while s.log > -25 # upto precision e^-25
				k += 1
				## s *= 1 / k / (k+ν) * (z/2)^2
				s *= halfz_power_2 / (@const k*(k+ν))
				if k%2 == 0
					out_anc += convert(s)
				else
					out_anc -= convert(s)
				end
			end
		end
		y! += out_anc
		~@routine
	end
end

# To obtain gradients, one call **Grad(ibesselj)**

y, x = 0.0, 1.0
Grad(ibesselj)(Val(1), y, 2, x)

# Here, **Grad(ibesselj)** is a callable instance of type **Grad{typeof(ibesselj)}}**.
# The first parameter `Val(1)` indicates the first argument is the loss.

# To obtain second order gradients, one can Feed dual numbers to this gradient function.
_, hxy, _, hxx = Grad(ibesselj)(Val(1), Dual(y, zero(y)), 2, Dual(x, one(x)))
println("The hessian dy^2/dx^2 is $(grad(hxx).partials[1])")

# Here, the gradient field is a Dual number, it has a field partials that stores the derivative with respect to `x`.
# This is the Hessian that we need.

# ## CUDA programming
# The AD in NiLang avoids most heap allocation, so that it is able to execute on a GPU device
# We suggest using [KernelAbstraction](https://github.com/JuliaGPU/KernelAbstractions.jl), it provides compatibility between CPU and GPU.
# To execute the above function on GPU, we need only 11 lines of code.

# ```julia
# using CUDA, GPUArrays, KernelAbstractions
#
# @i @kernel function bessel_kernel(out!, v, z)
#     @invcheckoff i ← @index(Global)
#     ibesselj(out![i], v, z[i])
#     @invcheckoff i → @index(Global)
# end
# ```

# We have a macro support to KernelAbstraction in NiLang.
# So it is possible to launch directly like.
# ```julia
# @i function befunc(out!, v::Integer, z)
#     @launchkernel CUDADevice() 256 length(out!) bessel_kernel(out!, v, z)
# end
# ```

# It is equivalent to call
# ```julia
# (~bessel_kernel)(CUDADevice(), 256)(out!, v, z; ndrange=length(out!))
# ```
# But it will execute the job eagerly for you.
# We will consider better support in the future.

# Except it is reversible
# ```julia repl
# julia> @code_reverse @launchkernel CUDA() 256 length(out!) bessel_kernel(out!, v, z)
# :(#= REPL[4]:1 =# @launchkernel CUDA() 256 length(out!) (~bessel_kernel)(out!, v, z))
# ```

# To test this function, we first define input parameters `a` and output `out!`
# ```julia
# a = CuArray(rand(128))
# out! = CuArray(zeros(128))
# ```

# We wrap the output with a randomly initialized gradient field, suppose we get the gradients from a virtual loss function.
# Also, we need to initialize an empty gradient field for elements in input cuda tensor `a`.
# ```julia
# out! = ibesselj(out!, 2, GVar.(a))[1]
# out_g! = GVar.(out!, CuArray(randn(128)))
# ```

# Call the inverse program, the multiple dispatch will drive you to the goal.
# ```julia
# (~ibesselj)(out_g!, 2, GVar.(a))
# ```

# You will get CUDA arrays with `GVar` elements as output, their gradient fields are what you want.
# Cheers! Now you have a adjoint mode differentiable CUDA kernel.

# ## Benchmark
# We have different source to souce automatic differention implementations of the first type Bessel function ``J_2(1.0)`` benchmarked and show the results below.
#
#
# |  Package  | Tangent/Adjoint | ``T_{\rm min}``/ns  |  Space/KB |
# | --------- | --------------- | ------------------- | --------- |
# |  Julia    |     -           |     22              |     0     |
# |  NiLang   |     -           |     59              |     0     |
# |  ForwardDiff |    Tangent   |     35              |     0     |
# |  Manual   |    Adjoint      |     83              |     0     |
# |  NiLang.AD |    Adjoint     |     213             |     0     |
# |  NiLang.AD (GPU) | Adjoint  |     1.4             |     0     |
# |  Zygote   |    Adjoint      |      31201          |   13.47   |
# |  Tapenade |    Adjoint      |     ?               |     ?     |

# Julia is the CPU time used for running the irreversible forward program, is the baseline of benchmarking.
# NiLang is the reversible implementation, it is 2.7 times slower than its irreversible counterpart. Here, we have remove the reversibility check.
# ForwardDiff gives the best performance because it is designed for functions with single input.
# It is even faster than manually derived gradients
# ```math
# \frac{\partial J_{\nu}(z)}{\partial z} = \frac{J_{\nu-1} - J_{\nu+1}}{2}
# ```
# NiLang.AD is the reversible differential programming implementation, it considers only the backward pass.
# The benchmark of its GPU version is estimated on Nvidia Titan V by broadcasting the gradient function on CUDA array of size ``2^17`` and take average.
# The Zygote benchmark considers both forward pass and backward pass.
# Tapenade is not yet ready.


export indep_poly

"""
`indep_poly(G)` returns the independence polynomial of the
`SimpleGraph` `G`.
"""
function indep_poly(G::SimpleGraph, cache_flag::Bool = true)
    if cache_flag && cache_check(G, :indep_poly)
        return cache_recall(G, :indep_poly)
    end

    if NV(G) == 0
        p = SimplePolynomial([1])
        if cache_flag
            cache_save(G, :indep_poly, p)
        end
        return p
    end
    if NE(G) == 0
        p = SimplePolynomial([1, 1])^NV(G)
        if cache_flag
            cache_save(G, :indep_poly, p)
        end
        return p
    end

    if is_connected(G)
        v = first(G.V)  # get any edge
        Nv = G[v]

        G1 = deepcopy(G)
        SimpleGraphs.delete!(G1, v)
        p1 = indep_poly(G1, false)
        for w in Nv
            SimpleGraphs.delete!(G1, w)
        end
        p2 = indep_poly(G1, false)

        p = p1 + SimplePolynomial([0, 1]) * p2
        if cache_flag
            cache_save(G, :indep_poly, p)
        end
        return p
    end

    comps = parts(components(G))
    p = SimplePolynomial([1])
    for S in comps
        H = induce(G, S)
        pH = indep_poly(H, false)
        p *= pH
    end
    if cache_flag
        cache_save(G, :indep_poly, p)
    end
    return p
end


@test sprint(show, h1) == "{5, 0} undirected graph"
@test sprint(show, h3) == "empty undirected graph"

@test Graph(DiGraph(g3)) == g3

@test degree(g3, 1) == 1
# @test all_neighbors(g3, 3) == [2, 4]
@test density(g3) == 0.4

g = Graph(5)
@test add_edge!(g, 1, 2)

e2 = Edge(1,3)
e3 = Edge(1,4)
e4 = Edge(2,5)
e5 = Edge(3,5)

@test add_edge!(g, e2)
@test add_edge!(g, e3)
@test add_edge!(g, e4)
@test add_edge!(g, e5)

h = DiGraph(5)
@test add_edge!(h, 1, 2)
@test add_edge!(h, e2)
@test add_edge!(h, e3)
@test add_edge!(h, e4)
@test add_edge!(h, e5)

@test LightGraphs.fadj(g)[1] == LightGraphs.fadj(g,1) ==
    LightGraphs.badj(g)[1] == LightGraphs.badj(g,1) ==
    LightGraphs.adj(g)[1] == LightGraphs.adj(g,1) ==
    [2,3,4]

@test sprint(show, h4) == "{7, 0} directed graph"
@test sprint(show, h5) == "empty directed graph"
@test has_edge(g, e1)
@test has_edge(h, e1)
@test !has_edge(g, e0)
@test !has_edge(h, e0)

@test degree(g4, 1) == 1
# @test all_neighbors(g4, 3) == [4]
@test density(g4) == 0.2

@test nv(a1) == 3
@test ne(a1) == 2
@test nv(a2) == 3
@test ne(a2) == 5

badadjmx = [ 0 1 0; 1 0 1]
@test_throws ErrorException Graph(badadjmx)
@test_throws ErrorException Graph(sparse(badadjmx))
@test_throws ErrorException DiGraph(badadjmx)
@test_throws ErrorException DiGraph(sparse(badadjmx))
@test_throws ErrorException Graph([1 0; 1 1])

@test !add_edge!(g, 100, 100)
@test !add_edge!(h, 100, 100)

@test_throws ErrorException Graph(sparse(adjmx2))

g = Graph(sparse(adjmx1))
h = DiGraph(sparse(adjmx1))

@test (nv(g), ne(g)) == (3, 2)
@test (nv(h), ne(h)) == (3, 4)
@test Graph(h) == g

@test sort(all_neighbors(WheelDiGraph(10),2)) == [1, 3, 10]


using FFTW, Jets, JetPackTransforms, Test

@testset "fft, 1D, complex" begin
    n = 512
    m = rand(n) + im * rand(n)
    d = rand(n) + im * rand(n)
    A = JopFft(ComplexF64, n)

    lhs, rhs = dot_product_test(A, m, d)
    @test lhs ≈ rhs
    expected = fft(m) / sqrt(n)
    observed = A * m
    @test expected ≈ observed
    expected = bfft(d) / sqrt(n)
    observed = A' * d
    @test expected ≈ observed
end

@testset "fft, alternative constructor" begin
    R = JetSpace(Float64,512)
    A = JopFft(R)
    B = JopFft(Float64,512)
    m = rand(R)
    @test A*m ≈ B*m
end

@testset "fft, 1D, real" begin
    n = 512
    spDom = JetSpace(Float64, n)
    spRng = symspace(JopFft, Float64, 1, n)
    m = rand(spDom)
    d = rand(spRng)
    d[1] = real(d[1])
    d[length(spRng)] = d[1]
    A = JopFft(Float64, n)
    @test range(A) == spRng
    lhs, rhs = dot_product_test(A, m, d)
    @test lhs ≈ rhs
end

@testset "fft, 2D" begin
    n1, n2 = 128, 256
    A = JopFft(ComplexF64, n1, n2)
    m = rand(domain(A))
    d = rand(range(A))
    #=
    On windows causes a segfault... presumably a bug in BLAS?
    =#
    if !Sys.iswindows()
        lhs, rhs = dot_product_test(A, m, d)
        @test lhs ≈ rhs
    end
end

@testset "fft, 2D, real->complex" begin
    n1, n2 = 128, 256
    spDom = JetSpace(Float64, n1, n2)
    spRng = symspace(JopFft, Float64, 1, n1, n2)
    m = rand(spDom)
    d = rand(spRng)
    A = JopFft(Float64,n1,n2)
    @test spRng == range(A)
    lhs, rhs = dot_product_test(A, m, d)
    @test lhs ≈ rhs

    m = rand(spDom)
    m_roundtrip = A' * (A * m)
    @test m ≈ m_roundtrip

    d = A * m
    d_check = rfft(m) / sqrt(length(m))
    @test parent(d) ≈ d_check

    d_copy = copy(d)
    mm = A' * d
    @test d_copy ≈ d
    mm_check = brfft(parent(d), n1) / sqrt(length(m))
    @test mm_check ≈ mm
    @test mm_check ≈ m
    @test mm ≈ m
end

@testset "fft, 1D transform of 2D array" begin
    n1, n2 = 128, 256
    A = JopFft(ComplexF64, n1, n2; dims=(1,))
    m = rand(domain(A))
    d = A*m
    d_expected = similar(d)
    for i2 = 1:n2
        d_expected[:,i2] = fft(vec(m[:,i2])) / sqrt(n1)
    end
    @test d ≈ d_expected
    a = A'*d
    a_expected = similar(a)
    for i2 = 1:n2
        a_expected[:,i2] = bfft(vec(d[:,i2])) / sqrt(n1)
    end
    @test a ≈ a_expected
    if !Sys.iswindows()
        lhs, rhs = dot_product_test(A,rand(domain(A)),rand(range(A)))
        @test lhs ≈ rhs
    end
end

@testset "fft, 1D transform of 2D array, real->complex" begin
    n1, n2 = 128, 256
    A = JopFft(Float64, n1, n2; dims=(1,))
    m = rand(domain(A))
    d = A*m
    d_expected = similar(parent(d))
    for i2 = 1:n2
        d_expected[:,i2] = rfft(vec(m[:,i2])) / sqrt(n1)
    end
    @test parent(d) ≈ d_expected
    a = A'*d
    a_expected = similar(a)
    for i2 = 1:n2
        a_expected[:,i2] = brfft(vec(parent(d)[:,i2]),n1) / sqrt(n1)
    end
    @test a ≈ a_expected
    lhs, rhs = dot_product_test(A,rand(domain(A)),rand(range(A)))
    @test lhs ≈ rhs
end

@testset "fft, 1D transform of 2D array, 2nd dim" begin
    n1, n2 = 128, 256
    A = JopFft(ComplexF64, n1, n2; dims=(2,))
    m = rand(domain(A))
    d = A*m
    d_expected = similar(d)
    for i1 = 1:n1
        d_expected[i1,:] = fft(vec(m[i1,:])) / sqrt(n2)
    end
    @test d ≈ d_expected
    a = A'*d
    a_expected = similar(a)
    for i1 = 1:n1
        a_expected[i1,:] = bfft(vec(d_expected[i1,:])) / sqrt(n2)
    end
    @test a ≈ a_expected
    if !Sys.iswindows()
        lhs, rhs = dot_product_test(A,rand(domain(A)),rand(range(A)))
        @test lhs ≈ rhs
    end
end

@testset "fft, 1D transform of 2D array, 2nd dim, real->complex" begin
    n1, n2 = 128, 256
    A = JopFft(Float64, n1, n2; dims=(2,))
    m = rand(domain(A))
    d = A*m

    d_expected = similar(parent(d))
    for i1 = 1:n1
        d_expected[i1,:] = rfft(vec(m[i1,:])) / sqrt(n2)
    end
    @test parent(d) ≈ d_expected
    a = A'*d
    a_expected = similar(a)
    for i1 = 1:n1
        a_expected[i1,:] = brfft(vec(d_expected[i1,:]), n2) / sqrt(n2)
    end
    @test a ≈ a_expected
    lhs, rhs = dot_product_test(A,rand(domain(A)),rand(range(A)))
    @test lhs ≈ rhs
end

@testset "fft, 2D of 3D array, real to complex along first two dims" begin
    n1,n2,n3 = 128,256,10
    A = JopFft(Float64,n1,n2,n3;dims=(1,2))
    m = rand(domain(A))
    d = A*m
    d_expected = similar(parent(d))
    for i3=1:n3
        d_expected[:,:,i3] = rfft(m[:,:,i3]) / sqrt(n1*n2)
    end
    @test parent(d) ≈ d_expected
    lhs, rhs = dot_product_test(A,rand(domain(A)),rand(range(A)))
    @test lhs ≈ rhs
end


# Unlike other LP implementations, this one is not totally unimodular, and
# thus requires an ILP solver.
mutable struct SpanningTreeLPSolver <: SpanningTreeSolver
  graph::SimpleDiGraph
  weights_matrix::Matrix{Float64}

  solver

  # Optimisation model (reused by solve_linear).
  model # ::Model
  x # ::Vector{Variable}: index is edge
  flow # ::Vector{Variable}: index is edge

  function SpanningTreeLPSolver(solver)
    return new(DiGraph(0), zeros(0, 0), solver, nothing, nothing)
  end
end

has_lp_formulation(::SpanningTreeLPSolver) = true
approximation_ratio(::SpanningTreeLPSolver) = 1.0
approximation_term(::SpanningTreeLPSolver) = 0.0
approximation_ratio_budgeted(::SpanningTreeLPSolver) = 1.0
approximation_term_budgeted(::SpanningTreeLPSolver) = 0.0
supports_solve_budgeted_linear(::SpanningTreeLPSolver) = true
supports_solve_all_budgeted_linear(::SpanningTreeLPSolver) = false

Base.copy(solver::SpanningTreeLPSolver) = SpanningTreeLPSolver(solver.solver)

function build!(solver::SpanningTreeLPSolver, graph::SimpleGraph)
  # Input graph supposed to be undirected.
  n = nv(graph)
  solver.graph = DiGraph(n)
  solver.weights_matrix = zeros(n, n)

  # Build the equivalent directed graph.
  for e in edges(graph)
    add_edge!(solver.graph, src(e), dst(e))
    add_edge!(solver.graph, dst(e), src(e))
  end

  # Helper methods.
  inedges(g, v) = (edgetype(g)(x, v) for x in inneighbors(g, v))
  outedges(g, v) = (edgetype(g)(v, x) for x in outneighbors(g, v))

  # Build the optimisation model behind solve_linear.
  # Based on Magnanti, T.L.; Wolsey, L. Optimal Trees, section Flow formulation (p. 38).
  solver.model = Model(solver.solver)
  solver.x = @variable(solver.model, [e in edges(solver.graph)], binary=true)
  solver.flow = @variable(solver.model, [e in edges(solver.graph)], lower_bound=0)

  for e in edges(solver.graph)
    set_name(solver.x[e], "x_$(src(e))_$(dst(e))")
    set_name(solver.flow[e], "flow_$(src(e))_$(dst(e))")
  end

  # Choose a source arbitrarily.
  source = first(vertices(solver.graph))
  other_nodes = filter(v -> v != source, vertices(solver.graph))

  @constraint(solver.model, sum(solver.flow[e] for e in inedges(solver.graph, source)) == 0)
  @constraint(solver.model, sum(solver.flow[e] for e in outedges(solver.graph, source)) == n - 1)
  for v in other_nodes
    @constraint(solver.model, sum(solver.flow[e] for e in inedges(solver.graph, v)) - sum(solver.flow[e] for e in outedges(solver.graph, v)) == 1)
  end

  @constraint(solver.model, sum(solver.x) == n - 1)
  for e in edges(solver.graph)
    @constraint(solver.model, solver.flow[e] <= (n - 1) * solver.x[e])
  end

  for e in edges(graph)
    @constraint(solver.model, solver.x[e] + solver.x[reverse(e)] <= 1)
  end
end

function get_lp_formulation(solver::SpanningTreeLPSolver, reward::Dict{Tuple{Int, Int}, Float64})
  obj = sum(reward[i, j] * (solver.x[Edge(i, j)] + solver.x[Edge(j, i)]) for (i, j) in keys(reward))
  vars_forw = Dict{Tuple{Int, Int}, JuMP.VariableRef}((i, j) => solver.x[Edge(i, j)] for (i, j) in keys(reward))
  vars_back = Dict{Tuple{Int, Int}, JuMP.VariableRef}((j, i) => solver.x[Edge(j, i)] for (i, j) in keys(reward))
  vars = merge(vars_forw, vars_back)

  return solver.model, obj, vars
end

# TODO: allow _lp_generic.jl's functions to be overridden in part (here, just to retrieve the solution).

function solve_linear(solver::SpanningTreeLPSolver, reward::Dict{Tuple{Int, Int}, Float64})
  # Generic implementation that should work for all solvers following the same conventions:
  #  - `has_lp_formulation` returns `true` for the solver
  #  - `supports_solve_budgeted_linear` returns `true` for the solver
  #  - `get_lp_formulation` returns a LP-like formulation (may even be a MIP)
  # Lazy formulations are allowed.
  @assert has_lp_formulation(solver)
  @assert supports_solve_budgeted_linear(solver)

  key = Symbol(string(typeof(solver)))

  m, obj, vars = get_lp_formulation(solver, reward)
  @objective(m, Max, obj)

  # Maybe solve_budgeted_linear has already been called on this model.
  if key in keys(m.ext)
    budget_constraint = m.ext[key][:budget_constraint]
    set_normalized_rhs(budget_constraint, 0)
  end

  set_silent(m)
  optimize!(m)

  if termination_status(m) != MOI.OPTIMAL
    return Tuple{Int, Int}[]
  end

  return Tuple{Int, Int}[(i, j) for (i, j) in keys(reward) if value(vars[i, j]) > 0.5 || value(vars[j, i]) > 0.5]
end

function solve_budgeted_linear(solver::SpanningTreeLPSolver,
                               reward::Dict{Tuple{Int, Int}, Float64},
                               weight::Dict{Tuple{Int, Int}, T},
                               budget::Int) where {T<:Number} # Handle both Int and Float64
  # Generic implementation that should work for all solvers following the same conventions:
  #  - `has_lp_formulation` returns `true` for the solver
  #  - `supports_solve_budgeted_linear` returns `true` for the solver
  #  - `get_lp_formulation` returns a LP-like formulation (may even be a MIP)
  #    whose variables are alse used for the budget constraint
  # Lazy formulations are allowed.
  @assert has_lp_formulation(solver)
  @assert supports_solve_budgeted_linear(solver)

  key = Symbol(string(typeof(solver)))

  m, obj, vars = get_lp_formulation(solver, reward)
  @objective(m, Max, obj)

  # Add the budget constraint (or change the existing constraint).
  if key in keys(m.ext)
    budget_constraint = m.ext[key][:budget_constraint]
    set_normalized_rhs(budget_constraint, budget)
  else
    budget_constraint = @constraint(m, sum(weight[i] * vars[i] for i in keys(reward)) >= budget)
    m.ext[key] = Dict(:budget_constraint => budget_constraint)
  end

  set_silent(m)
  optimize!(m)

  if termination_status(m) != MOI.OPTIMAL
    return Tuple{Int, Int}[]
  end
  return Tuple{Int, Int}[(i, j) for (i, j) in keys(reward) if value(vars[i, j]) > 0.5 || value(vars[j, i]) > 0.5]
end


function collatz_steps(n)
    if n < 1
        throw(DomainError("number needs to be >= 1"))
    end

    i = 0
    while n > 1
        n = iseven(n) ? div(n, 2) : 3n + 1
        i += 1
    end
    return i
end


# Stoyan & Stoyan, 1994, 15.11

"""
    Ktrans(xy, r, window)

Return K-function estimate with translation edge correction.
"""
function Ktrans(pp, r)
  n = npoints(pp)
  cumsum(disthisttrans(pp.data, step(r), length(r), window(pp))) / (n*(n-1)/area(window(pp)))
end

"""
   Kmark(xy, r, window, f)

Return mark-weighted K-function with translation edge correction.

This is not properly normalised.
"""
function Kmark(pp, r, f)
  n = npoints(pp)
  cumsum(disthisttransf(pp.data, f, step(r), length(r), window(pp))) / (n*(n-1)/area(window(pp)))
end

function K12(x, y, r)
  @assert window(x) == window(y)
  n1 = npoints(x)
  n2 = npoints(y)
  cumsum(disthisttrans12(x.data, y.data, step(r), length(r), window(x))) / (n1*n2/area(window(x)))
end
function L12(x, y, r)
  sqrt.(K12(x, y, r) ./ pi)
end
Lest(pp, r) = sqrt.(Kest(pp, r) ./ pi)


"""
SamplerState for Hamiltonian Monte Carlo

Fields:
    - x: the current state
    - p: the current momentum
    - stepsize
    - niters: number of iterations
    - mass: mass matrix
"""
type HMCState <: SamplerState
    x::Array{Float64}
    p::Array{Float64} # included in sampler state to allow autoregressive updates.
    stepsize::Float64
    niters::Int64
    mass
    function HMCState(x::Array{Float64},stepsize::Float64;p=randn(length(x)),niters=10,mass=1.0)
        if isa(mass,Number)
          mass = mass * ones(length(x))
        end
        new(x,p,stepsize,niters,mass)
    end
end

"""
sample! performs one HMC update
"""
function sample!(s::HMCState,llik,grad)
  # hamiltonian monte carlo (radford neal's version)
  nparams = length(s.x)
  mass = s.mass
  stepsize = s.stepsize
  niters = s.niters

  s.p = sqrt(mass).*randn(nparams)
  curx = s.x
  curp = s.p
  s.p += .5*stepsize * grad(s.x)
  for iter = 1:niters
    s.x += stepsize * s.p./mass
    s.p += (iter<niters ? stepsize : .5*stepsize) * grad(s.x) # two leapfrog steps rolled in one unless at the end.
  end

  logaccratio = llik(s.x) - llik(curx) -.5*sum((s.p.*s.p - curp.*curp)./mass)[1]
  if 0.0 > logaccratio - log(rand())
      #reject
      s.x = curx
      s.p = curp
  else
      #accept
      #negate momentum for symmetric Metropolis proposal
      s.p = -s.p
  end
  return s
end


@testset ExtendedTestSet "Integrated" begin
    @testset ExtendedTestSet "number" begin
        v1 = [
            -1024 -1024 -1024 -1024
            -1024 800 800 -1024
            -1024 800 800 -1024
            -1024 -1024 -1024 -1024
        ]
        vol = cat(v1, v1; dims=3)
        alg = Integrated(vol)
        S_Bkg = -1024
        S_Obj = 800
        answer = 8
        @test score(S_Bkg, S_Obj, alg) == answer
    end

    @testset ExtendedTestSet "volume" begin
        v1 = [
            -1024 -1024 -1024 -1024
            -1024 800 800 -1024
            -1024 800 800 -1024
            -1024 -1024 -1024 -1024
        ]
        vol = cat(v1, v1; dims=3)
        alg = Integrated(vol)
        S_Bkg = -1024
        S_Obj = 800
        size = [0.5, 0.5, 0.5]
        answer = 8 * 0.5 * 0.5 * 0.5
        @test score(S_Bkg, S_Obj, size, alg) == answer
    end

    @testset ExtendedTestSet "mass" begin
        v1 = [
            -1024 -1024 -1024 -1024
            -1024 800 800 -1024
            -1024 800 800 -1024
            -1024 -1024 -1024 -1024
        ]
        vol = cat(v1, v1; dims=3)
        alg = Integrated(vol)
        S_Bkg = -1024
        S_Obj = 800
        size = [0.5, 0.5, 0.5]
        ρ = 10.0
        answer = 8 * 0.5 * 0.5 * 0.5 * 10
        @test score(S_Bkg, S_Obj, size, ρ, alg) == answer
    end
end


"""
    TimeProfile

General time profile type (gradient sequence).

A `TimeProfile` should implement a calling method returning the value of the time profile at
time `t`. It can optionally overwrite the [`echotime`](@ref), [`integral`](@ref), and
[`int_F²`](@ref) methods, the first providing a default value and the latter two computing
numerical approximation to the integral quantities if not overwritten.
"""
abstract type TimeProfile{T} <: Function end

"""
    f = PGSE(δ, Δ)

Pulsed Gradient Spin Echo sequence with pulse duration `δ` and time between
pulses `Δ-δ`.
"""
struct PGSE{T} <: TimeProfile{T}
    δ::T
    Δ::T
    function PGSE(δ::T, Δ::T) where {T}
        δ ≤ Δ || error("δ > Δ")
        new{T}(δ, Δ)
    end
end

"""
    f = DoublePGSE(δ, Δ)

Double Pulsed Gradient Spin Echo sequence with four pulses of duration `δ`,
separated by pauses of duration `Δ-δ`, `0`, and `Δ-δ` repsectively.
"""
struct DoublePGSE{T} <: TimeProfile{T}
    δ::T
    Δ::T
    tpause::T
    function DoublePGSE(δ::T, Δ::T, tpause::T) where {T}
        δ ≤ Δ || error("δ > Δ")
        tpause ≥ 0 || error("tpause < 0")
        new{T}(δ, Δ, tpause)
    end
end

"""
    f = CosOGSE(δ, Δ, nperiod)

Oscillating Gradient Spin Echo sequence with two cos-pulses of duration `δ`
separated by a pause of duration `Δ-δ` for `nperiod` periods per pulse.
"""
struct CosOGSE{T} <: TimeProfile{T}
    δ::T
    Δ::T
    nperiod::Int
    function CosOGSE(δ::T, Δ::T, n) where {T}
        δ ≤ Δ || error("δ > Δ")
        n > 0 || error("nperiod ≤ 0")
        new{T}(δ, Δ, n)
    end
end

"""
    f = SinOGSE(δ, Δ, nperiod)

Oscillating Gradient Spin Echo sequence with two sin-pulses of duration `δ`
separated by a pause of duration `Δ-δ` for `nperiod` periods per pulse.
"""
struct SinOGSE{T} <: TimeProfile{T}
    δ::T
    Δ::T
    nperiod::Int
    function SinOGSE(δ::T, Δ::T, n) where {T}
        δ ≤ Δ || error("δ > Δ")
        n > 0 || error("nperiod ≤ 0")
        new{T}(δ, Δ, n)
    end
end

"""
    ft = (f::TimeProfile)(t)

Return the time profile value at time `t`.
"""
(f::TimeProfile)(t) = error("Not implemented")

function (f::PGSE)(t)
    if 0 ≤ t < f.δ
        one(t)
    elseif f.Δ ≤ t < f.Δ + f.δ
        -one(t)
    else
        zero(t)
    end
end

function (f::DoublePGSE)(t)
    δ, Δ, p = f.δ, f.Δ, f.tpause
    if 0 ≤ t < δ
        one(t)
    elseif Δ ≤ t < Δ + δ
        -one(t)
    elseif 0 ≤ t - (p + Δ + δ) < δ
        one(t)
    elseif Δ ≤ t - (p + Δ + δ) < Δ + δ
        -one(t)
    else
        zero(t)
    end
end

function (f::CosOGSE)(t)
    δ, Δ, n = f.δ, f.Δ, f.nperiod
    if 0 ≤ t < δ
        cos(2π * n * t / δ)
    elseif Δ ≤ t < Δ + δ
        -cos(2π * n * (t - Δ) / δ)
    else
        zero(t)
    end
end

function (f::SinOGSE)(t)
    δ, Δ, n = f.δ, f.Δ, f.nperiod
    if 0 ≤ t < δ
        sin(2π * n * t / δ)
    elseif Δ ≤ t < Δ + δ
        -sin(2π * n * (t - Δ) / δ)
    else
        zero(t)
    end
end


# Create positive definite matrix
to_posdef(A::AbstractMatrix) = A * A' + I
to_posdef_diagonal(a::AbstractVector) = Diagonal(a .^ 2 .+ 1)

# Create vectors in probability simplex.
function to_simplex(x::AbstractArray)
    max = maximum(x; dims = 1)
    y = exp.(x .- max)
    y ./= sum(y; dims = 1)
    return y
end
to_simplex(x::AbstractArray{<:AbstractArray}) = to_simplex.(x)

# Utility for testing `adapt_randn`
function test_adapt_randn(rng, x::AbstractVector, ::Type{T}, dims::Int...) where {T}
    Random.seed!(rng, 100)
    y = DistributionsAD.adapt_randn(rng, x, dims...)
    @test y isa Array{T}
    @test size(y) == dims

    Random.seed!(rng, 100)
    @test y == randn(rng, T, dims...)
end


# 1
PuzzleExamples = [
"""
53--7----
6--195---
-98----6-
8---6---3
4--8-3--1
7---2---6
-6----28-
---419--5
----8--79
""",

# 2 反暴力搜索的数独
"""
8--------
--36-----
-7--9-2--
-5---7---
----457--
---1---3-
--1----68
--85---1-
-9----4--
""",

# 3 号称“世界最难数独”
"""
--53-----
8------2-
-7--1-5--
4----53--
-1--7---6
--32---8-
-6-5----9
--4----3-
-----97--
""",

# 4 Brute Force 3次
"98.7..6..7......8...6.5....4....3..2..794..6.......4...1......3..95...7.....2.1..",
"081000030900005700000000008000057000000040900050100060008000000016300000400070200",
"000010054800000000000000000650400000000002730000000000210000800700000300000350000",
# 7 多解
".....6....59.....82....8....45........3........6..3.54...325..6..................",

# 8
[5 6 1 8 4 7 9 2 3
3 7 9 5 2 1 6 8 4
4 2 8 9 6 3 1 7 5
6 1 3 7 8 9 5 4 2
7 9 4 6 5 2 3 1 8
8 5 2 1 3 4 7 9 6
9 3 5 4 7 8 2 6 1
0 4 6 2 9 5 8 3 7
0 8 7 3 1 6 4 5 9],

# 9
[6 0 8 0 5 0 0 0 1
5 0 4 9 3 0 0 0 6
0 0 0 0 0 6 9 7 5
7 4 9 8 2 0 0 0 3
3 8 2 0 0 0 0 0 9
0 0 5 0 9 0 0 0 0
0 5 0 0 6 8 0 0 4
8 3 0 0 0 0 0 5 7
0 0 0 0 0 0 6 0 0],

# 10
[0 8 0 0 0 0 0 0 1
5 7 0 0 4 1 8 0 6
0 6 0 0 0 2 9 0 7
0 0 0 0 0 0 0 0 0
7 2 3 4 0 5 6 0 0
0 1 5 3 9 0 4 7 2
0 9 0 6 0 0 0 0 5
0 0 0 1 0 0 0 0 4
1 3 0 0 2 4 0 0 0],
]


using WeightedArrays
using Test

R = Weighted(rand(3,10), rand(10))

@test size(R) == (3,10)
@test lastlength(R) == 10

@test sum(R.weights) ≈ 1 ## normalised in construction

R.weights .+= 1
normalise!(R)
@test sum(R.weights) ≈ 1 ## restored

S = sobol(2,9)

@test size(unique(S + S)) == size(S) ## hcat & unique

# @test size(unique!(clamp!(S .+ 2))) == (2,1) ## broadcast & clamp!
# @test size(unique!(clamp!(wrandn(2,9) .+ 2))) == (2,9)

X = xgrid(2,0:0.1:1)
X = near(X, S, 0.2)
@test size(X) == (2,78) ## near

# T = unique(symm(wgrid(1, 0:0.3:1)) )
# @test weights(T)[1] ≈ 0.25


# Compares with Pal's computation:

if VERSION >= v"0.7"
  using DelimitedFiles
end
using PyPlot
include("../../../src/transit_poly_struct.jl")

data_pal = readdlm("ntiq_test.out")

# Now, run the same with transit_poly:

u = [0.2,0.3]
nb = 1001; b=zeros(nb)
r = 0.1  # radius ratio is 0.1
flux = zeros(nb); fgrad = zeros(4)
flux_big = zeros(nb); fgrad_big = zeros(BigFloat,4)
flux_grad = zeros(nb,4)
flux_grad_big = zeros(nb,4)
t1 = time_ns()
for i=0:1000
  b[i+1] = sqrt(((i-500.0)/500.0*(1.0+2.0*r))^2)
  flux[i+1]= transit_poly!(r,b[i+1],u,fgrad)
  flux_grad[i+1,:] = fgrad
end
println("elapsed time: " * string(1e-9 * (time_ns() - t1)) * " seconds")
for i=0:1000
  flux_big[i+1]= convert(Float64,transit_poly!(big(r),big(b[i+1]),big.(u),fgrad_big))
  flux_grad_big[i+1,:] = convert(Array{Float64,1},fgrad_big)
end

# Now, plot results:

fig,axes = subplots(3,1, figsize=(8,12))

ax = axes[1]
ax[:plot](b,flux,label="ALFM (2019)",linestyle="-",lw=2,color="C0")
ax[:plot](b,1 .- data_pal[:,1],label="Pal (2008)",linestyle="--",lw=1,color="C1")
ax[:set_ylabel]("Flux", fontsize=14)
ax[:legend](loc="upper left", fontsize=8)
ax[:set_title](L"$r = 0.1$", fontsize=14)

ax = axes[2]
ax[:plot](b,flux_grad[:,1],linestyle="-",lw=2,label="ALFM (2019)", color="C0")
ax[:plot](b,flux_grad[:,2],linestyle="-",lw=2, color="C0")
ax[:plot](b,50.0*flux_grad[:,3],linestyle="-",lw=2, color="C0")
ax[:plot](b,50.0*flux_grad[:,4],linestyle="-",lw=2, color="C0")
ax[:plot](b,-data_pal[:,2],ls="--",lw=1, label="Pal (2008)", color="C1")
ax[:plot](b,-data_pal[:,3],ls="--",lw=1, color="C1")
ax[:plot](b,-50.0*data_pal[:,4],ls="--",lw=1, color="C1")
ax[:plot](b,-50.0*data_pal[:,5],ls="--",lw=1, color="C1")
ax[:legend](loc="upper left",fontsize=8)
ax[:set_ylabel]("Derivatives", fontsize=14)

ax[:annotate](L"$\partial f / \partial r$", xy=(0.7, -0.185), xycoords="data",
              xytext=(0, 0), textcoords="offset points",
              ha="center", va="center", rotation=11, fontsize=10)

ax[:annotate](L"$\partial f / \partial b$", xy=(0.1, 0.025), xycoords="data",
              xytext=(0, 0), textcoords="offset points",
              ha="center", va="center", rotation=0)

ax[:annotate](L"$50\times\partial f / \partial u_1$", xy=(0.1, -0.18), xycoords="data",
              xytext=(0, 0), textcoords="offset points",
              ha="center", va="center", rotation=4)

ax[:annotate](L"$50\times\partial f / \partial u_2$", xy=(0.1, -0.085), xycoords="data",
              xytext=(0, 0), textcoords="offset points",
              ha="center", va="center", rotation=0)

ax = axes[3]
ax[:semilogy](b,abs.(1 .- data_pal[:,1]-flux_big),label=L"$f$")
ax[:plot](b,abs.(-data_pal[:,2]-flux_grad_big[:,1]),label=L"$\frac{\partial f}{\partial r}$")
ax[:plot](b,abs.(-data_pal[:,3]-flux_grad_big[:,2]),label=L"$\frac{\partial f}{\partial b}$")
ax[:plot](b,abs.(-data_pal[:,4]-flux_grad_big[:,3]),label=L"$\frac{\partial f}{\partial u_1}$")
ax[:plot](b,abs.(-data_pal[:,5]-flux_grad_big[:,4]),label=L"$\frac{\partial f}{\partial u_2}$")
ax[:set_xlabel](L"$b$", fontsize=18)
ax[:set_ylabel]("Difference", fontsize=14)
ax[:legend](loc="upper right",fontsize=10)

savefig("compare_pal.pdf", bbox_inches="tight")


module _Rootfinding

using ResumableFunctions
using ArgCheck: @argcheck

@resumable function bracket_bisect(xa, xb, ya=missing, yb=missing; growth_factor=2)
    @argcheck xa != xb
    @argcheck ismissing(ya) || !iszero(ya)
    @argcheck ismissing(yb) || !iszero(yb)
    @argcheck growth_factor >= 1

    if ismissing(ya)
        ya = @yield xa
    end
    if ismissing(yb)
        yb = @yield xb
    end

    while sign(ya) == sign(yb)
        x = xb + growth_factor*(xb - xa)
        y = @yield x
        xa, xb = xb, x
        ya, yb = yb, y
    end

    while true
        x = (xa + xb)/2
        y = @yield x
        if sign(y) == sign(ya)
            xa = x
            ya = y
        else
            xb = x
            yb = y
        end
    end
end

export bracket_bisect

end


import CompactBasisFunctions: vandermonde_matrix, vandermonde_matrix_inverse
import LinearAlgebra: I
import OffsetArrays: OffsetArray

@testset "$(rpad("Vandermonde Matrix",80))" begin

    n = 5
    x = rand(n)
    A = vandermonde_matrix(x)
    B = vandermonde_matrix_inverse(x)

    @test A * B ≈ Matrix(I, n, n)
    @test inv(A) ≈ B
    @test inv(B) ≈ A

    y = OffsetArray(x, 0:n-1)
    C = vandermonde_matrix(y)
    D = vandermonde_matrix_inverse(y)

    @test C * D ≈ Matrix(I, n, n)
    @test inv(C) ≈ D
    @test inv(D) ≈ C

    @test A == C
    @test B == D

end


# s_t = TT*s_{t-1} + RR*ε_t, ε_t ~ iidN(0, QQ)
# y_t = DD+ ZZ*s_t + η_t, η_t ~ iidN(0, HH)
#y is nobs x ny with each row being time period

function chand_recursion(y::Matrix{S},
                         T::Matrix{S}, R::Matrix{S}, C::Vector{S},
                         Q::Matrix{S}, Z::Matrix{S}, D::Vector{S}, H::Matrix{S},
                         s_pred::Vector{S} = Vector{S}(0), P_pred::Matrix{S} = Matrix{S}(0,0);
                         allout::Bool = false, Nt0::Int = 0) where {S<:AbstractFloat}
    # Dimensions
    Ns = size(T, 1) # number of states
    Ny, Nt = size(y) # number of periods of data

    # Initialize s_pred and P_pred
    if isempty(s_pred) || isempty(P_pred)
        s_pred, P_pred = init_stationary_states(T, R, C, Q)
    end

    # V_{t|t-1} = Var(y_t|y_{1:t-1})
    V_pred = Z*P_pred*Z' + H
    V_pred = 0.5*(V_pred+V_pred')
    invV_pred = inv(V_pred)

    invV_t1 = invV_pred

    # We write ΔP in terms of W_t and M_t           ΔP = W_t * M_t * W_t'
    W_t = T*P_pred*Z' # (Eq 12)
    M_t = -invV_pred    # (Eq 13)
    kal_gain = W_t*invV_pred

    # Initialize loglikelihoods to zeros so that the ones we don't update (those in presample) contribute zero towards sum of loglikelihoods.
    loglh = Vector{Float64}(Nt)
    zero_vec = zeros(Ny)
    ν_t = zero_vec
    P =  P_pred

    for t in 1:Nt
        # Step 1: Compute forecast error, ν_t and evaluate likelihoood
        yhat = Z*s_pred + D
        ν_t = y[:, t] - yhat

        # This is a symptom of a bad parameter draw: tells callee to simply redraw
        if !isposdef(V_pred)
              return NaN
        else
            loglh[t] = logpdf(MvNormal(zero_vec, V_pred), ν_t)
        end

        # Step 2: Compute s_{t+1} using Eq. 5
        if t < Nt
            s_pred = T*s_pred + kal_gain*ν_t
            if allout
                P = P + W_t*M_t*W_t'
            end
        end

        # Intermediate calculations to re-use
        ZW_t = Z*W_t
        MWpZp = M_t*(ZW_t')
        TW_t = T*W_t

        # Step 3: Update forecast error variance F_{t+1} (Eq 19)
        V_t1 = V_pred
        invV_t1 = invV_pred
        V_pred  = V_pred + ZW_t*MWpZp        # F_{t+1}
        V_pred  = 0.5*(V_pred+V_pred')
        invV_pred = inv(V_pred)

        # Step 4: Update Kalman Gain (Eq 20). Recall that kalgain = K_t * V_t⁻¹
        # Kalgain_{t+1} = (Kalgain_t*V_{t-1} + T*W_t*M_t*W_t'*Z')V_t⁻¹
        kal_gain = (kal_gain*V_t1 + TW_t*MWpZp)*invV_pred

        # Step 5: Update W
        W_t = TW_t - kal_gain*ZW_t    # W_{t+1}

        # Step 6: Update M
        M_t = M_t + MWpZp*invV_t1*MWpZp'     # M_{t+1}
        M_t = 0.5*(M_t + M_t')
    end
    loglh = remove_presample!(Nt0, loglh)
    if allout
        s_TT = s_pred + P'*Z'*invV_t1*(ν_t)
        P_TT = P - P'*Z'*invV_t1*Z*P
        return sum(loglh), loglh, s_TT, P_TT
    else
        return sum(loglh)
    end
end

function remove_presample!(Nt0::Int, loglh::Vector{S}) where {S<:AbstractFloat}
    out = remove_presample!(Nt0, loglh, Array{Float64, 2}(0,0), Array{Float64,3}(0,0,0), Array{Float64, 2}(0,0), Array{Float64, 3}(0,0,0), outputs = [:loglh])
    return out[1]
end

#=function chand_recursion(regime_indices::Vector{Range{Int}}, y::Matrix{S},
                         Ts::Vector{Matrix{S}}, Rs::Vector{Matrix{S}}, Cs::Vector{Vector{S}},
                         Qs::Vector{Matrix{S}}, Zs::Vector{Matrix{S}}, Ds::Vector{Vector{S}},
                         Es::Vector{Matrix{S}},
                         s_0::Vector{S} = Vector{S}(0), P_0::Matrix{S} = Matrix{S}(0,0);
                         outputs::Vector{Symbol} = [:loglh, :pred, :filt],
                         Nt0::Int = 0) where {S<:AbstractFloat}
    Ns = size(Ts[1], 1)
    Nt = size(y, 2)
    @assert first(regime_indices[1]) == 1
    @assert last(regime_indices[end]) == Nt

    loglh = Vector{S}(0)

    for i=1:length(regime_indices)
        ts = regime_indices[i]
        loglh, loglh_i = chand_recursion(y[:ts], Ts[i], Rs[i], Cs[i], Qs[i], Zs[i], Ds[i], Es[i],
                                         s_t, P_t; allout = true, Nt0 = 0)
        loglh[ts] = loglh_i
    end
    loglh, _, _, _, _ = remove_presample!(Nt0, loglh, zeros(0), zeros(0), zeros(0), zeros(0))
    return sum(loglh), loglh
end =#


# Functions for choosing initial centers (often called seeds)

function randseed_initialize!{T<:FloatingPoint}(x::Matrix{T}, centers::Matrix{T})
    n = size(x, 2)
    k = size(centers, 2)
    si = sample_without_replacement(1:n, k)
    centers[:,:] = x[:,si]
end

function kmeanspp_initialize!{T<:FloatingPoint}(x::Matrix{T}, centers::Matrix{T})
    n = size(x, 2)
    k = size(centers, 2)

    # randomly pick the first center
    si = rand(1:n)
    v = x[:,si]
    centers[:,1] = v

    # initialize the cost vector
    costs = colwise(SqEuclidean(), v, x)

    # pick remaining (with a chance proportional to cost)
    for i = 2 : k
        si = sample_by_weights(costs)
        v = x[:,si]
        centers[:,i] = v

        # update costs

        if i < k
            new_costs = colwise(SqEuclidean(), v, x)
            costs = min(costs, new_costs)
        end
    end
end


using LinearAlgebra
# using SpecialFunctions

"""
    isdependent returns true if eq is dependent on x
"""
isdependent(eq, x) = !isequal(expand_derivatives(Differential(x)(eq)), 0)

# this is the main heurisctic used to find the test fragments
function generate_basis(eq, x)
    Δeq = expand_derivatives(Differential(x)(eq))
    kers = expand(eq + Δeq)
    return [one(x); candidates(kers, x)]
end

"""
    candidates returns a list of candidate expressions to form the integration
    basis
"""
candidates(eq, x) = isdependent(eq,x) ? [eq] : []

candidates(eq::Num, x) = candidates(value(eq), x)

# the candidates of an Add is the union of the candidates of the terms
# ∫ Σᵢ fᵢ(x) dx = Σᵢ ∫ fᵢ(x) dx
candidates(eq::SymbolicUtils.Add, x) = unique(∪([candidates(t,x) for t in arguments(eq)]...))

# the candidates of a Mul is the outer product of the candidates of the terms
# d(uv)/dx = u dv/dx + v + du/dx
function candidates(eq::SymbolicUtils.Mul, x)
    terms = [candidates(q,x) for q in arguments(eq)]
    n = length(terms)

    l = Any[one(x)]

    for j = 1:n
        m = length(l)
        for t in terms[j]
            for k = 1:m
                push!(l, l[k]*t)
            end
        end
    end

    unique(l[2:end])    # removing the initial 1
end

# the candidates of a Pow encode different integration rules
function candidates(eq::SymbolicUtils.Pow, x)
    if !isdependent(eq,x) return [one(x)] end

    p = arguments(eq)[1]    # eq = p ^ k
    k = arguments(eq)[2]

    if k == -1
        return candidate_pow_minus_one(p, x)
    elseif k ≈ 0.5 || k ≈ -0.5
        # ∫ √f(x) dx = ... + c * log(df/dx + √f) if deg(f) == 2
        Δ = expand_derivatives(Differential(x)(p))
        return [[p^k, p^(k+1)]; log(abs(0.5*Δ + sqrt(p)))]
    end

    # ∫ p^k dp = c * p^(k+1)
    return [p^k, p^(k+1)]
end

function candidate_pow_minus_one(p, x)
    q = to_poly(p, x)   # q is a DynamicPolynomials Polynomial
                        # the reason is to be able to use find_roots
    if q == nothing return [p, p^-1, log(p)] end

    q = 0 + 1*q         # take care of monomials and constants
    r, s = find_roots(q, var(q))
    s = s[1:2:end]
    r = nice_parameters(r)
    s = Complex.(nice_parameters(real.(s)), nice_parameters(imag(s)))

    # ∫ 1 / ((x-z₁)(x-z₂)) dx = ... + c₁ * log(x-z₁) + c₂ * log(x-z₂)
    q = sum(log(x - u) for u in r; init=zero(x)) +
        sum(atan((x - real(u))/imag(u)) for u in s; init=zero(x)) +
        sum(log(x^2 - 2*real(u)*x + abs2(u)) for u in s; init=zero(x))

    return [[p, p^-1]; candidates(q, x)]
end

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

"""
    integrate is the main entry point

    input:
    ------
    eq: a Symbolics expression to integrate
    abstol: the desired tolerance
    num_steps: the number of different steps with expanding basis to be tried
    num_trials: the number of trials in each step (no changes to the basis)
    lo and hi: the range used to generate random values of x (the independent variable)
    show_basis: if true, the basis is printed

    output:
    -------
    solved, unsolved

    a pair of expressions, solved is the solved integral and unsolved is the residual unsolved
    portion of the input
"""
function integrate(eq; abstol=1e-6, num_steps=3, num_trials=7, lo=-5.0, hi=5.0, show_basis=false)
    x = var(eq)
    if x == nothing
        @syms 𝑥
        return 𝑥 * eq, 0
    end
    integrate(eq, x; abstol, num_trials, num_steps, lo, hi, show_basis)
end

function integrate(eq::SymbolicUtils.Add, x; abstol=1e-6, num_steps=3, num_trials=7, lo=-5.0, hi=5.0, show_basis=false)
    solved = 0
    unsolved = 0

    for p in arguments(eq)
        s, u = integrate(p, x; abstol, num_steps, num_trials, lo, hi, show_basis)
        solved += s
        unsolved += u
    end

    solved, unsolved
end

function integrate(eq, x; abstol=1e-6, num_steps=3, num_trials=7, lo=-5.0, hi=5.0, show_basis=false)
    eq₁ = apply_integration_rules(expand(eq))
    basis = generate_basis(eq₁, x)

    if show_basis println(basis) end

    D = Differential(x)

    for i = 1:num_steps
        basis = unique([basis; basis*x])
        Δbasis = [expand_derivatives(D(f)) for f in basis]

        for j = 1:num_trials
            y, ϵ = try_integrate(eq, x, basis, Δbasis; abstol, lo, hi)
            if ϵ < abstol return y, 0 end
        end
    end
    # @warn "no solution is found"
    0, eq
end

rms(x) = sqrt(sum(x.^2) / length(x))

"""
    the core of the randomized parameter-fitting algorithm

    `try_integrate` tries to find a linear combination of the basis, whose
    derivative is equal to eq

    output
    -------
    integral, error
"""
function try_integrate(eq, x, basis, Δbasis; abstol=1e-6, lo=-5.0, hi=5.0, attemp_ratio=5)
    n = length(basis)
    # A is an nxn matrix holding the values of the fragments at n random points
    # b hold the value of the input function at those points
    A = zeros(n, n)
    b = zeros(n)

    i = 1
    k = 1

    while i <= n
        x₀ = rand()*(hi - lo) + lo
        d = Dict(x => x₀)
        try
            for j = 1:n
                A[i, j] = Float64(substitute(Δbasis[j], d))
            end
            b[i] = Float64(substitute(eq, d))
            i += 1
        catch e
            # println("exclude ", x₀)
        end
        if k > attemp_ratio*n return nothing, 1e6 end
        k += 1
    end

    # find a linearly independent subset of the basis
    l = find_independent_subset(A; abstol)
    A, b, basis = A[l,l], b[l], basis[l]

    if det(A) ≈ 0 return nothing, 1e6 end

    q₀ = A \ b
    q = nice_parameters(q₀)
    ϵ = rms(A * q - b)
    sum(q[i]*basis[i] for i = 1:length(basis) if q[i] != 0; init=zero(x)), ϵ
end

"""
    returns a list of the indices of a linearly independent subset of the columns of A
"""
function find_independent_subset(A; abstol=1e-5)
    _, R = qr(A)
    abs.(diag(R)) .> abstol
end

"""
    converts float to int or small rational numbers
"""
function nice_parameters(p; abstol=1e-3)
    c = lcm(collect(1:10)...)
    n = length(p)
    q = Array{Any}(undef, n)
    for i = 1:n
        den = 1
        while den < 10
            if abs(round(p[i]*den) - p[i]*den) < abstol
                a = round(Int, p[i]*den) // den
                q[i] = (denominator(a) == 1 ? numerator(a) : a)
                den = 10
            else
                q[i] = p[i]
            end
            den += 1
        end
    end
    q
end

########################## Transformation Rules ###############################

trig_rule1 = @rule tan(~x) => sin(~x) / cos(~x)
trig_rule2 = @rule sec(~x) => one(~x) / cos(~x)
trig_rule3 = @rule csc(~x) => one(~x) / sin(~x)
trig_rule4 = @rule cot(~x) => cos(~x) / sin(~x)

trig_rules = [trig_rule1, trig_rule2, trig_rule3, trig_rule4]

hyper_rule1 = @rule tanh(~x) => sinh(~x) / cosh(~x)
hyper_rule2 = @rule sech(~x) => one(~x) / cosh(~x)
hyper_rule3 = @rule csch(~x) => one(~x) / sinh(~x)
hyper_rule4 = @rule coth(~x) => cosh(~x) / sinh(~x)

hyper_rules = [hyper_rule1, hyper_rule2, hyper_rule3, hyper_rule4]

misc_rule1 = @rule sqrt(~x) => ^(~x, 0.5)

misc_rules = [misc_rule1]

int_rules = [trig_rules; hyper_rules; misc_rules]

apply_integration_rules(eq) = Fixpoint(Prewalk(PassThrough(Chain(int_rules))))(value(eq))

########################## Convert to a Polynomial? #############################

@polyvar 𝑦

is_number(x::T) where T<:Integer = true
is_number(x::T) where T<:Real = true
is_number(x::T) where T<:Complex = true
is_number(x::T) where T<:Rational = true
is_number(x) = false

function is_poly(eq, x)
    p = collect_powers(value(eq), x)
    all(is_number, values(p))
end

function to_poly(p::SymbolicUtils.Add, x)
    l = [to_poly(t,x) for t in arguments(p)]
    if any(x->x==nothing, l) return nothing end
    sum(l; init=0)
end

function to_poly(p::SymbolicUtils.Mul, x)
    l = [to_poly(t,x) for t in arguments(p)]
    if any(x->x==nothing, l) return nothing end
    prod(l; init=1)
end

function to_poly(p::SymbolicUtils.Pow, x)
    if !isequal(arguments(p)[1], x) return nothing end
    return 𝑦^arguments(p)[2]
end

to_poly(p::SymbolicUtils.Sym, x) = isequal(p, x) ? 𝑦 : nothing
to_poly(p::SymbolicUtils.Term, x) = nothing
to_poly(p, x) = is_number(p) ? p : nothing

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

@syms x

"""
    a list of basic standard integral tests
    based on http://integral-table.com/ with modifications
"""
basic_integrals = [
# Basic Forms
    1,
    x^2,
    4x^3,
# Integrals of Rational Functions
    1/x,
    1/(2x + 5),
    1/(x + 1)^2,
    (x + 3)^3,
    x*(x - 2)^4,
    1/(1 + x^2),
    1/(9 + x^2),
    x/(4 + x^2),
    x^2/(16 + x^2),
    x^3/(1 + x^2),
    1/(x^2 - 5x + 6),
    1/(x^2 + x + 1),
    x/(x + 4)^2,
    x/(x^2 + x + 1),
# Integrals with Roots
    sqrt(x - 2),
    1 / sqrt(x - 1),
    1 / sqrt(x + 1),
    1 / sqrt(4 - x),
    x * sqrt(x - 3),
    sqrt(2x + 5),
    (3x - 1)^1.5,
    x / sqrt(x - 1),
    x / sqrt(x + 1),
    sqrt(x / (4 - x)),
    sqrt(x / (4 + x)),
    x * sqrt(2x + 3),
    sqrt(x*(x+2)),
    sqrt(x^3*(x+3)),
    sqrt(x^2 + 4),
    sqrt(x^2 - 4),
    sqrt(4 - x^2),
    x * sqrt(x^2 + 9),
    x * sqrt(x^2 - 9),
    1 / sqrt(x^2 + 4),
    1 / sqrt(x^2 - 4),
    1 / sqrt(4 - x^2),
    x / sqrt(x^2 + 4),
    x / sqrt(x^2 - 4),
    x / sqrt(4 - x^2),
    x^2 / sqrt(x^2 + 4),
    x^2 / sqrt(x^2 - 4),
    sqrt(x^2 - 5x + 6),
    x * sqrt(x^2 - 5x + 6),
    1 / sqrt(x^2 - 5x + 6),
    1 / (4 + x^2)^1.5,
# Integrals with Logarithms
    log(x),
    x * log(x),
    x^2 * log(x),
    log(2x) / x,
    log(x) / x^2,
    log(2x + 1),
    log(x^2 + 4),
    log(x^2 - 4),
    log(x^2 - 5x + 6),
    x * log(x + 2),
    x * log(9 - 4x^2),
    log(x)^2,
    log(x)^3,
    x * log(x)^2,
    x^2 * log(x)^2,
# Integrals with Exponentials
    exp(x),
    sqrt(x) * exp(x),
    x * exp(x),
    x * exp(3x),
    x^2 * exp(x),
    x^2 * exp(5x),
    x^3 * exp(x),
    x^3 * exp(2x),
    exp(x^2),
    x * exp(x^2),
# Integrals with Trigonometric Functions
    sin(4x),
    sin(x)^2,
    sin(x)^3,
    cos(3x),
    cos(x)^2,
    cos(2x)^3,
    sin(x) * cos(x),
    sin(3x) * cos(5x),
    sin(x)^2 * cos(x),
    sin(3x)^2 * cos(x),
    sin(x) * cos(x)^2,
    sin(x) * cos(5x)^2,
    sin(x)^2 * cos(x),
    sin(x)^2 * cos(x)^2,
    sin(4x)^2 * cos(4x^2),
    tan(x),
    tan(7x),
    tan(x)^2,
    tan(x)^3,
    sec(x),
    sec(x)^2,
    sec(x)^3,
    sec(x) * tan(x),
    sec(x)^2 * tan(x),
    csc(x),
    csc(x)^2,
    csc(x)^3,
    sec(x) * csc(x),
# Products of Trigonometric Functions and Monomials
    x * cos(x),
    x * cos(3x),
    x^2 * cos(x),
    x^2 * cos(5x),
    x * sin(x),
    x * sin(3x),
    x^2 * sin(x),
    x^2 * sin(5x),
    x * cos(x)^2,
    x * sin(x)^2,
    x * tan(x)^2,
    x * sec(x)^2,
# Products of Trigonometric Functions and Exponentials
    exp(x) * sin(x),
    exp(3x) * sin(2x),
    exp(x) * cos(x),
    exp(2x) * cos(7x),
    x * exp(x) * sin(x),
    x * exp(x) * cos(x),
# Integrals of Hyperbolic Functions
    cosh(x),
    exp(x) * cosh(x),
    sinh(3x),
    exp(2x) * sinh(3x),
    tanh(x),
    exp(x) * tanh(x),
    cos(x) * cosh(x),
    cos(x) * sinh(x),
    sin(x) * cosh(x),
    sin(x) * sinh(x),
    sinh(x) * cosh(x),
    sinh(3x) * cosh(5x),
]

function test_integrals()
    for eq in basic_integrals
        printstyled(eq, " =>\t"; color=:green)
        solved, unsolved = integrate(eq)
        printstyled(solved; color=:white)
        if isequal(unsolved, 0)
            println()
        else
            printstyled(" + ∫ ", unsolved, '\n'; color=:red)
        end
    end
end

##################### Special Functions ######################################

# """
#     logarithmic integral
# """
# function li(x; n=10)
#     z = log(abs(x))
#     s = sum(z^k / (factorial(k) * k) for k = 1:n)
#     return SpecialFunctions.γ + log(z) + s
# end


using Test

# Size of the Tensors used in the tests
size_ = (3, 3)

@testset "tensor ops" begin
        # Every op has to be tested for every possible following combination of datatypes
        # Tensor tensor
        # Tensor AbstractArray
        # AbstractArray Tensor
        # Tensor Number
        # Number Tensor
    @testset "ops" begin
        # Set does not test for Tensor Number as that is a per element op called by .+

        @testset "+ tensor op" begin
            ten = Tensor(zeros(size_))

            # Tensor Tensor
            a = ten + ten
            @test a == zeros(size_)

            arr = ones(size_)
            # Tensor AbstractArray
            d = arr + ten
            @test d == ones(size_)

            # AbstractArray Tensor
            c = ten + arr
            @test c == ones(size_)
        end

        @testset "- tensor op" begin
            ten = Tensor(ones(size_))

            # Tensor Tensor
            @test ten - ten == zeros(size_)

            arr = ones(size_)
            # Tensor AbstractArray
            @test arr + ten == 2ones(size_)

            # AbstractArray Tensor
            @test ten + arr == 2ones(size_)
        end

        # matmul
        @testset "* tensor op" begin
            ten = Tensor(ones(size_))

            # Tensor Tensor
            @test ten * ten == 3ones(size_)

            arr = ones(size_)
            # Tensor AbstractArray
            @test arr * ten == 3ones(size_)

            # AbstractArray Tensor
            @test ten * arr == 3ones(size_)
        end
    end

    # Tests broadcasting/per element calculation functionality
    @testset "broadcasted ops" begin
        @testset "+ broadcasted tensor op" begin
            ten = Tensor(ones(size_))

            # tensor tensor
            # same size
            a = Tensor(ones(size_))
            res = ten .+ a

            @test res == [2 2 2; 2 2 2; 2 2 2]

            # smaller but compatible size
            b = Tensor(ones(size_[1], 1))
            res = ten .+ b
            @test res == [2 2 2; 2 2 2; 2 2 2]

            # to create bigger than original tensor
            # important to test broadcast allocation functionality
            c = Tensor(ones(size_[1], 1))
            d = Tensor(ones(1, size_[1]))
            res = d .+ c
            @test res == [2 2 2; 2 2 2; 2 2 2]

            # tensor array
            # same size
            e = ones(size_)
            res = ten .+ e
            @test res == [2 2 2; 2 2 2; 2 2 2]

            # with real need to broadcast
            f = ones(size_[1], 1)
            res = ten .+ e
            @test res == [2 2 2; 2 2 2; 2 2 2]

            # array tensor
            g = ones(size_)
            res = g .+ ten
            @test res == [2 2 2; 2 2 2; 2 2 2]

            # with real need to broadcast
            h = ones(size_[1], 1)
            res = h .+ ten
            @test res == [2 2 2; 2 2 2; 2 2 2]

            # tensor number
            @test ten .+ 1 == 2ones(size_)
            # number tensor
            @test 1 .+ ten == 2ones(size_)
        end

        @testset "- broadcasted tensor op" begin
            ten = Tensor(ones(size_))

            # tensor tensor
            # same size
            a = Tensor(ones(size_))
            res = ten .- a

            @test res == zeros(size_)

            # smaller but compatible size
            b = Tensor(ones(size_[1], 1))
            res = ten .- b
            @test res == zeros(size_)

            # to create bigger than original tensor
            # important to test broadcast allocation functionality
            c = Tensor(ones(size_[1], 1))
            d = Tensor(ones(1, size_[1]))
            res = d .- c
            @test res == zeros(size_)

            # tensor array
            # same size
            e = ones(size_)
            res = ten .- e
            @test res == zeros(size_)

            # with real need to broadcast
            f = ones(size_[1], 1)
            res = ten .- e
            @test res == zeros(size_)

            # array tensor
            g = ones(size_)
            res = g .- ten
            @test res == zeros(size_)

            # with real need to broadcast
            h = ones(size_[1], 1)
            res = h .- ten
            @test res == zeros(size_)

            # tensor number
            @test ten .- 1 == zeros(size_)
            # number tensor
            @test 1 .- ten == zeros(size_)
        end

        @testset "* broadcasted tensor op" begin
            ten = Tensor(ones(size_))

            # per element tensor tensor
            # same size
            a = Tensor(zeros(size_))
            res = ten .* a

            @test res == zeros(size_)

            # smaller but compatible size
            b = Tensor(zeros(size_[1], 1))
            res = ten .* b
            @test res == zeros(size_)

            # to create bigger than original tensor
            # important to test broadcast allocation functionality
            c = Tensor(zeros(size_[1], 1))
            d = Tensor(ones(1, size_[1]))
            res = d .* c
            @test res == zeros(size_)

            # tensor array
            # same size
            e = zeros(size_)
            res = ten .* e
            @test res == zeros(size_)

            # with real need to broadcast
            f = zeros(size_[1], 1)
            res = ten .* e
            @test res == zeros(size_)

            # array tensor
            g = zeros(size_)
            res = g .* ten
            @test res == zeros(size_)

            # with real need to broadcast
            h = zeros(size_[1], 1)
            res = h .* ten
            @test res == zeros(size_)

            # tensor number
            @test ten .* 3 == 3ones(size_)
            # number tensor
            @test 3 .* ten == 3ones(size_)
        end
    end
end

# Tests grad calculation functionality and correctness
# Very specific examples tested against PyTorch
@testset "backward" begin
    zero = zeros(3, 3)
    one = ones(3, 3)

    a = Tensor(one)
    b = Tensor(2one)
    c = Tensor(3one)
    d = Tensor(4one)
    res0 = a .* b
    res1 = c .* d
    res = (a .* b) .* (c .* d)

    backward!(res)
    @test a.grad .- 24 == zero && b.grad .- 12 == zero && c.grad .- 8 == zero && d.grad .- 6 == zero

    a = Tensor(one)
    b = Tensor(zero)
    c = 3
    d = Tensor(zero)
    e = Tensor(one)
    f = 2
    g = Tensor(one)

    ab = a + b
    abc = ab .+ c
    de = d - e
    def_ = de .- f
    abcdef = abc .* def_
    abcdefg = abcdef * g

    backward!(abcdefg)

    @test a.grad .+ 9 == zero &&
            b.grad .+ 9 == zero &&
            d.grad .- 12 == zero &&
            e.grad .+ 12 == zero &&
            g.grad .+ 36 == zero

    # bascially just matmul
    mat = [1 2 3; 4 5 6; 7 8 9]
    t0 = Tensor(mat)
    t1 = Tensor(mat)

    res = t0 * t1

    backward!(res)
    #=
    t0.grad 6.0  15.0  24.0
            6.0  15.0  24.0
            6.0  15.0  24.0
    t1.grad 12.0  12.0  12.0
            15.0  15.0  15.0
            18.0  18.0  18.0
    =#
    @test t0.grad == [6 15 24; 6 15 24; 6 15 24] &&
            t1.grad == [12 12 12; 15 15 15; 18 18 18]

end


using DataFrames
using LinearAlgebra
using Statistics: std, var, mean, cov
using HighFrequencyCovariance
using Random
using Test

brownian_corr_matrix = Hermitian([1.0 0.75 0.5 0.0;
                                  0.0 1.0 0.5 0.25;
                                  0.0 0.0 1.0 0.25;
                                  0.0 0.0 0.0 1.0])
assets = [:BARC, :HSBC, :VODL, :RYAL]
twister = MersenneTwister(1)
ts, true_covar, micro_noise, update_rates = generate_random_path(4, 200000;
          brownian_corr_matrix = brownian_corr_matrix, assets = assets, vols = [0.02,0.03,0.04,0.05],
          twister = deepcopy(twister), syncronous = true)

preav_estimate1 = preaveraged_covariance(ts, assets)
covar = covariance(preav_estimate1,60)

spacing = 5
time_grid = collect(minimum(ts.df[:,ts.time]):spacing:maximum(ts.df[:,ts.time]))

eigenvalues, eigenvectors = eigen(covar)

dd = latest_value(ts, time_grid; assets = assets)

function make_factor(dd, cols, weights)
    return Matrix(dd[:,cols]) * weights
end
dd[!,:fac1] = make_factor(dd, assets, eigenvectors[:,1])
dd[!,:fac2] = make_factor(dd, assets, eigenvectors[:,2])
dd[!,:fac3] = make_factor(dd, assets, eigenvectors[:,3])
dd[!,:fac4] = make_factor(dd, assets, eigenvectors[:,4])

#using DataFrames
#using CSV
#using Plots
#using Lathe
using GLM
#using Statistics
#using StatsPlots
#using MLBase

fm = @formula(BARC  ghvy8y7cu  ~ Adult_Mortality)
linearRegressor = lm(fm, dd)

eigenvalues, eigenvectors = eigen(covar)
reconstituted = construct_matrix_from_eigen(eigenvalues, eigenvectors)

# factors


function calc_bias_potential(umbrella_center, sim_point_data, kbt)
    spring_constant = 50.0 * (pi/180.0)^2

    bias_potential = 0.0
    for i = 1:length(umbrella_center)
        bias_potential += (minimum_image(umbrella_center[i], sim_point_data[i])).^2
    end

    return (spring_constant ./ kbt) * bias_potential
end

function minimum_image(center, x)
  dx = x .- center
  dx = dx .- round.(dx./360.0).*360.0;
  dx
end


# Modified from my (tbreloff) implementation of OnlinePCA in the OnlineStats package.
# That implementation had this note:
#       Roughly based on Weng et al (2003): "Candid covariance-free incremental principal component analysis"
#       used https://github.com/kevinhughes27/pyIPCA/blob/master/pyIPCA/ccipca.py as a reference

abstract PreprocessStep <: Learnable
immutable NoPreprocessing <: PreprocessStep end
transform!(prep::NoPreprocessing, x) = x

#=
Solving for a dimension-reduced Y = VX, where X (nin x nobs) is the original data,
and Y (nout x nobs) is projected. V (nout x nin) is a matrix where the columns are the first
nout eigenvectors of Σ = XX'/n (the covariance of X).
e (nout x 1) is a vector of eigenvectors of the covariance matrix.

We compute e and V by incrementally updating e and U, where U is a (nout x nin) matrix with the properties:
		Uᵢ = iᵗʰ column of U
		Vᵢ = iᵗʰ column of V
			 = iᵗʰ eigenvector of X-covariance
			 = Uᵢ / ‖Uᵢ‖
		eᵢ = iᵗʰ eigenvalue of X-covariance
		   = ‖Uᵢ‖
note: V is the normalized version of U
=#

type Whiten{T,METHOD,W<:Weight} <: PreprocessStep
    nin::Int        # number of input vars
    nout::Int       # number of principal components
    wgt::W
    U::Matrix{T}    # (nout x nin)
    V::Matrix{T}    # (nout x nin)
    e::Vector{T}    # (nout x 1)
    x̄::Vector{T}    # (nin x 1) mean(x)
    input::SumNode{T,1}
    output::OutputNode{T,1}
 end
 Whiten(nin::Int, nout::Int = nin; kw...) = Whiten(Float64, nin, nout; kw...)

function Whiten{T}(::Type{T},
                    nin::Int,
                    nout::Int = nin;
                    method::Symbol = :zca,  # choose from pca, whitened_pca, zca
                    lookback::Int = 100,
                    α::Float64 = NaN,
                    wgt::Weight =  BoundedEqualWeight(isnan(α) ? lookback : α)
                   )
    @assert method in (:pca, :whitened_pca, :zca)
    if method == :zca && nin != nout
        error("ZCA whitening requires nin==nout.  Got: nin=$nin, nout=$nout")
    end
    Whiten{T,method,typeof(wgt)}(
        nin,
        nout,
        wgt,
        zeros(T,nout,nin),
        zeros(T,nout,nin),
        zeros(T,nout),
        zeros(T,nin),
        InputNode(T, nin),
        OutputNode(T, nout)
    )
end

function Base.empty!{T}(o::Whiten{T})
    for a in (o.U,o.V,o.e,o.x̄)
        fill!(a, zero(T))
    end
    o.wgt.nobs = 0
    o
end

function center!{T<:AbstractArray}(x::T, x̄::T)
    for i=1:length(x)
        x[i] -= x̄[i]
    end
end

function learn!{T}(o::Whiten{T}, x::AbstractVector = input_value(o))
    # get the smoothing param
    OnlineStats.updatecounter!(o.wgt)
	λ = weight(o.wgt)

    # update the mean
    smooth!(o.x̄, x, λ)

    # copy and center x
    x = x - o.x̄
    # center!(x, o.x̄)

    #=
    TODO:
        - remove temporaries
        - handle ZCA whitening: for V'V = Σ⁻¹, D = diag(e)
            y_pca = D⁻¹/² * V'
            y_zca = V * y_pca
    =#

    n = min(o.nout, o.wgt.nobs)
	@inbounds for i in 1:n
        # TODO this could be more robust?
		if o.e[i] == zero(T)
			# initialize ith principal component
			Uᵢ = x
		else
			# update the ith eigvec/eigval
			Uᵢ = view(o.U, i, :)
			smooth!(Uᵢ, x * (dot(x, Uᵢ) / o.e[i]), λ)
        end
		eᵢ = norm(Uᵢ)
        # if eᵢ ≈ zero(T)
        #     warn("Got small e[$i] ($(eᵢ)) while computing iteration $(o.wgt.nobs) for:\n  $o")
        # end
		Vᵢ = Uᵢ / max(eᵢ, T(1e-8))

		# store these updates
		o.U[i,:] = Uᵢ
		o.V[i,:] = Vᵢ
		o.e[i] = eᵢ

        if i < n
    		# subtract projection on ith PC for the next step
    		x -= dot(x, Vᵢ) * Vᵢ
        end
	end
end

grad!(o::Whiten) = error("Can't backprop through Whiten")

# -------------------------------------------------------------------
# transform defs

# TODO: cut down on temporaries

function normal_pca{T}(o::Whiten{T})
    x = input_value(o) - o.x̄
    y = output_value(o)
    y[:] = o.V * x
    y
end

function whitened_pca{T}(o::Whiten{T})
    y = normal_pca(o)
    y[:] = y ./ sqrt.(o.e .+ T(1e-8))
    y
end

function zca{T}(o::Whiten{T})
    y = whitened_pca(o)
    y[:] = o.V' * y
    y
end

# "normal" pca... should match MultivariateStats.fit(PCA, x, maxdimout=k)
#   y = V (x - mean(x))
transform!{T}(o::Whiten{T,:pca}) = normal_pca(o)

# whitened pca... multiply by D⁻¹/², where D = diag(e) is a diagonal matrix
# with the eigenvalues e on the diagonal
#   yᵢ = D⁻¹/² V (x - mean(x))
transform!{T}(o::Whiten{T,:whitened_pca}) = whitened_pca(o)

# whitened zca... multiply by (V' * D⁻¹/²), where D = diag(e) is a diagonal matrix
# with the eigenvalues e on the diagonal.
# Note: ZCA is similar to whitened PCA, except that it rotates the whitened projection
#       back into "input space".  If you want whitening while maintaining the directions
#       of the inputs, this is the one to use.
#   yᵢ = D⁻¹/² V (x - mean(x))
transform!{T}(o::Whiten{T,:zca}) = zca(o)


include("rijke.jl")
using PyPlot
N = 18
u0Init = zeros(N)
u0Init[1] = 1.
s = [7.0, 0.2]
nRunUp = 40000
#u0 = Rijke(u0Init, s, nRunUp)
t = LinRange(-1.01,-0.99,500)
heat_sqrt = qfun_sqrt.(t)
heat_poly = qfun_poly.(t)
fig, ax = subplots(1,1)
ax.plot(t, heat_sqrt, label=L"$\dot{q}$, using sqrt", ".")
ax.plot(t, heat_poly, label=L"$\dot{q}$, using poly", ".")
ax.xaxis.set_tick_params(labelsize=18)
ax.yaxis.set_tick_params(labelsize=18)
ax.set_xlabel("delayed velocity",fontsize=18)
fig.legend(fontsize=18)


## characters defined by actions/homomorphisms

function _action_class_fun(
    conjugacy_cls::AbstractVector{CCl},
) where {CCl <: AbstractOrbit{<:PermutationGroups.AbstractPerm}}
    vals = Int[PermutationGroups.nfixedpoints(first(cc)) for cc in conjugacy_cls]
    # in general:
    # vals = [tr(matrix_representative(first(cc))) for cc in conjugacy_cls]
    return Characters.ClassFunction(vals, conjugacy_cls)
end

function _action_class_fun(
    conjugacy_cls::AbstractVector{CCl},
) where {CCl <: AbstractOrbit{<:AbstractMatrix}}
    vals = [tr(first(cc)) for cc in conjugacy_cls]
    return Characters.ClassFunction(vals, conjugacy_cls)
end

function _action_class_fun(
    hom::InducedActionHomomorphism{<:ByPermutations},
    conjugacy_cls
)
    vals = Int[PermutationGroups.nfixedpoints(induce(hom, first(cc))) for cc in conjugacy_cls]
    # in general:
    # vals = [tr(matrix_representative(first(cc))) for cc in conjugacy_cls]
    return Characters.ClassFunction(vals, conjugacy_cls)
end

function _action_class_fun(
    hom::InducedActionHomomorphism{<:ByLinearTransformation},
    conjugacy_cls,
)
    vals = [tr(induce(hom, first(cc))) for cc in conjugacy_cls]
    return Characters.ClassFunction(vals, conjugacy_cls)
end

"""
    action_character(conjugacy_cls, tbl::CharacterTable)
Return the character of the representation given by the elements in the
conjugacy classes `conjugacy_cls`.

This corresponds to the classical definition of characters as a traces of the
representation matrices.
"""
function action_character(conjugacy_cls, tbl::CharacterTable)
    ac_char = _action_class_fun(conjugacy_cls)
    constituents = Int[dot(ac_char, χ) for χ in irreducible_characters(tbl)]
    return Character(tbl, constituents)
end

"""
    action_character(hom::InducedActionHomomorphism, tbl::CharacterTable)
Return the character of the representation given by the images under `hom` of
elements in `conjugacy_classes(tbl)`.

This corresponds to the classical definition of characters as a traces of the
representation matrices.
"""
function action_character(hom::InducedActionHomomorphism, tbl::CharacterTable)
    ac_char = _action_class_fun(hom, conjugacy_classes(tbl))
    constituents = Int[dot(ac_char, χ) for χ in irreducible_characters(tbl)]
    return Character(tbl, constituents)
end