# Exported functions

# TODO:
# Add tests for the sensibleness of mi and cmi
# Offer alternative target
# Look into standardising the documentation
# Implement (or find implementation of) Bayesian blocks in Julia,
# 	to not be dependent on the Python module

export get_shrinkage_entropy, get_shrinkage_mi, get_shrinkage_cmi, get_shrinkage_total_correlation

"""
Calculates a James-Stein shrinkage estimate for the entropy of a set
of observed values, via the following steps:
(1) convert the observed values to frequencies of discrete bins
(2) convert the frequencies to probablities using shrinkage estimation
(3) run the probabilities through the entropy formula.
Can calculate joint entropy for up to 3 sets of observed values.

Parameters:

valuesX - dxn Array{Float64,2} - The observed values, where d is
the number of dimensions and n is the number of values.

[valuesY] - dxn Array{Float64,2} - The observed values, where d is
the number of dimensions and n is the number of values.

[valuesZ] - dxn Array{Float64,2} - The observed values, where d is
the number of dimensions and n is the number of values.

Keyword arguments:

[base=2] - the base of the logarithm, which determines the units
[lambda=nothing] - the shrinkage intensity, between 0 and 1 inclusive, which
		will be calculated if this value is not given
[mode="uniformwidth"] - the discretization mode: "uniformwidth", "uniformcount" or
		"bayesianblocks"
"""
function get_shrinkage_entropy(valuesX::Array{Float64,2}; base::Number = 2, lambda = nothing, mode::ASCIIString = "uniformwidth")
	probabilities = getprobabilitiesshrinkage(getfrequencies(valuesX, mode), lambda)
	return applyentropyformula(probabilities, base)
end
function get_shrinkage_entropy(valuesX::Array{Float64,2}, valuesY::Array{Float64,2}; base::Number = 2, lambda = nothing, mode::ASCIIString = "uniformwidth")
	jointprobabilities = getprobabilitiesshrinkage(getjointfrequencies(valuesX, valuesY, mode), lambda)
	return applyentropyformula(jointprobabilities, base)
end
function get_shrinkage_entropy(valuesX::Array{Float64,2}, valuesY::Array{Float64,2}, valuesZ::Array{Float64,2}; base::Number = 2, lambda = nothing, mode::ASCIIString = "uniformwidth")
	jointprobabilities = getprobabilitiesshrinkage(getjointfrequencies(valuesX, valuesY, valuesZ, mode), lambda)
	return applyentropyformula(jointprobabilities, base)
end

"""
Calculates a James-Stein shrinkage estimate for the mutual information
between two sets of observed values.

Parameters:

valuesX - dxn Array{Float64,2} - A set of observed values, where d is
the number of dimensions and n is the number of values.

valuesY - dxn Array{Float64,2} - Another set of observed values, where
d is the number of dimensions and n is the number of values.

Keyword arguments:

[base=2] - the base of the logarithm, which determines the units
[lambda=nothing] - the shrinkage intensity, between 0 and 1 inclusive, which
		will be calculated if this value is not given
[mode="uniformwidth"] - the discretization mode: "uniformwidth", "uniformcount" or
		"bayesianblocks"
"""
# NB "mi" stands for mutual information
function get_shrinkage_mi(valuesX::Array{Float64,2}, valuesY::Array{Float64,2}; base::Number = 2, lambda = nothing, mode::ASCIIString = "uniformwidth")
	entropyX = get_shrinkage_entropy(valuesX, base = base, lambda = lambda, mode = mode)
	entropyY = get_shrinkage_entropy(valuesY, base = base, lambda = lambda, mode = mode)
	entropyXY = get_shrinkage_entropy(valuesX, valuesY, base = base, lambda = lambda, mode = mode)
	return applymutualinformationformula(entropyX, entropyY, entropyXY)
end

"""
Calculates a James-Stein shrinkage estimate for the conditional mutual
information between two sets of observed values, given a third set of
observed values.

Parameters:

valuesX - dxn Array{Float64,2} - A set of observed values, where d is
the number of dimensions and n is the number of values.

valuesY - dxn Array{Float64,2} - Another set of observed values, where
d is the number of dimensions and n is the number of values.

valuesZ - dxn Array{Float64,2} - The set of observed values, that will
be conditioned on, where d is the number of dimensions and n is the
number of values.

Keyword arguments:

[base=2] - the base of the logarithm, which determines the units
[lambda=nothing] - the shrinkage intensity, between 0 and 1 inclusive, which
		will be calculated if this value is not given
[mode="uniformwidth"] - the discretization mode: "uniformwidth", "uniformcount" or
		"bayesianblocks"
"""
# NB "cmi" stands for conditional mutual information
function get_shrinkage_cmi(valuesX::Array{Float64,2}, valuesY::Array{Float64,2}, valuesZ::Array{Float64,2}; base::Number = 2, lambda = nothing, mode::ASCIIString = "uniformwidth")
	entropyZ = get_shrinkage_entropy(valuesZ, base = base, lambda = lambda, mode = mode)
	entropyXZ = get_shrinkage_entropy(valuesX, valuesZ, base = base, lambda = lambda, mode = mode)
	entropyYZ = get_shrinkage_entropy(valuesY, valuesZ, base = base, lambda = lambda, mode = mode)
	entropyXYZ = get_shrinkage_entropy(valuesX, valuesY, valuesZ, base = base, lambda = lambda, mode = mode)
	return applyconditionalmutualinformationformula(entropyZ, entropyXZ, entropyYZ, entropyXYZ)
end

function get_shrinkage_total_correlation(valuesX::Array{Float64,2}, valuesY::Array{Float64,2}, valuesZ::Array{Float64,2}; base::Number = 2, lambda = nothing, mode::ASCIIString = "uniformwidth")
	entropyX = get_shrinkage_entropy(valuesX, base = base, lambda = lambda, mode = mode)
	entropyY = get_shrinkage_entropy(valuesY, base = base, lambda = lambda, mode = mode)
	entropyZ = get_shrinkage_entropy(valuesZ, base = base, lambda = lambda, mode = mode)
	entropyXYZ = get_shrinkage_entropy(valuesX, valuesY, valuesZ, base = base, lambda = lambda, mode = mode)
	return applytotalcorrelationformula(entropyX, entropyY, entropyZ, entropyXYZ)
end


function analysis_functions(c::Coupler)
    # Extract number of branches
    n = length(c.a)

    # Form ABCD matrix
    M = _ABCD_open_stub(c.a[1])
    for i in eachindex(c.b)
        M *= (_ABCD_line(c.b[i]) * _ABCD_open_stub(c.a[i+1]))
    end
    A_ = M[1,1]
    B_ = M[1,2]
    C_ = M[2,1]
    D_ = M[2,2]

    # Γₑ_N
    Γₑ_N = (A_ - D_) + (B_ - C_)

    # Γₑ_D
    Γₑ_D = (A_ + D_ + B_ + C_)

    # Tₑ_N
    Tₑ_N = 2*(Polynomial([1, 0, -1], :t)^(n-1))

    # Create functions
    # Reflection coefficient
    Γ(t) = (1/2) * (Γₑ_N(t)/Γₑ_D(t) + Γₑ_N(1/t)/Γₑ_D(1/t))

    # Directivity
    D(t) = (1/2) * (Γₑ_N(t)/Γₑ_D(t) - Γₑ_N(1/t)/Γₑ_D(1/t))

    # Insertion Loss
    IL(t)  = (1/2) * (Tₑ_N(t)/Γₑ_D(t) + (-1)^(n) * Tₑ_N(1/t)/Γₑ_D(1/t))

    # Coupling
    C(t) = (1/2) * (Tₑ_N(t)/Γₑ_D(t) - (-1)^(n) * Tₑ_N(1/t)/Γₑ_D(1/t))

    return (Γ = Γ, D = D, IL = IL, C = C)
end

db(x) = 10*log10(abs2(x))

function _ABCD_open_stub(a)
    A = Polynomial(Float64[1], :t)
    B = Polynomial(Float64[0], :t)
    C = Polynomial(Float64[0, a], :t)
    D = Polynomial(Float64[1], :t)
    return [A B;C D]
end

function _ABCD_line(b)
    A = Polynomial(Float64[1, 0, 1], :t)
    B = Polynomial(Float64[0, 2/b], :t)
    C = Polynomial(Float64[0, 2*b], :t)
    D = Polynomial(Float64[1, 0, 1], :t)
    return [A B;C D]
end


##
## Transformations are
## * givens rotations  Find [c s; -conj(s) conj(c)][a,b] = [x 0];
## * fuse U,V -> UV or UV,D (when complex real)
## * turnover: [   [ -->  [
##               [      [   [
## to this we add abstract "passthrough!" functions

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

"""
    c,s,r  = givensrot(a,b)

where `[c s; -s conj(c)] * [a,b] = [r, 0]
"""
@inline function givensrot(a::T, b::T) where {T <: Real}
    G, r = givens(a,b,1,2)
    s = r >= 0 ? one(T) : -one(T)   #sign(r)
    s*G.c, s*G.s, s*r
end

## This leaves s real
## a, b= rand(Complex{Float64},2)
## G,r = givens(b,a, 1,2)
## c, s = G.s, real(G.c)
## [c s; -s conj(c)] * [a,b]
function givensrot(a::Complex{T},b::Complex{T}) where {T <: Real}
    G, r = givens(b, a, 1, 2)
    G.s, real(G.c), r
end

givensrot(a::Complex{T},b::T) where {T <: Real} =  givensrot(a, complex(b, zero(T)))

## This is suggested in the book. A bit slower, but only one function needed...
## function givensrotX(a::S, b::T) where {T <: Real, S <: Union{T, Complex{T}}}
##     m = max(norm(a), norm(b))
##     x,y = a/m, b/m
##     r = sqrt(norm(x)^2 + y^2)
##     c = x/r
##     s = y/r
##     r = m*r
##     return conj(c),real(s),r
## end

# avoid square root when norm is near 1
#
# 5.1. Backward Stability of the Turnover. Suggests
# this modification to c,s to ensure unitarity
# Assuming that c and s have been produced by valid formulas, the
# magnitude of the correction will be extremely tiny, on the order of
# the unit roundoff. This seemingly trivial correction is crucial. If it
# is neglected, the core transformations will gradually drift away from
# unitarity over the course of many operations, and the algorithm will
# fail. If it is not neglected, the programmer has a good chance of
# producing an accurate turnover.
function approx_givensrot(a::S,b::T) where {S,T}
    delta::T = abs(a*conj(a)) + b^2 - one(T)
    r::T = one(T) - delta/(one(T)+one(T))  # taylor approx for 1/sqrt(1 + delta) when delta ~ 0
    c = a*r
    s = b*r
    return conj(c),s

end

## extra polish; not used
function polish_givens(c::S, s::T) where {S, T}

    delta = real(conj(c)*c) + s*s - one(T)
    r = one(T) -  delta/(one(T) + one(T))
    r * c, r * s

end

##################################################
## Fuse
## fuse combines two rotations, a and b, into one,
## For general rotation, we have no phase consideration, so return an identity rotator
@inline function fuse(a::AbstractRotator{T,S}, b::AbstractRotator{T,S}) where {T,S <: Real}
    #    idx(a) == idx(b) || error("can't fuse")
    i = idx(a)

    c1, s1 = vals(a)
    c2, s2 = vals(b)

    u = c1 * c2 - s1 * conj(s2)
    v = c1 * s2 + s1 * conj(c2)

    u,v = polish_givens(u,v)
    return Rotator(u,v,i), IdentityRotator{T, S}(i)

end

## U * V -> UV [ alpha 0; 0 conj(alpha)]

## For ComplexRealRotator, the result of a*b will not have a real sign
## we output by rotating by alpha.
## We have U*V = (UV) * D
## To get D * (UV) a flip is needed
@inline function fuse(a::Rotator{T,S}, b::Rotator{T,S}) where {T, S <: Complex{T}}

    i = idx(a)
    #    idx(a) == idx(b) || error("can't fuse")
    c1, s1 = vals(a)
    c2, s2 = vals(b)

    u = c1 * c2 - s1 * s2
    v = c1*s2 + s1 * conj(c2)

    s = norm(v)
    alpha =  iszero(v) ? one(Complex{T}) : v/s
    c = u * alpha

    c, s = polish_givens(c, s)

    Rotator(c, s, i), DiagonalRotator(conj(alpha), i)
end

## Fuse for diagonal matrices
fuse(D::IdentityRotator, U) = U, D
function fuse(D::DiagonalRotator, U)
    alpha,_ = vals(D)
    i = idx(D)
    @assert i == idx(U)
    c, s = vals(U)
    Rotator(c*alpha, s, i)
end

##################################################
# Turnover: Q1    Q3   | x x x |      U1
#              Q2    = | x x x | = U3    U2
#                      | x x x |
#
# This is the key computation once matrices are written as rotators
# This function returns
# c4,s4,c5,s5,c6,s6

# following pp15-17 in book
# and the paper
# M = U    W
#       V
# we find U1 by computing a givens rotation to clear [3,1]
#
#     * M =| a x x |
# U1'      | b x x |
#          | 0 x x |
# Then use fact that a^2 + b^2 ~ 1 to use approx givens to find V1 to clear [2,1]
# V1' *     * M  = | 1  0 0   |
#       U1'        | 0  c s   |
#                  | 0 -s cbar|
# The 1 in [1,1] must be 1 and not -1;  so the  lower part is a rotator
# from here [2:3,2:3] is unitary, we pick [3,3] for c and use
# the invariant s1*s2 = s5*s6 to find s6 unless s5=0
#
# question: polishing slows things down, and doesn't seem necessary,
# as approx_givensrot does this; but does make Wilkinson case a bit better
@inline function _turnover(c1::S, s1::T, c2::S, s2::T, c3::S, s3::T) where {S,T}

    ## UVW has first row [UVW11, UVW21, UVW31]
    UVW21 = -c2 * s3 * conj(c1) - c3*s1
    UVW31 = s2 * s3

    if iszero(UVW21) && iszero(UVW31)
        if iszero(s2*conj(c3)) ## this is the case of diagonal matrix Q1*Q2*Q3
            c4 = one(S)
            c5 = c1*c3 - c2*s1*s3
            c6 = c2
            return c4, zero(T), c5, zero(T), c6, zero(T)
        end
    end

    c4, s4, r4 = givensrot(UVW21, UVW31)

    c4, s4 = conj(c4), -s4  # conjugate, as [c4 s4; -s4 conj(s4)]*[a,b] = [r,0]
#    c4, s4 =  polish_givens(c4, s4)

    UVW11 = c1*c3 - c2*s1*s3
    c5, s5::T = approx_givensrot(UVW11, real(r4))
    c5, s5 = conj(c5), -s5
#    c5, s5 =  polish_givens(c5, s5)

    ##  use s1*s2 = s5 * s6 if we can
    ## get a from  M  = V1' * U1' * U *  V  *  W; a = M[3,3]
    if !iszero(s5)
        a = conj(c1) * s2 * s4 + conj(c2) * c4
        b = s1*s2/s5
        c6, s6 = approx_givensrot(a, b)
    else
        a = c2*conj(c1)
        b = -s2
        c6, s6 = givensrot(a, b)
    end
#    c6, s6 = approx_givensrot(a, b)
#    c6, s6 =  polish_givens(c6, s6)

    return (c4, s4, c5, s5, c6, s6)

end

##
##  Turnover interface for rotators
##
function turnover(Q1::R1t,
                  Q2::R2t,
                  Q3::R3t) where {R1t, R2t, R3t}

    c1, s1 = vals(Q1); c2, s2 = vals(Q2); c3,s3 = vals(Q3)
    i,j,k = idx(Q1), idx(Q2), idx(Q3)

    # @assert i == k && (abs(j-i) == 1)

    c4,s4,c5,s5,c6,s6 = _turnover(c1,s1, c2,s2, c3,s3)

    R1 = R3t(c4, s4, j)
    R2 = R2t(c5, s5, i)
    R3 = R1t(c6, s6, j)

    R1, R2, R3

end


function solve(prob::FVMPDEProblem;kwargs...)
    prob.t = 0
    if haskey(kwargs,:scheme)
        scheme = kwargs[:scheme]
        kwargs = Dict([p for p in pairs(kwargs) if p[1] != :scheme])
    else
        scheme = uw1
    end

    if haskey(kwargs,:showstep)
        dt_show = kwargs[:showstep]
        kwargs = Dict([p for p in pairs(kwargs) if p[1] != :showstep])
    else
        dt_show = Inf
    end

    U0 = copy(prob.U0)

    ODE_Function(u,p,t) = FVMPDE_∂u∂t(prob,u,t,scheme,dt_show)
    ODE_Problem = ODEProblem(ODE_Function,U0,prob.tspan;prob.ode_problem_kwargs...);

    if haskey(kwargs,:algorithm)
        alg = kwargs[:algorithm]
        kwargs = Dict([p for p in pairs(kwargs) if p[1] != :algorithm])
        sol = DifferentialEquations.solve(ODE_Problem,alg;kwargs...)
    else
        sol = DifferentialEquations.solve(ODE_Problem;kwargs...)
    end

    return sol
end

function FVMPDE_∂u∂t(prob::FVMPDEProblem,u,t,scheme,dt_show)
    if t - prob.t > dt_show
        println(t)
        prob.t = t
    end
    θ = prob.param
    indices = prob.grid.indices
    Δx = prob.grid.dX
    Δx_c = prob.grid.dX_C
    ∂F∂x = FVMPDE_∂F∂x(scheme,prob.F,u,prob.grid,1,t,θ,prob.jac)

    if prob.grid.dimension == _2D
        Δy = prob.grid.dY[1]
        ∂G∂y = FVMPDE_∂F∂x(scheme,prob.G,u,prob.grid,2,t,θ,prob.jac)
        ∂H∂z = 0
    elseif prob.grid.dimension == _3D
        Δy = prob.grid.dY[1]
        Δz = prob.grid.dZ[1]
        ∂G∂y = FVMPDE_∂F∂x(scheme,prob.G,u,prob.grid,2,t,θ,prob.jac)
        ∂H∂z = FVMPDE_∂F∂x(scheme,prob.H,u,prob.grid,3,t,θ,prob.jac)
    else
        ∂G∂y = 0
        ∂H∂z = 0
    end

    S = prob.S(u)

    ∂u∂t = S .- (∂F∂x .+ ∂G∂y .+ ∂H∂z)

    return ∂u∂t
end

function FVMPDE_∂F∂x(scheme,func_F,u,grid::FVMPDEGrid,_dim,t,p,jac)
    Uₓᴸ, Uₓᴿ = GeneralFluxLimiter(scheme,u,grid,_dim)

    ∂f∂x = zeros(size(u))

    for i in grid.indices
        A⁺ , A⁻ = LRJacobian(func_F,u[i,:],t,p,jac)
        ∂f∂x[i,:] = (A⁻ * Uₓᴸ[i,:] .+ A⁺ * Uₓᴿ[i,:])
    end
    return ∂f∂x
end

function LRJacobian(func_F,u,t,p,jac)
    if jac
        A = ForwardDiff.jacobian(func_F, u)
    else
        A = func_F(u)
    end

    if size(A,1) == 1
        A = reshape(A,1,1)
        A⁺ = 0.5 .* (A .+ abs.(A))
        A⁻ = 0.5 .* (A .- abs.(A))
    else
        Λ = diagm(eigvals(A))
        @assert typeof(Λ) === Matrix{Float64}
        Λ⁺ = 0.5 .* (Λ .+ abs.(Λ))
        Λ⁻ = 0.5 .* (Λ .- abs.(Λ))
        P = eigvecs(A)
        P⁻¹ = inv(P)
        A⁺ = P * Λ⁺ * P⁻¹
        A⁻ = P * Λ⁻ * P⁻¹

    end

    return A⁺, A⁻
end


#Gravity model of commuting in Kenya

#Calculate interaction force for gravity model --- assuming Urban pop
function TupleDist(x1::Tuple{Float64,Float64},x2::Tuple{Float64,Float64})
    return sqrt((x1[1]-x2[1])^2 + (x1[2]-x2[2])^2 )
end

function get_transport_matrix(filename)
    KenyaTbl,n = get_Kenyadata(filename)
    transport_matrix = zeros(n,n)
    for i = 1:n
        for j = 1:n
            if i != j
            transport_matrix[i,j] = KenyaTbl[:Urban][i]*KenyaTbl[:Urban][j]/(TupleDist(KenyaTbl[:Loc][i],KenyaTbl[:Loc][j])^2)
        end
        end
    end
    #Normalise the columns
    for j = 1:n
        transport_matrix[:,j] = transport_matrix[:,j]/sum(transport_matrix[:,j])
    end
    return transport_matrix,n
end

function get_transport_matrix(KenyaTbl,n)
    transport_matrix = zeros(n,n)
    for i = 1:n
        for j = 1:n
            if i != j
            transport_matrix[i,j] = KenyaTbl[:Urban][i]*KenyaTbl[:Urban][j]/(TupleDist(KenyaTbl[:Loc][i],KenyaTbl[:Loc][j])^2)
        end
        end
    end
    #Normalise the columns
    for j = 1:n
        transport_matrix[:,j] = transport_matrix[:,j]/sum(transport_matrix[:,j])
    end
    return transport_matrix
end


__precompile__()

"""
    MMD Kernels
"""
module Kernels

export expect, RBFKernel, loss

abstract type AbstractKernel end

struct RBFKernel <: AbstractKernel
    sigmas::Vector{Float64}
    matrix::Matrix{Float64}
end

function RBFKernel(nqubits::Int, sigmas::Vector{Float64}, isbinary::Bool)
    basis = collect(0:(1<<nqubits - 1))
    return RBFKernel(sigmas, rbf_kernel_matrix(basis, basis, sigmas, isbinary))
end

expect(kernel::RBFKernel, px::Vector{Float64}, py::Vector{Float64}) = px' * kernel.matrix * py

# RBF Kernel

function rbf_kernel_matrix(x::Vector, y::Vector, sigmas::Vector{Float64}, isbinary::Bool)
    if length(sigmas) == 0
        throw(DimensionMismatch("At least 1 sigma prameters is required for RBM kernel!"))
    end

    if isbinary
        dx2 = map(count_ones, xor.(x, y'))
    else
        dx2 = (x .- y').^2
    end

    K = 0
    for sigma in sigmas
        gamma = 1.0 / (2 * sigma)
        K = K + exp.(-gamma * dx2)
    end
    return K
end

function loss(px::AbstractVecOrMat{Float64}, kernel::RBFKernel, py::AbstractVecOrMat{Float64})
    pxy = px - py
    return expect(kernel, pxy, pxy)
end

end


mutable struct LeastSquares{T} <: AbstractLoss
    A::Matrix{T}
    b::Vector{T}
    temp::Vector{T}

    function (::Type{LeastSquares})(A::Matrix, b::Vector)
        T = promote_type(eltype(Q), eltype(b))
        new{T}(A, b)
    end

    function (::Type{LeastSquares})(N::Integer, d::Integer)
        @assert d ≥ 1 "d must be at least 1"
        @assert N ≥ d "N must be at least $(d)"

        A = randn(N, d)
        A *= 2/norm(A)
        b = Float64[rand() ≤ 0.5 ? -1. : 1. for idx in 1:N]
        T = promote_type(eltype(A), eltype(b))
        new{T}(A, b, copy(b))
    end
end

nfeatures(loss::LeastSquares) = size(loss.A,2)
nsamples(loss::LeastSquares) = size(loss.A,1)

function loss!(loss::LeastSquares{T}, x::Vector{T}, g::Vector{T}) where T
    mul!(loss.temp,loss.A,x)
    result = convert(T, 0.5) * norm(loss.temp-loss.b,2)^2
    mul!(g, adjoint(loss.A), loss.temp)
    return result
end


using MittagLefflerFunctions
using PyPlot
using SpecialFunctions: erfcx
using Printf

α = 1/2
β = 1.0

x = range(0, 10, length=201)
y = erfcx.(x)

rvals = collect(4:2:10)
maxerr = Float64[]

figure(1)
@printf("%5s  %5s  %10s\n\n", "m", "n", "max error")
style = ("-", "--", ":", "-.")
for row = 1:length(rvals)
    r = rvals[row]
    m = r + 1
    n = r
    E = MLPade(α, β, m, n, :svd)
    err = abs.( y - E.(x) )
    semilogy(x, err, style[row], label="m, n = $m, $n")
    push!(maxerr, maximum(err))
    @printf("%5d  %5d  %10.2e\n", m, n, maxerr[row])
end
legend()
grid(true)
xlabel(L"$x$")
savefig("figure09.pdf")


# SKCE

# predicted Laplace distributions with exponential kernel for the targets
function CalibrationErrors.unsafe_skce_eval_targets(
    kernel::ExponentialKernel{Euclidean}, p::Laplace, y::Real, p̃::Laplace, ỹ::Real
)
    # extract the parameters
    μ = p.μ
    β = p.θ
    μ̃ = p̃.μ
    β̃ = p̃.θ

    res =
        kernel(y, ỹ) - laplace_laplacian_kernel(β, abs(μ - ỹ)) -
        laplace_laplacian_kernel(β̃, abs(μ̃ - y)) +
        laplace_laplacian_kernel(β, β̃, abs(μ - μ̃))

    return res
end

# z = abs(μ - y)
function laplace_laplacian_kernel(β, z)
    if isone(β)
        (1 + z) * exp(-z) / 2
    else
        (β * exp(-z / β) - exp(-z)) / (β^2 - 1)
    end
end

# z = abs(μ - μ̃)
function laplace_laplacian_kernel(β, β̃, z)
    if isone(β)
        if isone(β̃)
            return (3 + 3 * z + z^2) * exp(-z) / 8
        else
            c = β̃^2 - 1
            csq = c^2
            return β̃^3 * exp(-z / β̃) / csq - ((1 + z) / (2 * c) + β̃^2 / csq) * exp(-z)
        end
    end

    if isone(β̃)
        c = β^2 - 1
        csq = c^2
        return β^3 * exp(-z / β) / csq - ((1 + z) / (2 * c) + β^2 / csq) * exp(-z)
    elseif β̃ == β
        c = β^2 - 1
        csq = c^2
        return exp(-z) / csq + ((β + z) / (2 * c) - β / csq) * exp(-z / β)
    else
        c1 = β^2 - 1
        c2 = β̃^2 - 1
        c3 = β^2 - β̃^2
        return β^3 * exp(-z / β) / (c1 * c3) - β̃^3 * exp(-z / β̃) / (c2 * c3) +
               exp(-z) / (c1 * c2)
    end
end

# UCME

function CalibrationErrors.unsafe_ucme_eval_targets(
    kernel::ExponentialKernel{Euclidean}, p::Laplace, y::Real, ::Laplace, testy::Real
)
    return kernel(y, testy) - laplace_laplacian_kernel(p.θ, abs(p.μ - testy))
end

# kernels with input transformations
# TODO: scale upfront?
function CalibrationErrors.unsafe_skce_eval_targets(
    kernel::TransformedKernel{
        ExponentialKernel{Euclidean},<:Union{ScaleTransform,ARDTransform}
    },
    p::Laplace,
    y::Real,
    p̃::Laplace,
    ỹ::Real,
)
    # obtain the transform
    t = kernel.transform

    return CalibrationErrors.unsafe_skce_eval_targets(
        ExponentialKernel(), apply(t, p), t(y), apply(t, p̃), t(ỹ)
    )
end

function CalibrationErrors.unsafe_ucme_eval_targets(
    kernel::TransformedKernel{
        ExponentialKernel{Euclidean},<:Union{ScaleTransform,ARDTransform}
    },
    p::Laplace,
    y::Real,
    testp::Laplace,
    testy::Real,
)
    # obtain the transform
    t = kernel.transform

    # `testp` is irrelevant for the evaluation and therefore not transformed
    return CalibrationErrors.unsafe_ucme_eval_targets(
        ExponentialKernel(), apply(t, p), t(y), testp, t(testy)
    )
end

# utilities

# internal `apply` avoids type piracy of transforms
apply(t::Union{ScaleTransform,ARDTransform}, d::Laplace) = Laplace(t(d.μ), t(d.θ))


m = @model x0 begin
    μ ~ HalfNormal(1)
    σ ~ HalfCauchy(1)
    x ~ Normal(x0 + μ, σ)
end

m_fwd = m(:μ,:σ)
m_inv = m(:x)

r = makeRand(m_fwd)

grt = (μ = 0.1, σ = 1.2735923758662662, x0=0.0)

function realize(n; pars...)
    pars = Dict(pars)
    xs = zeros(n)
    for j in 1:n
        newx = r(;pars...).x
        xs[j] = newx
        pars = merge(pars, Dict(:x0=>newx))

    end
    xs
end

obs = realize(1000;grt...)

using Plots
pyplot()
plot(obs)


f(a,b) = 333.75*b^6+a^2*(11*a^2*b^2 - b^6 - 121*b^4 - 2) + 5.5*b^8 + a/(2*b)

a=77617.0
b=33096.0

println("Float64: ", f(a, b))
println("BigFloat: ", f(BigFloat(a), BigFloat(b)))


### log pdf of a multivariate normal
function logpfun_mv(x, μ, σ, w)
    logw = log.(w)
    K = length(logw)
    N = size(x, 2)
    sum(logsumexp(logw[k] + logpdf(MvNormal(μ[k], σ), x[:,n]) for k in 1:K) for n in 1:N)
end

@model function (x, nc, nsites, σ,  p,  wc=ones(nc)/nc, wp=ones(nsites-1)/(nsites-1), w=wc.*wp', q = 2)

    n = size(x, 2)
    startpos ~ filldist(Uniform(-1500, 1500), nc)
    delta    ~ filldist(Gamma(q*p, 1/q), nc, nsites)
    configs = startpos .+ cumsum(delta, dims=2)

    μ = [[configs[i, j], configs[i, j+1]] for i = 1:nc, j=1:(nsites-1)]

    #x ~ filldist(MixtureModel(Normal.(vec(μ), w)), n)
    Turing.@addlogprob! logpfun_mv(x, vec(μ), σ, w)
end


include("hyperelasticity.jl")

ts = ExplicitSolverTime(Δt, NT)
ubd, abd = compute_boundary_info(domain, globaldata, ts)
Fext = compute_external_force(domain, globaldata, ts)

d0 = zeros(2domain.nnodes)
v0 = zeros(2domain.nnodes)
a0 = zeros(2domain.nnodes)

# general neural network
function nn_law(ε)
  ae(ε, [20,20,20,3])
end

# # free energy
# function nn_law(ε)
#   φ = squeeze(ae(ε, [20,20,20,1]))
#   tf.gradients(φ, ε)
# end

# # parameter
# c1 = Variable(0.5) * 1e-1
# c2 = Variable(0.5) * 1e-1
# function nn_law(ε)
#   compute_stress_rivlin_saunders(ε, c1, c2)
# end

d, v, a= ExplicitSolver(globaldata, domain, d0, v0, a0, Δt, NT, nn_law, Fext, ubd, abd; strain_type="finite")

idx = @. (0:n)*m + 1
idx = [idx; (@. idx + domain.nnodes)]
loss = sum((d[:, idx] - d_[:, idx])^2)
sess = Session(); init(sess)
@info run(sess, loss)
# error()

for i = 1:100
  BFGS!(sess, loss, 1000)
  ADCME.save(sess, "data1_$i.mat")
end


function [D,J,X]=JacobiCart1(ksi,F,Grid)

x1=Grid.Nodes(F.N(1)).P(1);
y1=Grid.Nodes(F.N(1)).P(2);
x2=Grid.Nodes(F.N(2)).P(1);
y2=Grid.Nodes(F.N(2)).P(2);
x3=Grid.Nodes(F.N(3)).P(1);
y3=Grid.Nodes(F.N(3)).P(2);
x4=Grid.Nodes(F.N(4)).P(1);
y4=Grid.Nodes(F.N(4)).P(2);
xBar=ksi(1);
yBar=ksi(2);

J=[ (x1 - x2 + x3 - x4)*yBar + x2 - x1, (x1 - x2 + x3 - x4)*xBar + x4 - x1
     (y1 - y2 + y3 - y4)*yBar + y2 - y1, (y1 - y2 + y3 - y4)*xBar + y4 - y1];
J=[J
   0 0];
D=J(1,1)*J(2,2)-J(1,2)*J(2,1);
if nargout > 2
  X=zeros(1,3);
  X(1)=Grid.Nodes(F.N(1)).P(1)...
    +(Grid.Nodes(F.N(2)).P(1)-Grid.Nodes(F.N(1)).P(1))*ksi(1)...
    +(Grid.Nodes(F.N(4)).P(1)-Grid.Nodes(F.N(1)).P(1))*ksi(2)...
    +(Grid.Nodes(F.N(3)).P(1)-Grid.Nodes(F.N(4)).P(1)-Grid.Nodes(F.N(2)).P(1)+Grid.Nodes(F.N(1)).P(1))*ksi(1)*ksi(2);
  X(2)=Grid.Nodes(F.N(1)).P(2)...
    +(Grid.Nodes(F.N(2)).P(2)-Grid.Nodes(F.N(1)).P(2))*ksi(1)...
    +(Grid.Nodes(F.N(4)).P(2)-Grid.Nodes(F.N(1)).P(2))*ksi(2)...
    +(Grid.Nodes(F.N(3)).P(2)-Grid.Nodes(F.N(4)).P(2)-Grid.Nodes(F.N(2)).P(2)+Grid.Nodes(F.N(1)).P(2))*ksi(1)*ksi(2);
  X(3)=Grid.Nodes(F.N(1)).P(3)...
    +(Grid.Nodes(F.N(2)).P(3)-Grid.Nodes(F.N(1)).P(3))*ksi(1)...
    +(Grid.Nodes(F.N(4)).P(3)-Grid.Nodes(F.N(1)).P(3))*ksi(2)...
    +(Grid.Nodes(F.N(3)).P(3)-Grid.Nodes(F.N(4)).P(3)-Grid.Nodes(F.N(2)).P(3)+Grid.Nodes(F.N(1)).P(3))*ksi(1)*ksi(2);
end
end


immutable Hyperboloid{T} <: Model
    a::T
    α::T
    μ::T
end

Base.show(io::IO, ::Type{Hyperboloid}) = print(io, "Classical Hyperboloid")

function Base.show(io::IO, m::Hyperboloid)
    @printf(
        io,
        "Classical Hyperboloid (a = %.3f, α = %.3f, μ = %.3f)",
        m.a,
        m.α,
        m.μ,
    )
end

function initial_parameters(::Type{Hyperboloid})
    return Float64[log(1.0), log(1.0), log(1.0)]
end

canonical_parameters(m::Hyperboloid) = (m.a, m.α, m.μ)

function Hyperboloid(parameters::Vector)
    a = exp(parameters[1])
    α = exp(parameters[2])
    μ = exp(parameters[3])
    return Hyperboloid(a, α, μ)
end

@inline function predict(
    m::Hyperboloid,
    x2::Real,
    t2::Real,
    x1::Real,
    t1::Real,
)
    a, α, μ = m.a, m.α, m.μ
    u1 = x1 * (1 + α * t1)^(-μ)
    u2 = x2 * (1 + α * t2)^(-μ)
    z = a * (u2 - u1)
    p = invlogit(z)
    return p
end

@inline function gradient_component(
    m::Hyperboloid,
    x2::Real,
    t2::Real,
    x1::Real,
    t1::Real,
)
    a, α, μ = m.a, m.α, m.μ
    u1 = x1 * (1 + α * t1)^(-μ)
    u2 = x2 * (1 + α * t2)^(-μ)
    gr = (
        (u2 - u1) * a,
        a * (x2 * -μ * (1 + α * t2)^(-μ - 1) * t2 - x1 * -μ * (1 + α * t1)^(-μ - 1) * t1) * α,
        a * (-x2 * (1 + α * t2)^(-μ) * log(1 + α * t2) + x1 * (1 + α * t1)^(-μ) * log(1 + α * t1)) * μ,
    )
    return gr
end


using LinearAlgebra
function mwe()
  local Q::Matrix{Float64}

  #Q = Matrix{Float64}(undef, 800, 800)
  X = rand(200,13)
  f = qr(X)
  @time Q = Matrix(f.Q[:,:])

  X = rand(400,13)
  f = qr(X)
  @time Q = Matrix(f.Q[:,:])

  X = rand(600,13)
  f = qr(X)
  @time Q = Matrix(f.Q[:,:])

  X = rand(800,13)
  f = qr(X)
  @time Q = Matrix(f.Q[:,:])
end

@time mwe()


using SymPy

mu, eta, dt, lambda = @vars mu eta dt lambda
S = [
    1+2/3*mu*dt/eta -1/3*mu*dt/eta 0.0
    -1/3*mu*dt/eta 1+2/3*mu*dt/eta 0.0
    0.0 0.0 1+mu*dt/eta
]
invS = inv(S)
H = invS * [
    2*mu+lambda lambda 0.0
    lambda 2mu+lambda 0.0
    0.0 0.0 mu
]
S = simplify(invS)
H = simplify(H)

k = -1
s = ""
for i = 1:3
    for j = 1:3
        global k += 1
        c = sympy.ccode(S[i,j])
        global s *= "s[9*i+$k] = $c;\n"
    end
end
s = s*"\n\n"

k = -1
for i = 1:3
    for j = 1:3
        global k += 1
        c = sympy.ccode(H[i,j])
        global s *= "h[9*i+$k] = $c;\n"
    end
end

println(s)

gs11, gs12, gs13, gs21, gs22, gs23, gs31, gs32, gs33 = @vars gs11 gs12 gs13 gs21 gs22 gs23 gs31 gs32 gs33
gs = [
    gs11 gs12 gs13
    gs21 gs22 gs23
    gs31 gs32 gs33
]

hs11, hs12, hs13, hs21, hs22, hs23, hs31, hs32, hs33 = @vars hs11 hs12 hs13 hs21 hs22 hs23 hs31 hs32 hs33
hs = [
    hs11 hs12 hs13
    hs21 hs22 hs23
    hs31 hs32 hs33
]

a = sum(S .* gs)
gmu = diff(a, mu)|>simplify
geta = diff(a, eta)|>simplify
glambda = diff(a, lambda)|>simplify

s = ""
k = -1
for i = 1:3
    for j = 1:3
        global k += 1
        global s *= "double gs$i$j = grad_s[9*i+$k];\n"
    end
end

s *= "grad_mu_all[i] += "*sympy.ccode(gmu)*";\n"
s *= "grad_eta_all[i] += "*sympy.ccode(geta)*";\n"
s *= "grad_lambda_all[i] += "*sympy.ccode(glambda)*";\n\n"

a = sum(H .* hs)
gmu = diff(a, mu)|>simplify
geta = diff(a, eta)|>simplify
glambda = diff(a, lambda)|>simplify

k = -1
for i = 1:3
    for j = 1:3
        global k += 1
        global s *= "double hs$i$j = grad_h[9*i+$k];\n"
    end
end

s *= "grad_mu_all[i] += "*sympy.ccode(gmu)*";\n"
s *= "grad_eta_all[i] += "*sympy.ccode(geta)*";\n"
s *= "grad_lambda_all[i] += "*sympy.ccode(glambda)*";\n"

println(s)


using Einsum
using LinearAlgebra

@testset "simple" begin

    ## the very basics
    bc = rand(2,3)
    cd = rand(3,4)

    @mul A[b,d] := bc[b,c] * cd[c,d]
    @test A == bc * cd

    B = similar(A)
    @mul B[b,d] = bc[b,c] * cd[c,d] # in-place
    @test B == A

    ## and with vectors
    c = randn(3)
    @mul b[b] := bc[b,c] * c[c]
    @test b == bc * c

    @mul c[c] := b[b] * bc[b,c]
    @test c == vec(b' * bc)

    @mul z[] := b[b] * b[b]
    @test z[] == dot(b,b)

    b′ = similar(b)
    @mul b′[b] = bc[b,c] * c[c] # in-place
    @test b′ == bc * c

    c′ = similar(c)
    @mul c′[c] = b[b] * bc[b,c]

end
@testset "permutedims" begin

    bbb = rand(2,2,2);
    bbb2 = randn(2,2,2);

    @mul    A[i,j] := bbb[i,k,k2] * bbb2[j,k,k2]
    @reduce B[i,j] := sum(k,k2) bbb[i,k,k2] * bbb2[j,k,k2]
    @einsum C[i,j] := bbb[i,k,k2] * bbb2[j,k,k2]
    @test A ≈ B ≈ C

    D = similar(A);
    @mul    D[i,j] = bbb[i,k,k2] * bbb2[j,k,k2] # in-place
    @test A == D

    @mul    A[i,j] := bbb[k2,i,k] * bbb2[j,k,k2]
    @reduce B[i,j] := sum(k,k2) bbb[k2,i,k] * bbb2[j,k,k2]
    @einsum C[i,j] := bbb[k2,i,k] * bbb2[j,k,k2]
    @test A ≈ B ≈ C

    D = similar(A);
    @mul D[i,j] = bbb[k2,i,k] * bbb2[j,k,k2] # in-place
    @test A == D

    cccc = rand(3,3,3,3);
    cccc2 = randn(3,3,3,3);

    @mul    A[i,j] := cccc[x,y,j,z] * cccc2[z,i,y,x]
    @reduce B[i,j] := sum(x,y,z) cccc[x,y,j,z] * cccc2[z,i,y,x]
    @einsum C[i,j] := cccc[x,y,j,z] * cccc2[z,i,y,x]
    @test A ≈ B ≈ C

    D = similar(A);
    @mul    D[i,j] = cccc[x,y,j,z] * cccc2[z,i,y,x]  # in-place
    @test A ≈ D

    @mul    A[i,j,k,l] := cccc[i,x,y,k] * cccc2[l,j,y,x]
    @reduce B[i,j,k,l] := sum(x,y) cccc[i,x,y,k] * cccc2[l,j,y,x]
    @einsum C[i,j,k,l] := cccc[i,x,y,k] * cccc2[l,j,y,x]
    @test A ≈ B ≈ C

    D = similar(A);
    @mul    D[i,j,k,l] = cccc[i,x,y,k] * cccc2[l,j,y,x]   # in-place
    @test A ≈ D

    @mul    A[i,j,k,l,m,n] := cccc[x,k,i,m] * cccc2[l,x,j,n]
    @reduce B[i,j,k,l,m,n] := sum(x) cccc[x,k,i,m] * cccc2[l,x,j,n]
    @einsum C[i,j,k,l,m,n] := cccc[x,k,i,m] * cccc2[l,x,j,n]
    @test A ≈ B  ≈ C

    D = similar(A);
    @mul    D[i,j,k,l,m,n] = cccc[x,k,i,m] * cccc2[l,x,j,n]  # in-place
    @test A ≈ D

    ## repeat with non-alphabeitcal LHS!
    @mul    A[j,i,l,k,n,m] := cccc[x,k,i,m] * cccc2[l,x,j,n]
    @reduce B[j,i,l,k,n,m] := sum(x) cccc[x,k,i,m] * cccc2[l,x,j,n]
    @einsum C[j,i,l,k,n,m] := cccc[x,k,i,m] * cccc2[l,x,j,n]
    @test A ≈ B  ≈ C

    D = similar(A);
    @mul    D[j,i,l,k,n,m] = cccc[x,k,i,m] * cccc2[l,x,j,n]  # in-place
    @test A ≈ D

    ## Mason Potter's example:
    # W = rand(2,2,2,2); M = rand(2,2,2);
    W = rand(6,7,4,5); M = rand(7,2,3);

    @reduce N[σ, b\a, b′\a′] := sum(σ′) W[σ,σ′,b,b′] * M[σ′,a,a′]
    @mul N2[σ, b\a, b′\a′] := W[σ,σ′,b,b′] * M[σ′,a,a′]
    @test N ≈ N2

    @reduce R[σ, b,a, b′\a′] := sum(σ′) W[σ,σ′,b,b′] * M[σ′,a,a′]  lazy
    @mul R2[σ, b,a, b′\a′] := W[σ,σ′,b,b′] * M[σ′,a,a′]
    @test R ≈ R2
    # invperm((4,2,1,3)) == (3, 2, 4, 1)

    @reduce S[σ, b,a, b′,a′] := sum(σ′) W[σ,σ′,b,b′] * M[σ′,a,a′]
    @mul S2[σ, b,a, b′,a′] := W[σ,σ′,b,b′] * M[σ′,a,a′]
    @test S ≈ S2
    # invperm((4,2,1,3,5)) == (3, 2, 4, 1, 5)

    ## in-place version
    N3 = similar(N); N4 = similar(N);
    @reduce N3[σ, b\a, b′\a′] = sum(σ′) W[σ,σ′,b,b′] * M[σ′,a,a′]  lazy
    @mul N4[σ, b\a, b′\a′] = W[σ,σ′,b,b′] * M[σ′,a,a′]
    @test N ≈ N3 ≈ N4

    ## these were OK:
    M = rand(3,3,3);

    @reduce Q[σ, b\a, a′] := sum(σ′) M[σ,σ′,b] * M[σ′,a,a′];
    @mul Q2[σ, b\a, a′] := M[σ,σ′,b] * M[σ′,a,a′];
    @test Q ≈ Q2

    @reduce Q3[σ, b,a, a′] := sum(σ′) M[σ,σ′,b] * M[σ′,a,a′];
    @mul    Q4[σ, b,a, a′] := sum(σ′) M[σ,σ′,b] * M[σ′,a,a′];
    @test Q3 ≈ Q4

end
@testset "in & out shapes" begin

    bc = rand(2,3)
    cd = rand(3,4)

    @mul A[b,_,d] := bc[b,c] * cd[c,d]
    @mul A1[b,1,d] := bc[b,c] * cd[c,d]
    @test size(A) == size(A1) == (2,1,4)

    B = similar(A)
    @mul B[b,_,d] = bc[b,c] * cd[c,d] # in-place
    @test B == A

    @mul C[b] := bc[b,c] * cd[c,3]
    @test C == bc * cd[:,3]

    D = similar(C)
    @mul D[b] = bc[b,c] * cd[c,3] # in-place
    @test D == C

end
@testset "batch" begin

    bbb = rand(2,2,2)
    bbb2 = randn(2,2,2)

    @mul A[i,k,n] := bbb[i,j,n] * bbb2[j,k,n]
    @reduce B[i,k,n] := sum(j) bbb[i,j,n] * bbb2[j,k,n]
    @einsum C[i,k,n] := bbb[i,j,n] * bbb2[j,k,n]
    @test A == B == C

    D = similar(A);
    E = similar(B);
    F = similar(C);
    @mul D[i,k,n] = bbb[i,j,n] * bbb2[j,k,n] # in-place
    @reduce E[i,k,n] = sum(j) bbb[i,j,n] * bbb2[j,k,n]
    @einsum F[i,k,n] = bbb[i,j,n] * bbb2[j,k,n]
    @test A == D == E == F

    ## permute
    @mul A[n,k,i] := bbb[i,j,n] * bbb2[n,k,j]
    @reduce B[n,k,i] := sum(j) bbb[i,j,n] * bbb2[n,k,j]
    @einsum C[n,k,i] := bbb[i,j,n] * bbb2[n,k,j]
    @test A == B == C

    D = similar(A);
    E = similar(B);
    F = similar(C);
    @mul D[n,k,i] := bbb[i,j,n] * bbb2[n,k,j] # in-place
    @reduce E[n,k,i] := sum(j) bbb[i,j,n] * bbb2[n,k,j]
    @einsum F[n,k,i] := bbb[i,j,n] * bbb2[n,k,j]
    @test A == D == E == F

    ## readme & help
    X = rand(2,3,4)
    Y = rand(3,2,4)
    @mul W[β][i,j] := X[i,k,β] * Y[k,j,β]
    @test W[2] ≈ X[:,:,2] * Y[:,:,2]

    B = rand(8)
    C = rand(3,2,4)
    @mul A[_,i][j] := B[i\k] * C[j,2,k]
    rB = reshape(B,2,4)
    @einsum D[i,j] := rB[i,k] * C[j,2,k]
    @test A[1,2] ≈ D[2,:]

end
@testset "recursion" begin

    A = rand(2,3); B = rand(3,4); C = rand(4,5);

    @einsum V[i,l] := A[i,j] * B[j,k] * C[k,l]
    @reduce V2[i,l] := sum(j,k) A[i,j] * B[j,k] * C[k,l]

    @reduce W[i,l] := sum(j) A[i,j] * @mul [j,l] := B[j,k] * C[k,l]

    # @mul W2[i,l] := sum(j) A[i,j] * @mul [j,l] := B[j,k] * C[k,l]
    # maybe I decided not to allow this for now.

    @test V ≈ V2 ≈ W #≈ W2

end


using Flux
println("Integro Differential equation tests")
using DiffEqFlux
println("Starting Soon!")
using ModelingToolkit
using DiffEqBase
using Test, NeuralPDE
println("Starting Soon!")
using GalacticOptim
using Optim
using Quadrature,Cubature, Cuba
using QuasiMonteCarlo
using SciMLBase
import ModelingToolkit: Interval, infimum, supremum
using DomainSets

using Random
Random.seed!(100)

cb = function (p,l)
    println("Current loss is: $l")
    return false
end

#Integration Tests
println("Integral Tests")
@parameters t
@variables i(..)
Di = Differential(t)
Ii = Integral(t in DomainSets.ClosedInterval(0, t))
eq = Di(i(t)) + 2*i(t) + 5*Ii(i(t)) ~ 1
bcs = [i(0.) ~ 0.0]
domains = [t ∈ Interval(0.0,2.0)]
chain = Chain(Dense(1,15,Flux.σ),Dense(15,1))
initθ = Float64.(DiffEqFlux.initial_params(chain))
strategy_ = NeuralPDE.GridTraining(0.1)
discretization = NeuralPDE.PhysicsInformedNN(chain,
                                             strategy_;
                                             init_params = nothing,
                                             phi = nothing,
                                             derivative = nothing,
                                             )
@named pde_system = PDESystem(eq,bcs,domains,[t],[i(t)])
sym_prob = NeuralPDE.symbolic_discretize(pde_system, discretization)
prob = NeuralPDE.discretize(pde_system,discretization)
res = GalacticOptim.solve(prob, BFGS(); cb = cb, maxiters=100)
ts = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1]
phi = discretization.phi

analytic_sol_func(t) = 1/2*(exp(-t))*(sin(2*t))
u_real  = [analytic_sol_func(t) for t in ts]
u_predict  = [first(phi([t],res.minimizer)) for t in ts]
@test Flux.mse(u_real, u_predict) < 0.001
# plot(ts,u_real)
# plot!(ts,u_predict)

## Simple Integral Test
println("Simple Integral Test")

@parameters x
@variables u(..)
Ix = Integral(x in DomainSets.ClosedInterval(0, x))
# eq = Ix(u(x)) ~ (x^3)/3
eq = Ix(u(x)*cos(x))~ (x^3)/3

bcs = [u(0.) ~ 0.0]
domains = [x ∈ Interval(0.0,1.00)]
# chain = Chain(Dense(1,15,Flux.σ),Dense(15,1))
chain = FastChain(FastDense(1,15,Flux.σ),FastDense(15,1))
initθ = Float64.(DiffEqFlux.initial_params(chain))
strategy_ = NeuralPDE.GridTraining(0.1)
discretization = NeuralPDE.PhysicsInformedNN(chain,
                                             strategy_;
                                             init_params = initθ,
                                             phi = nothing,
                                             derivative = nothing,
                                             )
@named pde_system = PDESystem(eq,bcs,domains,[x],[u(x)])
prob = NeuralPDE.discretize(pde_system,discretization)
res = GalacticOptim.solve(prob, BFGS(); cb = cb, maxiters=200)
xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1]
phi = discretization.phi
u_predict  = [first(phi([x],res.minimizer)) for x in xs]
u_real  = [x^2/cos(x) for x in xs]
@test Flux.mse(u_real, u_predict) < 0.001

# plot(xs,u_real)
# plot!(xs,u_predict)

#simple multidimensitonal integral test
println("simple multidimensitonal integral test")

@parameters x,y
@variables u(..)
Dx = Differential(x)
Dy = Differential(y)
Ix = Integral((x,y) in DomainSets.UnitSquare())
eq = Ix(u(x,y)) ~ 1/3
bcs = [u(0., 0.) ~ 1, Dx(u(x,y)) ~ -2*x , Dy(u(x ,y)) ~ -2*y ]
domains = [x ∈ Interval(0.0,1.00), y ∈ Interval(0.0,1.00)]
chain = Chain(Dense(2,15,Flux.σ),Dense(15,1))
initθ = Float64.(DiffEqFlux.initial_params(chain))
strategy_ = NeuralPDE.GridTraining(0.1)
discretization = NeuralPDE.PhysicsInformedNN(chain,
                                             strategy_;
                                             init_params = nothing,
                                             phi = nothing,
                                             derivative = nothing,
                                             )
@named pde_system = PDESystem(eq,bcs,domains,[x,y],[u(x, y)])
prob = NeuralPDE.discretize(pde_system,discretization)
res = GalacticOptim.solve(prob, BFGS(); cb = cb, maxiters=100)
xs = 0.00:0.01:1.00
ys = 0.00:0.01:1.00
phi = discretization.phi

u_real = collect(1 - x^2 - y^2 for y in ys, x in xs);
u_predict = collect(Array(phi([x,y], res.minimizer))[1] for y in ys, x in xs);
@test Flux.mse(u_real, u_predict) < 0.001

# error_ = u_predict .- u_real
# p1 = plot(xs,ys,u_real,linetype=:contourf,label = "analytic")
# p2 = plot(xs,ys,u_predict,linetype=:contourf,label = "predict")
# p3 = plot(xs,ys,error_,linetype=:contourf,label = "error")
# plot(p1,p2,p3)

@parameters x,y
@variables u(..)
Dx = Differential(x)
Dy = Differential(y)
Ix = Integral((x,y) in DomainSets.ProductDomain(UnitInterval(),ClosedInterval(0 ,x)))
eq = Ix(u(x,y)) ~ 5/12
bcs = [u(0., 0.) ~ 0, Dy(u(x,y)) ~ 2*y , u(x, 0) ~ x ]
domains = [x ∈ Interval(0.0,1.00), y ∈ Interval(0.0,1.00)]
chain = Chain(Dense(2,15,Flux.σ),Dense(15,1))
initθ = Float64.(DiffEqFlux.initial_params(chain))
strategy_ = NeuralPDE.GridTraining(0.1)
discretization = NeuralPDE.PhysicsInformedNN(chain,
                                             strategy_;
                                             init_params = nothing,
                                             phi = nothing,
                                             derivative = nothing,
                                             )
@named pde_system = PDESystem(eq,bcs,domains,[x,y],[u(x, y)])
prob = NeuralPDE.discretize(pde_system,discretization)
res = GalacticOptim.solve(prob, BFGS(); cb = cb, maxiters=100)
xs = 0.00:0.01:1.00
ys = 0.00:0.01:1.00
phi = discretization.phi

u_real = collect( x + y^2 for y in ys, x in xs);
u_predict = collect(Array(phi([x,y], res.minimizer))[1] for y in ys, x in xs);
@test Flux.mse(u_real, u_predict) < 0.01

# error_ = u_predict .- u_real
# p1 = plot(xs,ys,u_real,linetype=:contourf,label = "analytic")
# p2 = plot(xs,ys,u_predict,linetype=:contourf,label = "predict")
# p3 = plot(xs,ys,error_,linetype=:contourf,label = "error")
# plot(p1,p2,p3)

## Two variables Integral Test
println("Two variables Integral Test")

@parameters x
@variables u(..) w(..)
Dx = Differential(x)
Ix = Integral(x in DomainSets.ClosedInterval(1, x))
# Iinf = Integral(x in DomainSets.ClosedInterval(1, Inf))

eqs =  [Ix(u(x)*w(x)) ~ log(abs(x)),
        Dx(w(x)) ~ -2/(x^3),
        u(x) ~ x ]

bcs = [u(1.) ~ 1.0, w(1.) ~ 1.0]
domains = [x ∈ Interval(1.0,2.0)]

chains = [FastChain(FastDense(1,15,Flux.σ),FastDense(15,1)) for _ in 1:2]
initθ = map(chain -> Float64.(DiffEqFlux.initial_params(chain)),chains)
strategy_ = NeuralPDE.GridTraining(0.1)
discretization = NeuralPDE.PhysicsInformedNN(chains,
                                             strategy_;
                                             init_params = initθ
                                             )
@named pde_system = PDESystem(eqs,bcs,domains,[x],[u(x), w(x)])
prob = NeuralPDE.discretize(pde_system,discretization)
res = GalacticOptim.solve(prob, BFGS(); cb = cb, maxiters=200)
xs = [infimum(d.domain):0.01:supremum(d.domain) for d in domains][1]
phi = discretization.phi
initθ = discretization.init_params
acum =  [0;accumulate(+, length.(initθ))]
sep = [acum[i]+1 : acum[i+1] for i in 1:length(acum)-1]
minimizers = [res.minimizer[s] for s in sep]

u_predict  = [(phi[1]([x],minimizers[1]))[1] for x in xs]
w_predict  = [(phi[2]([x],minimizers[2]))[1] for x in xs]
u_real  = [x for x in xs]
w_real  = [1/x^2 for x in xs]
@test Flux.mse(u_real, u_predict) < 0.001
@test Flux.mse(w_real, w_predict) < 0.001

# plot(xs,u_real)
# plot!(xs,u_predict)
# plot(xs,w_real)
# plot!(xs,w_predict)


"""
  Returns the a_j,k and c_k,p Coefficients considering the approximation of the
  exponential term exp(-t_j * z_p) = sum( k : a_j,k * c_k,p )
"""
function get_AC_coefficients(K::Int64
                            , times::Vector
                            , z_p::Vector
                            , alpha::Float64
                            , m::Float64
                            ; method = "nfft")

  if method == "leastsquare"
    A,C = get_AC_oefficients_lsqr(K,times,z_p)
  elseif method == "nfft"
    A,C = get_AC_coefficients_nfft(K,times,z_p,alpha,m)
  elseif method == "hist"
    A,C = get_AC_coefficients_hist_lsqr(K,times,z_p)
  else
    A = zeros(ComplexF64,length(times),K)
    C = zeros(ComplexF64,K,length(z_p))
    error("Not implemented yet")
  end
  return A,C
end

function _precompute_coffs_nfft(K::Int64
                                , times::Vector     # readout times
                                , z_p::Vector       # Correction map
                                , alpha::Float64    # oversampling factor of NFFT
                                , m::Float64        # kernel size of NFFT
                                ; adaptK::Bool=true
                                , centering::Bool=true)

  # Centering time and correctionterm to lower computational effort
  t_hat = (times[1] + times[end])/2
  z_hat = minimum(real(z_p)) + maximum(real(z_p)) + 1im*(minimum(imag(z_p)) + maximum(imag(z_p)))
  z_hat *= 0.5
  if !centering
    t_hat=0.0
    z_hat=0.0
  end
  # Centering readouttimes and correctionmap
  times = times .- t_hat
  z_p = z_p .- z_hat

  # Calculating timescaling factor so it satisfies: t_j / T in [0.5,0.5)
  T = maximum(abs.(times)) / (0.5-1e-7)

  # determining N and K,
  # K can be choosen freely with tradeoff comp. complexity vs precision
  N = Int64(4*ceil( maximum(abs.(z_p)) * maximum(abs.(times)) / (2*pi) ))
  if adaptK
    K = ceil(Int64, alpha*N + 2*m)
    # println("Calculated K for NFFT: ",K)
  end

  win, win_hat = NFFT.getWindow(:kaiser_bessel)

  return win, win_hat, K, N, T, times, z_p
end

function get_A_coefficients_nfft(K::Int64
                                , times::Vector     # readout times
                                , z_p::Vector       # Correction map
                                , alpha::Float64    # oversampling factor(NFFT)
                                , m::Float64        # kernel size(NFFT)
                                ; adaptK::Bool=true
                                , centering::Bool=true)

  win, win_hat, K, N, T, times, z_p = _precompute_coffs_nfft(K, times, z_p, alpha, m; adaptK=adaptK,centering=centering)

  A = zeros(ComplexF64, length(times), K)

  for kappa=1:K
    for j=1:length(times)
        argument = times[j]/T - (kappa - K/2)/(K - 2*m)
        A[j,kappa] = win(argument,N*alpha,m,alpha)
    end
  end
  return A
end

function get_C_coefficients_nfft(K::Int64
                                , times::Vector     # readout times
                                , z_p::Vector       # Correction map
                                , alpha::Float64    # oversampling factor(NFFT)
                                , m::Float64        # kernel size(NFFT)
                                ; adaptK::Bool=true
                                , centering::Bool=true)

  win, win_hat, K, N, T, times, z_p = _precompute_coffs_nfft(K, times, z_p, alpha, m; adaptK=adaptK,centering=centering)

  C = zeros(ComplexF64, K, length(z_p))

  for p=1:length(z_p)
    for kappa=1:K
        fourierarg = 1im * z_p[p]*T/(2pi)
        exparg = -z_p[p]*T * ( (kappa - K/2)/(K-2*m) )
        C[kappa,p] = (1/ ( (K-2*m)* 1/(N*alpha)*win_hat(fourierarg,N*alpha,m,alpha))) * exp(exparg)
    end
  end
  return C
end

"""
  Returns the a_j,k and c_k,p Coefficients considering the approximation of the
  exponential term exp(-t_j * z_p) = sum( k : a_j,k * c_k,p )

  Using the nfft method
"""
function get_AC_coefficients_nfft(K::Int64
                                , times::Vector     # readout times
                                , z_p::Vector       # Correction map
                                , alpha::Float64    # oversampling factor(NFFT)
                                , m::Float64        # kernel size(NFFT)
                                ; adaptK::Bool=true
                                , centering::Bool=true)

  A = get_A_coefficients_nfft(K, times, z_p, alpha, m; adaptK=adaptK, centering=centering)
  C = get_C_coefficients_nfft(K, times, z_p, alpha, m; adaptK=adaptK, centering=centering)
  return A, C
end

"""
  Returns the a_j,k and c_k,p Coefficients considering the approximation of the
  exponential term exp(-t_j * z_p) = sum( k : a_j,k * c_k,p )

  Using the least squares method
"""
function get_AC_coefficients_lsqr(K::Int64, times::Vector, z_p::Vector)
  A = zeros(ComplexF64,length(times),K)
  C = zeros(ComplexF64,K,length(z_p))
  G = zeros(ComplexF64,length(z_p),K)

  progr = Progress(length(times), 1, "Using leastsquares to get Coefficients...")

  t_hat = zeros(K)
  t_hat = collect(range(times[1], stop=times[length(times)], length=K))

  z_p = vec(z_p)
  for kappa=1:K
    for p=1:length(z_p)
      G[p,kappa] = exp(-t_hat[kappa]*z_p[p])
      C[kappa,p] = exp(-t_hat[kappa]*z_p[p])
    end
  end

  systemMat = G'*G
  err = zeros(length(times))
  for j=1:length(times)
    b = [exp(-times[j]*z_p[p]) for p=1:length(z_p) ]
    A[j,:] = systemMat \ (G'*b)
    next!(progr)
  end
  return A,C
end

"""
  Returns the a_j,k Coefficients considering the approximation of the
  exponential term exp(-t_j * z_p) = sum( k : a_j,k * c_k,p )

  Using the least squares method with histogramms
"""
function get_A_coefficients_hist_lsqr(K::Int64, times::Vector, z_p::Vector)

  numBins = 10*K
  if real.(z_p)== zeros(length(z_p))
    fmax_im = maximum(imag.(z_p))
    fmin_im = minimum(imag.(z_p))
    δf = (fmax_im-fmin_im)/(numBins-1)
    edges_im = [fmin_im-δf/2 + i*δf for i = 0:numBins]
    h = fit(Histogram, imag.(z_p),edges_im,closed=:left)
    z_center = 0.5*1im*vec( [edges_im[i]+edges_im[i+1] for i=1:numBins] )
    z_weights = h.weights
  elseif imag.(z_p)== zeros(length(z_p))
    fmax_re = maximum(real.(z_p))
    fmin_re = minimum(real.(z_p))
    δf = (fmax_re-fmin_re)/(numBins-1)
    edges_re = collect(fmin_re-δf/2.:δf:fmax_re+δf/2)
    h = fit(Histogram, real.(z_p),edges_re,closed=:left)
    z_center = 0.5*vec( [edges_re[i]+edges_re[i+1] for i=1:numBins] )
    z_weights = h.weights
  else
    nbins = ceil(Int, sqrt(numBins))
    fmax_re = maximum(real.(z_p))
    fmin_re = minimum(real.(z_p))
    δf_re = (fmax_re-fmin_re)/(nbins-1)
    edges_re = collect(fmin_re-δf/2.:δf:fmax_re+δf/2)
    fmax_im = maximum(imag.(z_p))
    fmin_im = minimum(imag.(z_p))
    δf = (fmax_im-fmin_im)/(nbins-1)
    edges_im = collect(fmin_im-δf/2.:δf:fmax_im+δf/2)
    h = fit(Histogram, (real.(z_p),imag.(z_p)),(edges_re,edges_im),closed=:left)
    z_center = 0.5*vec( [edges_re[i]+edges_re[i+1]+1im*(edges_im[j]+edges_im[j+1]) for i=1:numBins[1], j=1:numBins[2] ] )
    z_weights = vec(h.weights)
    numBins = nbins^2
  end

  A = zeros(ComplexF64,length(times),K)
  G = zeros(ComplexF64,numBins,K)

  # progr = Progress(length(times), 1, "Using leastsquares to get Hist-coefficients...")

  t_hat = zeros(K)
  # t_hat = collect(linspace(times[1],times[length(times)],K))
  t_hat = collect(range(times[1], stop=times[end], length=K))

  z_p = vec(z_p)
  for p=1:numBins
    for kappa=1:K
      G[p,kappa] = exp(-t_hat[kappa]*z_center[p])
    end
  end

  # systemMat = G'*diagm(z_weights)*G
  systemMat = G'*Diagonal(z_weights)*G
  err = zeros(length(times))
  @showprogress 1 "Using leastsquares to get Hist-coefficients..." for j=1:length(times) #length(times)
    b = [z_weights[p]*exp(-times[j]*z_center[p]) for p=1:numBins ]
    A[j,:] = systemMat \ (G'*b)
    # next!(progr)
  end
  return A
end

"""
  Returns the c_k,p Coefficients considering the approximation of the
  exponential term exp(-t_j * z_p) = sum( k : a_j,k * c_k,p )

  Using the least squares method with histogramms
"""
function get_C_coefficients_hist_lsqr(K::Int64, times::Vector, z_p::Vector)

  C = zeros(ComplexF64,K,length(z_p))
  t_hat = collect(range(times[1],stop=times[end],length=K))
  # z_p = vec(z_p)
  for p=1:length(z_p)
    for kappa=1:K
      C[kappa,p] = exp(-t_hat[kappa]*z_p[p])
    end
  end

  return C
end

function get_A_coefficients_one_term(K::Int64, times::Vector, z_p::Vector)
  A = zeros(ComplexF64,length(times),K)
  t_hat = collect(range(times[1],stop=times[end],length=K))
  for p=1:length(times)
    diff, idx = findmin(abs.(t_hat .- times[p]))
    A[p,idx]=1.0
  end
  return A
end

"""
  Returns the a_j,k Coefficients considering the approximation of the
  exponential term exp(-t_j * z_p) = sum( k : a_j,k * c_k,p )

  Using the least squares method with histogramms
"""
function get_AC_coefficients_two_terms(K::Int64, times::Vector, z_p::Vector)

  # build histogram
  numBins = 10*K
  if real.(z_p)== zeros(length(z_p))
    fmax_im = maximum(imag.(z_p))
    fmin_im = minimum(imag.(z_p))
    δf = (fmax_im-fmin_im)/(numBins-1)
    edges_im = [fmin_im-δf/2 + i*δf for i = 0:numBins]
    h = fit(Histogram, imag.(z_p),edges_im,closed=:left)
    z_center = 0.5*1im*vec( [edges_im[i]+edges_im[i+1] for i=1:numBins] )
    z_weights = h.weights
  elseif imag.(z_p)== zeros(length(z_p))
    fmax_re = maximum(real.(z_p))
    fmin_re = minimum(real.(z_p))
    δf = (fmax_re-fmin_re)/(numBins-1)
    edges_re = collect(fmin_re-δf/2.:δf:fmax_re+δf/2)
    h = fit(Histogram, real.(z_p),edges_re,closed=:left)
    z_center = 0.5*vec( [edges_re[i]+edges_re[i+1] for i=1:numBins] )
    z_weights = h.weights
  else
    nbins = ceil(Int, sqrt(numBins))
    fmax_re = maximum(real.(z_p))
    fmin_re = minimum(real.(z_p))
    δf_re = (fmax_re-fmin_re)/(nbins-1)
    edges_re = collect(fmin_re-δf/2.:δf:fmax_re+δf/2)
    fmax_im = maximum(imag.(z_p))
    fmin_im = minimum(imag.(z_p))
    δf = (fmax_im-fmin_im)/(nbins-1)
    edges_im = collect(fmin_im-δf/2.:δf:fmax_im+δf/2)
    h = fit(Histogram, (real.(z_p),imag.(z_p)),(edges_re,edges_im),closed=:left)
    z_center = 0.5*vec( [edges_re[i]+edges_re[i+1]+1im*(edges_im[j]+edges_im[j+1]) for i=1:numBins[1], j=1:numBins[2] ] )
    z_weights = vec(h.weights)
    numBins = nbins^2
  end

  A = zeros(ComplexF64,length(times),K)
  G = zeros(ComplexF64,numBins,K)

  progr = Progress(length(times), 1, "Using leastsquares to get Hist-coefficients...")

  # t_hat = zeros(K)
  t_hat = collect(range(times[1],stop=times[end],length=K))

  z_p = vec(z_p)
  for p=1:numBins
    for kappa=1:K
      G[p,kappa] = exp(-t_hat[kappa]*z_center[p])
    end
  end

  for j=1:length(times) #length(times)
    # find closest translate with smaller sampling time
    diff, idx = findmin(abs.(t_hat .- times[j]))
    if (t_hat[idx] >= times[j]) && idx > 1
      idx -= 1
    end

    systemMat = G[:,idx:idx+1]'*diagm(0 => z_weights)*G[:,idx:idx+1]
    b = [z_weights[p]*exp(-times[j]*z_center[p]) for p=1:numBins ]
    A[j,idx:idx+1] = systemMat \ (G[:,idx:idx+1]'*b)
    next!(progr)
  end
  return A
end


using Reexport

@reexport module UncertainStatistics
    using StatsBase
    using Distributions

    # Uncertain values and datasets types
    using ..UncertainValues
    using ..UncertainDatasets

    # Sampling constraints and resampling functions
    using ..Resampling

    import ..UVAL_COLLECTION_TYPES

    # Definitions of statistical methods for UncertainDatasets
    include("statsbase/uncertain_values/core_stats_point_estimates.jl")
    include("statsbase/uncertain_values/core_stats_pair_estimates.jl")
	include("statsbase/uncertain_datasets/core_stats_uncertaindatasets.jl")
    include("statsbase/uncertain_datasets/core_stats_uncertaindatasets_single_estimates.jl")
    include("statsbase/uncertain_datasets/core_stats_uncertaindatasets_pairwise_estimates.jl")

    include("hypothesis_tests/mann_whitney.jl")
    include("hypothesis_tests/t_tests.jl")
    include("hypothesis_tests/anderson_darling.jl")
    include("hypothesis_tests/kolmogorov_smirnov.jl")
    include("hypothesis_tests/jarque_bera.jl")

    include("hypothesis_tests/timeseries_tests.jl")

    #include("StatsBase_meanfunctions.jl")
    #include("StatsBase_scalarstatistics_moments.jl")
    #include("StatsBase_variation.jl")
    #include("StatsBase_zscores.jl")
    #include("StatsBase_entropies.jl")
    #include("StatsBase_quantiles.jl")
end

"""
    UncertainStatistics

A module defining function to compute various statistics for uncertain values
and uncertain datasets.
"""


using DataFrames
using CSV
using MLBase
using StatsModels
using StatsBase
using GLM

data = DataFrame(CSV.File("tacke40c.csv"))

fm = @formula(y ~ x)

k = 5

averageAbsMeanErrorTest = 0.0
subData = collect(Kfold(length(data.x), k))

for i in 1:k
    dataTrain = data[subData[i],:]
    dataTest = data[setdiff(1:end, subData[i]),:]

    linearRegressor = lm(fm, dataTest)
    dataPredictTest = predict(linearRegressor, dataTest)

    errorTest = dataTest.y - dataPredictTest
    absMeanErrorTest = mean(abs.(errorTest))
    println("Prosecna greska [$i]: $absMeanErrorTest")
    global averageAbsMeanErrorTest += absMeanErrorTest
end

averageAbsMeanErrorTest /= k

println("Prosecna absolutna greska je: $averageAbsMeanErrorTest")


##  Wavelet cross spectrum in julia
##  Author Bernd Blasius (bernd.blasius@gmail.com)

using DSP, Statistics

include("wavelet.jl")
include("tools.jl")

# calculate the wavelet cross spectrum (amplitude and phase) and
# wavelet coherence of two time series
function wcs(y1,y2,dt;
             pad=1,        # pad the time series with zeros
             nvoices = 32, # total number of voices
             noctave = 4,  # total number of octaves
             s0 = 2*dt,    # starting scale
             mother = "MORLET",
             param = 6)
    # x1, x2 : the two vectors of the time series
    # dt: sampling time interval
    # this function assumes that both time series cover the same time range
    # and are sampled equally
    # no checks are made to ensure this!!

    dj = 1 ./ nvoices
    j1 = noctave * nvoices

    # do the wavelet transform
    # period and scale should be returned the same
    wave1, period, scale, coi = wavelet(y1,dt,pad,dj,s0,j1,mother,param)
    wave2, period, scale, coi = wavelet(y2,dt,pad,dj,s0,j1,mother,param)

    power1 = abs2.(wave1)   # wavelet power spetrum
    power2 = abs2.(wave2)   # wavelet power spetrum

    # calculate wavelet cross spectrum (CWS)
    wcs = wave1 .* conj.(wave2)

    # smoothing
    spower1 = smooth(power1,dt,nvoices,period)
    spower2 = smooth(power2,dt,nvoices,period)
    swcs    = smooth(wcs,   dt,nvoices,period)

    amplitude = abs.(swcs)   # amplitude of wavelet cross spectrum
    phase     = angle.(swcs) # phase of wavelet cross spectrum

    # wavelet coherence
    wco = amplitude ./ (sqrt.(spower1) .* sqrt.(spower2))

    return spower1, spower2, amplitude, phase, wco, period, scale, coi
end


using MittagLefflerFunctions
using PyPlot

f(x) = exp(x)

nmax = 10
M = nmax + 2
a = chebyshev_coefs(Float64, f, nmax, M)

x = range(-1, 1, length=201)
y = Float64[ chebyshev_sum(a, xval) for xval in x ]

figure(1)
plot(x, y-f.(x))
grid(true)


type Knn
  k::Int
  x
  y
  means_train
  std_train
  norm_fx
  Knn() = new(3)
end

function train!{U<:Real,T}(model::Knn, x::Matrix{U}, y::Vector{T})
  model.x = x
  model.y = y
  nRows, nCols = size(x)
  model.means_train = Vector{Float64}(nCols)
  model.std_train = Vector{Float64}(nCols)
  for j in 1:nCols
      model.means_train[j] = mean(x[:,j])
      model.std_train[j] = std(x[:,j])
  end
end

function predict{U<:Real}(model::Knn,x::Matrix{U})
    nRows, nCols = size(x)
    x_norm = deepcopy(x)
    for j in 1:nCols
      for i in  1:nRows
         x_norm[i,j] = (x[i,j] - model.means_train[j])/(model.std_train[j])
      end
    end
    if (eltype(model.y)<:Real)
       #Regression
        resp = Vector{Float64}(nRows)
        for i in 1:size(x_norm,1)
          resp[i]= assign_value(x_norm, model.y, model.k, x_norm[i,:])
        end
        return resp
    else
      error("Classification is not implemented yet")
    end
end

function euclidean_distance(a, b)
 distance = 0.0
 for index in 1:size(a, 1)
  distance += (a[index]-b[index]) * (a[index]-b[index])
 end
 return distance
end

function get_k_nearest_neighbors(xTrain, rowTestItem, k)
   nRows, nCols = size(xTrain)
   distances = Array(Float32,nRows)
    for index in 1:nRows
     distances[index] = euclidean_distance(xTrain[index,:], rowTestItem)
    end
   sortedNeighbors = sortperm(distances)
   kNearestNeighbors = sortedNeighbors[1:k]
   return kNearestNeighbors
end

# assign_value(model.x_norm, model.y, model.k, x_norm[i,:])
function assign_value(xTrain, yTrain, k, imageI)
 kNearestNeighbors = get_k_nearest_neighbors(xTrain, imageI, k)

 soma::Float64 = 0
 highestCount = 0
 mostPopularLabel = 0
 for n in kNearestNeighbors
  soma = soma + yTrain[n]
 end
 return soma / k
end

function assign_label(xTrain, yTrain, k, imageI)
 kNearestNeighbors = get_k_nearest_neighbors(xTrain, imageI, k)
 counts = Dict{Int, Int}()
 highestCount = 0
 mostPopularLabel = 0
 for n in kNearestNeighbors
  labelOfN = yTrain[n]
  if !haskey(counts, labelOfN)
   counts[labelOfN] = 0
  end
  counts[labelOfN] += 1 #add one to the count
  if counts[labelOfN] > highestCount
   highestCount = counts[labelOfN]
   mostPopularLabel = labelOfN
  end
 end
 return mostPopularLabel
end


# Copyright (c) 2021 Idiap Research Institute, http://www.idiap.ch/
#  Niccolò Antonello <nantonel@idiap.ch>

for T in [Float64, Float32]

  #################
  # ProbabilityWeight semiring
  #################
  a = ProbabilityWeight(rand(T))
  b = ProbabilityWeight(rand(T))
  c = ProbabilityWeight(rand(T))
  d = ProbabilityWeight{T}(a) # check ProbabilityWeight(x::ProbabilityWeight)
  e = ProbabilityWeight{T}( ProbabilityWeight{T == Float64 ? Float32 : Float64}(get(a)) )
  f = ProbabilityWeight{T}( rand(T == Float64 ? Float32 : Float64) )
  @test typeof(e) == ProbabilityWeight{T}
  @test typeof(f) == ProbabilityWeight{T}

  @test (a*b).x == a.x*b.x
  @test (a + b).x == a.x + b.x
  @test (a/b).x == a.x/b.x
  @test reverse(a) == a
  @test FiniteStateTransducers.reverseback(a) == a
  o = one(ProbabilityWeight{T})
  @test o.x  == one(T)
  z = zero(ProbabilityWeight{T})
  @test z.x == zero(T)

  # + commutative monoid
  @test (a+b)+c ≈ a+(b+c)
  @test a+z ≈ z+a ≈ a
  @test a+b ≈ b+a
  # * monoid
  @test (a*b)*c ≈ a*(b*c)
  @test a*o == o*a == a # identity
  @test a*z == z*a == z # annihilator
  # distribution over addition
  @test a*(b+c) ≈ (a*b)+(a*c)
  @test (a+b)*c ≈ (a*c)+(b*c)
  # division
  c = a*b
  @test a*(c/a) ≈ c
  @test (c/b)*b ≈ c

  # properties
  W = ProbabilityWeight{T}
  @test FiniteStateTransducers.isidempotent(W) == false
  @test FiniteStateTransducers.iscommulative(W)
  @test FiniteStateTransducers.isleft(W)
  @test FiniteStateTransducers.isright(W)
  @test FiniteStateTransducers.isweaklydivisible(W)
  @test FiniteStateTransducers.iscomplete(W)
  @test FiniteStateTransducers.ispath(W) == false
  @test FiniteStateTransducers.reversetype(W) == W

  a = parse(W,"1.0")
  @test a.x == T(1.0)

  #################
  # LogWeight semiring
  #################
  a = LogWeight(1)
  b = LogWeight(2)
  @test (a + b).x ≈ log(exp(a.x)+exp(b.x))

  a = LogWeight(2)
  b = LogWeight(1)
  @test (a + b).x ≈ log(exp(a.x)+exp(b.x))

  a = LogWeight(rand(T))
  b = LogWeight(rand(T))
  c = LogWeight(rand(T))
  d = LogWeight{T}(a) # check ProbabilityWeight(x::ProbabilityWeight)
  e = LogWeight{T}( LogWeight{T == Float64 ? Float32 : Float64}(get(a)) )
  @test typeof(e) == LogWeight{T}

  @test (a*b).x ≈ a.x + b.x
  @test (a + b).x ≈ log(exp(a.x)+exp(b.x))
  @test (a / b).x ≈ log(exp(a.x)/exp(b.x))
  @test reverse(a) == a
  @test FiniteStateTransducers.reverseback(a) == a
  o = one(LogWeight{T})
  @test o.x  == 0
  z = zero(LogWeight{T})
  @test z.x == T(log(0))

  # + commutative monoid
  @test (a+b)+c ≈ a+(b+c)
  @test a+z ≈ z+a ≈ a
  @test a+b ≈ b+a
  # * monoid
  @test (a*b)*c ≈ a*(b*c)
  @test a*o == o*a == a # identity
  @test a*z == z*a == z # annihilator
  # distribution over addition
  @test a*(b+c) ≈ (a*b)+(a*c)
  @test (a+b)*c ≈ (a*c)+(b*c)
  # division
  c = a*b
  @test a*(c/a) ≈ c
  @test (c/b)*b ≈ c

  # properties
  W = LogWeight{T}
  @test FiniteStateTransducers.isidempotent(W) == false
  @test FiniteStateTransducers.iscommulative(W)
  @test FiniteStateTransducers.isleft(W)
  @test FiniteStateTransducers.isright(W)
  @test FiniteStateTransducers.isweaklydivisible(W)
  @test FiniteStateTransducers.iscomplete(W)
  @test FiniteStateTransducers.ispath(W) == false
  @test FiniteStateTransducers.reversetype(W) == W

  # parse
  a = parse(W,"1.0")
  @test a.x == T(1.0)

  #################
  # NLogWeight semiring
  #################
  a = NLogWeight(1)
  b = NLogWeight(2)
  @test (a + b).x ≈ -log(exp(-a.x)+exp(-b.x))

  a = NLogWeight(2)
  b = NLogWeight(1)
  @test (a + b).x ≈ -log(exp(-a.x)+exp(-b.x))

  a = NLogWeight(rand(T))
  b = NLogWeight(rand(T))
  c = NLogWeight(rand(T))
  d = NLogWeight{T}(a) # check ProbabilityWeight(x::ProbabilityWeight)
  e = NLogWeight{T}( NLogWeight{T == Float64 ? Float32 : Float64}(get(a)) )
  @test typeof(e) == NLogWeight{T}

  @test (a*b).x ≈ a.x + b.x
  @test (a + b).x ≈ -log(exp(-a.x)+exp(-b.x))
  @test (a / b).x ≈ log(exp(a.x)/exp(b.x))
  @test reverse(a) == a
  @test FiniteStateTransducers.reverseback(a) == a
  o = one(NLogWeight{T})
  @test o.x  == 0
  z = zero(NLogWeight{T})
  @test z.x == T(-log(0))

  # + commutative monoid
  @test (a+b)+c ≈ a+(b+c)
  @test a+z ≈ z+a ≈ a
  @test a+b ≈ b+a
  # * monoid
  @test (a*b)*c ≈ a*(b*c)
  @test a*o == o*a == a # identity
  @test a*z == z*a == z # annihilator
  # distribution over addition
  @test a*(b+c) ≈ (a*b)+(a*c)
  @test (a+b)*c ≈ (a*c)+(b*c)
  # division
  c = a*b
  @test a*(c/a) ≈ c
  @test (c/b)*b ≈ c

  # properties
  W = NLogWeight{T}
  @test FiniteStateTransducers.isidempotent(W) == false
  @test FiniteStateTransducers.iscommulative(W)
  @test FiniteStateTransducers.isleft(W)
  @test FiniteStateTransducers.isright(W)
  @test FiniteStateTransducers.isweaklydivisible(W)
  @test FiniteStateTransducers.iscomplete(W)
  @test FiniteStateTransducers.ispath(W) == false
  @test FiniteStateTransducers.reversetype(W) == W

  # parse
  a = parse(W,"1.0")
  @test a.x == T(1.0)

  #################
  # TropicalWeight semiring
  #################
  a = TropicalWeight(rand(T))
  b = TropicalWeight(rand(T))
  c = TropicalWeight(rand(T))
  d = TropicalWeight{T}(a) # check ProbabilityWeight(x::ProbabilityWeight)
  e = TropicalWeight{T}( TropicalWeight{T == Float64 ? Float32 : Float64}(get(a)) )
  @test typeof(e) == TropicalWeight{T}

  @test (a*b).x == a.x + b.x
  @test (a + b).x == min(a.x,b.x)
  @test (a / b).x ≈ log(exp(a.x)/exp(b.x))
  o = one(TropicalWeight{T})
  @test o.x  == 0
  z = zero(TropicalWeight{T})
  @test z.x == T(Inf)

  # + commutative monoid
  @test (a+b)+c ≈ a+(b+c)
  @test a+z ≈ z+a ≈ a
  @test a+b ≈ b+a
  # * monoid
  @test (a*b)*c ≈ a*(b*c)
  @test a*o == o*a == a # identity
  @test a*z == z*a == z # annihilator
  # distribution over addition
  @test a*(b+c) ≈ (a*b)+(a*c)
  @test (a+b)*c ≈ (a*c)+(b*c)
  # division
  c = a*b
  @test a*(c/a) ≈ c
  @test (c/b)*b ≈ c

  W = TropicalWeight{T}
  @test FiniteStateTransducers.isidempotent(W)
  @test FiniteStateTransducers.iscommulative(W)
  @test FiniteStateTransducers.isleft(W)
  @test FiniteStateTransducers.isright(W)
  @test FiniteStateTransducers.isweaklydivisible(W)
  @test FiniteStateTransducers.iscomplete(W)
  @test FiniteStateTransducers.ispath(W)
  @test FiniteStateTransducers.reversetype(W) == W

  # parse
  a = parse(W,"1.0")
  @test a.x == T(1.0)
end

a = BoolWeight(rand([true;false]))
b = BoolWeight(rand([true;false]))
c = BoolWeight(rand([true;false]))
d = BoolWeight(a) # check ProbabilityWeight(x::ProbabilityWeight)

@test (a*b).x == ( a.x && b.x)
@test (a + b).x == ( a.x || b.x )
o = one(BoolWeight)
@test o.x  == true
z = zero(BoolWeight)
@test z.x == false
@test_throws MethodError a/b

# + commutative monoid
@test (a+b)+c ≈ a+(b+c)
@test a+z ≈ z+a ≈ a
@test a+b ≈ b+a
# * monoid
@test (a*b)*c ≈ a*(b*c)
@test a*o == o*a == a # identity
@test a*z == z*a == z # annihilator
# distribution over addition
@test a*(b+c) ≈ (a*b)+(a*c)
@test (a+b)*c ≈ (a*c)+(b*c)

W = BoolWeight
@test FiniteStateTransducers.isidempotent(W)
@test FiniteStateTransducers.iscommulative(W)
@test FiniteStateTransducers.isleft(W)
@test FiniteStateTransducers.isright(W)
@test FiniteStateTransducers.isweaklydivisible(W) == false
@test FiniteStateTransducers.ispath(W)
@test FiniteStateTransducers.reversetype(W) == W

a = parse(W,"true")
@test a.x == true

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

a = LeftStringWeight("ciao")
b = LeftStringWeight("caro")
c = LeftStringWeight("mona")
d = LeftStringWeight(a) # check ProbabilityWeight(x::ProbabilityWeight)

@test (a*b).x == "ciaocaro"
@test (a*c).x == "ciaomona"
@test (a + b).x == "c"
@test (a + c).x == ""

o = one(LeftStringWeight)
@test o.x  == ""
@test o.iszero  == false
z = zero(LeftStringWeight)
@test z.x == ""
@test z.iszero

# + commutative monoid
@test (a+b)+c == a+(b+c)
@test a+z == z+a == a
@test a+b == b+a
# * monoid
@test (a*b)*c == a*(b*c)
@test a*o == o*a == a # identity
@test a*z == z*a == z # annihilator
# distribution over addition
@test a*(b+c) == (a*b)+(a*c)
# division
c = a*b
@test a*(c/a) == c
@test (c/b)*b != c

W = LeftStringWeight
@test FiniteStateTransducers.isidempotent(W)
@test FiniteStateTransducers.iscommulative(W) == false
@test FiniteStateTransducers.isleft(W)
@test FiniteStateTransducers.isright(W) == false
@test FiniteStateTransducers.isweaklydivisible(W)
@test FiniteStateTransducers.ispath(W) == false
@test FiniteStateTransducers.reversetype(W) == RightStringWeight

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

a = RightStringWeight("ciao")
b = RightStringWeight("caro")
c = RightStringWeight("mona")
d = RightStringWeight(a) # check ProbabilityWeight(x::ProbabilityWeight)

@test (a*b).x == "ciaocaro"
@test (a*c).x == "ciaomona"
@test (a + b).x == "o"
@test (a + c).x == ""

o = one(RightStringWeight)
@test o.x  == ""
@test o.iszero == false
z = zero(RightStringWeight)
@test z.x == ""
@test z.iszero

# + commutative monoid
@test (a+b)+c == a+(b+c)
@test a+z == z+a == a
@test a+b == b+a
# * monoid
@test (a*b)*c == a*(b*c)
@test a*o == o*a == a # identity
@test a*z == z*a == z # annihilator
# distribution over addition
@test (a+b)*c == (a*c)+(b*c)
# division
c = a*b
@test a*(c/a) != c
@test (c/b)*b == c

W = RightStringWeight
@test FiniteStateTransducers.isidempotent(W)
@test FiniteStateTransducers.iscommulative(W) == false
@test FiniteStateTransducers.isleft(W) == false
@test FiniteStateTransducers.isright(W)
@test FiniteStateTransducers.isweaklydivisible(W)
@test FiniteStateTransducers.ispath(W) == false
@test FiniteStateTransducers.reversetype(W) == LeftStringWeight

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

W1, W2 = LeftStringWeight, TropicalWeight{Float64}
a = ProductWeight(W1("ciao"),W2(1.0))
b = ProductWeight(W1("caro"),W2(2.0))
c = ProductWeight(W1("mona"),W2(3.0))
d = ProductWeight(a) # check ProbabilityWeight(x::ProbabilityWeight)
dd = typeof(d)(a)

@test (a*b) == ProductWeight(W1("ciaocaro"),W2(3.0))
@test (a*c) == ProductWeight(W1("ciaomona"),W2(4.0))
@test (a+b) == ProductWeight(W1("c"),W2(1.0))
@test (a+c) == ProductWeight(W1(""),W2(1.0))

W = typeof(a)
o = one(W)
@test get(o)[1]  == one(W1)
@test get(o)[2]  == one(W2)
z = zero(typeof(a))
@test get(z)[1] == zero(W1)
@test get(z)[2] == zero(W2)

# + commutative monoid
@test (a+b)+c == a+(b+c)
@test a+z == z+a == a
@test a+b == b+a
# * monoid
@test (a*b)*c == a*(b*c)
@test a*o == o*a == a # identity
@test a*z == z*a == z # annihilator
# distribution over addition
@test a*(b+c) == (a*b)+(a*c)
# division
c = a*b
@test a*(c/a) == c

@test FiniteStateTransducers.isidempotent(W)
@test FiniteStateTransducers.iscommulative(W) == false
@test FiniteStateTransducers.isleft(W)
@test FiniteStateTransducers.isright(W) == false
@test FiniteStateTransducers.isweaklydivisible(W)
@test FiniteStateTransducers.ispath(W) == false
@test FiniteStateTransducers.reversetype(W) == ProductWeight{2,Tuple{RightStringWeight,TropicalWeight{Float64}}}


# ====================  Find Classes  ====================
"""
    find_classes(adj::BitArray{2})

returns an array of length `size(adj,1)` with each entry with index `i` being
a unique identifier for the equivalency class of the node `i`.
Uses adjacency matrix as input. I.e. graph needs to be constructed before.

# Examples
```
julia> find_classes(convert(BitArray, [0 1 0; 1 0 0; 0 0 0]))
[1, 1, 2]
```
"""
function find_classes(adj::BitArray{2})
  vlR = UnitRange{UInt32}(1:size(adj,1))
  classes = collect(vlR)                        # every node has its own class at the start
  openL = trues(size(adj,1))                    # no vertices have been visited yet
  for i in vlR                                  # check all vertices at least once
    !openL[i] && continue                       # if node already checked, continue
    classes[i] = i                              # current node has class i
    indL = [i]                                  # add current vertex to search list
    while !isempty(indL)                        # search through all reachable vertices
      j = pop!(indL)                            # first vertex in index List is next candidate
      !openL[j] && continue
      openL[j] = false                          # set current node to visited
      neighbors = findall(adj[i,:] .& openL)    # add all adjacent AND open vertices to class
      indL = union(neighbors,indL)              # without double entries
      classes[indL] .= i                        # set all neighbors to current class
    end
    sum(openL) == 0 && break                    # no more open vertices left, return
  end
  return classes
end

"""
    find_classes(m::Mapping, vl; vl_len=length(vl), sorted=false)

returns an array of length `size(adj,1)` with each entry with index `i` being
a unique identifier for the equivalency class of the node `i`.
Uses a mapping function `m(x) = y` and an vertex list `vl` to construct graph i
n place.
`m` has to be symmetric (i.e. `m(x) = y => m(y) = x`)!
In case the vertex list is closed under the mapping (i.e. there exist no `x` in
`vl` such that `m(x)` is not in `vl`), `closed` can be set to `false` in order
to improve performace.
for that class. By specifying `continuous = true`,
`vl_len` can be specified explicitly in cases where `length(vl)` does not work
(e.g. nested generator objects).
For performace reasons this function can also be called with a plain function
instead of a Mapping.

# Examples
```
```
"""
find_classes(m::Mapping, vl::AbstractArray{T,1}; vl_len=length(vl)) where T = find_classes(m.f, vl, vl_len=vl_len)

function find_classes(m::Function, vl::AbstractArray{T,1}; vl_len=length(vl)) where T
  vlR = UnitRange{UInt32}(1:vl_len)
  classes = OrderedDict(zip(vl,vlR))
  openL = Dict(zip(vl,trues(vl_len)))           # mark all vertices as open (not visited)
  searchL = Stack{eltype(vl)}()
  @time @inbounds for i in vlR                  # visit each entry at least once
    vi = vl[i]
    if openL[vi]                              # if entry already checked, continue
      push!(searchL, vi)                    # add current vertex to search list
      while !isempty(searchL)               # search through all reachable entries
        vj = pop!(searchL)                # next vertex in index List is next candidate
        classes[vj] = i
        if openL[vj]
          openL[vj] = false             # set current node to visited
          for (j,el) in enumerate(m(vj))           # add all adjacent AND open vertices to class, skip entries outside vl
            if !(el in searchL) && haskey(openL,el) && openL[el]
                push!(searchL, el)
            end
          end
        end
      end
    end
  end
  return @inbounds ReduceMap(classes)
end

"""
    find_classes(pred::Predicate, vl; isSymmetric=false)

returns an array of length `length(vl)` with each entry with index `i` being
a unique identifier for the equivalency class of the node `i`.
See `find_classes(m::Mapping, vl)` for more information.
"""
function find_classes(pred::Predicate, vl::AbstractArray; isSymmetric=false)
    adj = isSymmetric ? build_adj_matrix(pred, vl) :
        build_adj_matrix(Predicate((x,y) ->(pred(x,y) || pred(y,x))), vl)
    @inbounds res = ReduceMap(OrderedDict(zip(vl, find_classes(adj))))
    return res
end

# ====================  Find Paths  ====================
function find_classes(m::Function, vl::AbstractArray{T,1}, operations::Array{UInt32,1}; vl_len=length(vl)) where T
  vlR = UnitRange{UInt32}(1:vl_len)
  parent = Dict(zip(vl,vl))
  parent_edge = Dict(zip(vl,zeros(UInt32,vl_len)))
  openL = Dict(zip(vl,trues(vl_len)))       # mark all vertices as open (not visited)
  searchL = Stack{eltype(vl)}()
  for i in vlR              # visit each entry at least once
    vi = vl[i]
    if openL[vi]                            # if entry already checked, continue
      push!(searchL, vi)                    # add current vertex to search list
      parent[vi] = vi
      parent_edge[vi] = 0x00000
      while !isempty(searchL)               # search through all reachable entries
        vj = pop!(searchL)                  # next vertex in index List is next candidate
        if openL[vj]
          openL[vj] = false                 # set current node to visited
          for (j,el) in enumerate(m(vj))    # add all adjacent AND open vertices to class, skip entries outside vl
              if !(el in searchL) && haskey(openL,el) && openL[el]
                push!(searchL, el)
                parent[el] = vj
                parent_edge[el] = operations[j]
            end
          end
        end
      end
    end
  end
  return parent, parent_edge
end

# ====================  Auxiliary functions ====================

"""
    build_adj_matrix(pred, vl)

Constructs adjacency matrix with vertices from `vl` and edges
from `pred(i,j)` where `i,j ∈ vl`
"""
function build_adj_matrix(pred::Predicate, vl)
  adj = BitArray(undef, length(vl),length(vl))
  for (i,ind) in enumerate(vl)
    adj[i,i] = true
    for (j,ind2) in enumerate(Iterators.drop(vl,i))
      adj[i,j+i] = pred(ind,ind2) || pred(ind2,ind)
      adj[j+i,i] = adj[i,j+i]
    end
  end
  return adj
end

#merge_classes()

# ====================  EquivalenceClasses Wrappers  ====================
"""
    EquivalenceClasses(pred::Predicate, vl::AbstractArray)

Constructs an EquivalenceClasses struct from a predicate `pred`
over `n` indices in the list `vl`.
This inputs requires `n^2` operations in order to create the
adjecency matrix. For large inputs a mapping is therefore
preferabel.

# Examples
```
julia> EquivalenceClasses(Predicate((x,y)->all(x .== -y)),
                          [(i,j) for i in -2:2 for j in 4:7])
```
"""
function EquivalenceClasses(pred::Predicate, vl::AbstractArray)
    return find_classes(pred,vl)
end

"""
    EquivalenceClasses(m::Mapping, vl::AbstractArray)

Constructs an EquivalenceClasses struct from a mapping `m`
over the indices in the list `vl`.

# Examples
```
julia> EquivalenceClasses(Mapping((x)->[-x]),
                          [(i,j) for i in -2:2 for j in 4:7])
```
"""
function EquivalenceClasses(m::Mapping, vl::AbstractArray)
  return find_classes(m,vl)
end

#function parent_to_classes(parent_map::Dict{T,T})
#end


"""
    SimpleRandomWalkEnv

A simple random walk environment for tutorial.

# Example
```julia
julia> env = SimpleRandomWalkEnv()
SimpleRandomWalkEnv(7, 3, 3, [-1, 1])

julia> env(1)
(observation = 2, reward = 0.0, isdone = false)

julia> reset!(env)
(observation = 3, isdone = false)
```
"""
mutable struct SimpleRandomWalkEnv <: AbstractSyncEnvironment{DiscreteSpace,DiscreteSpace,1}
    N::Int
    start::Int
    state::Int
    actions::Vector{Int}
    SimpleRandomWalkEnv(;N=7, actions=[-1, 1], start=3) = new(N, start, start, actions)
end

function (env::SimpleRandomWalkEnv)(a::Int)
    env.state = min(max(env.state + env.actions[a], 1), env.N)
    (observation = env.state,
     reward      = env.state == env.N ? 3.0 : (env.state == 1 ? 1. : 0.),
     isdone      = env.state == env.N || env.state == 1)
end

function reset!(env::SimpleRandomWalkEnv)
    env.state = env.start
    (observation = env.state, isdone = false)
end

observe(env::SimpleRandomWalkEnv) = (observation = env.state, isdone = env.state == env.N || env.state == 1)
observationspace(env::SimpleRandomWalkEnv) = DiscreteSpace(env.N)
actionspace(env::SimpleRandomWalkEnv) = DiscreteSpace(length(env.actions))


### Edge and Gradient related tests ###

using Images, Base.Test, Color

EPS = 1e-14

## Checkerboard array, used to test image gradients

let
    white{T}(::Type{T}) = one(T)
    black{T}(::Type{T}) = zero(T)
    white{T<:Unsigned}(::Type{T}) = typemax(T)
    black{T<:Unsigned}(::Type{T}) = typemin(T)

    global checkerboard
    function checkerboard{T}(::Type{T}, sq_width::Integer, count::Integer)
        wh = fill(white(T), (sq_width,sq_width))
        bk = fill(black(T), (sq_width,sq_width))
        bw = [wh bk; bk wh]
        vert = repmat(bw, (count>>1), 1)
        isodd(count) && (vert = vcat(vert, [wh bk]))
        cb = repmat(vert, 1, (count>>1))
        isodd(count) && (cb = hcat(cb, vert[:,1:sq_width]))
        cb
    end

    checkerboard(sq_width::Integer, count::Integer) = checkerboard(Uint8, sq_width, count)
end

SZ=5

cb_array    = checkerboard(SZ,3)
cb_image_xy = grayim(cb_array)
cb_image_yx = grayim(cb_array)
cb_image_yx["spatialorder"] = ["y","x"]

cb_image_gray = convert(Image{Gray}, cb_image_xy)
cb_image_rgb = convert(Image{RGB}, cb_image_xy)
cb_image_rgb2 = convert(Image{RGB{Float64}}, cb_image_xy)

for method in ["sobel", "prewitt", "ando3", "ando4", "ando5", "ando4_sep", "ando5_sep"]
    ## Checkerboard array

    (agx, agy) = imgradients(cb_array, method)
    amag = magnitude(agx, agy)
    agphase = phase(agx, agy)
    @test (amag, agphase) == magnitude_phase(agx, agy)

    @test agx[1,SZ]   < 0.0   # white to black transition
    @test agx[1,2*SZ] > 0.0   # black to white transition
    @test agy[SZ,1]   < 0.0   # white to black transition
    @test agy[2*SZ,1] > 0.0   # black to white transition

    # Test direction of increasing gradient
    @test cos(agphase[1,SZ])   - (-1.0) < EPS   # increasing left  (=  pi   radians)
    @test cos(agphase[1,2*SZ]) -   1.0  < EPS   # increasing right (=   0   radians)
    @test sin(agphase[SZ,1])   -   1.0  < EPS   # increasing up    (=  pi/2 radians)
    @test sin(agphase[2*SZ,1]) - (-1.0) < EPS   # increasing down  (= -pi/2 radians)

    # Test that orientation is perpendicular to gradient
    aorient = orientation(agx, agy)
    @test all((cos(agphase).*cos(aorient) .+ sin(agphase).*sin(aorient) .< EPS) |
                ((agphase .== 0.0) & (aorient .== 0.0)))  # this part is where both are
                                                          # zero because there is no gradient

    ## Checkerboard Image with row major order

    (gx, gy) = imgradients(cb_image_xy, method)
    mag = magnitude(gx, gy)
    gphase = phase(gx, gy)
    @test (mag, gphase) == magnitude_phase(gx, gy)

    @test gx[SZ,1]   < 0.0   # white to black transition
    @test gx[2*SZ,1] > 0.0   # black to white transition
    @test gy[1,SZ]   < 0.0   # white to black transition
    @test gy[1,2*SZ] > 0.0   # black to white transition

    @test cos(gphase[SZ,1])   - (-1.0) < EPS   # increasing left  (=  pi   radians)
    @test cos(gphase[2*SZ,1]) -   1.0  < EPS   # increasing right (=   0   radians)
    @test sin(gphase[1,SZ])   -   1.0  < EPS   # increasing up    (=  pi/2 radians)
    @test sin(gphase[1,2*SZ]) - (-1.0) < EPS   # increasing down  (= -pi/2 radians)

    # Test that orientation is perpendicular to gradient
    orient = orientation(gx, gy)
    @test all((cos(gphase).*cos(orient) .+ sin(gphase).*sin(orient) .< EPS) |
                ((gphase .== 0.0) & (orient .== 0.0)))  # this part is where both are
                                                        # zero because there is no gradient

    ## Checkerboard Image with column-major order

    (gx, gy) = imgradients(cb_image_yx, method)
    mag = magnitude(gx, gy)
    gphase = phase(gx, gy)
    @test (mag, gphase) == magnitude_phase(gx, gy)

    @test gx[1,SZ]   < 0.0   # white to black transition
    @test gx[1,2*SZ] > 0.0   # black to white transition
    @test gy[SZ,1]   < 0.0   # white to black transition
    @test gy[2*SZ,1] > 0.0   # black to white transition

    # Test direction of increasing gradient
    @test cos(gphase[1,SZ])   - (-1.0) < EPS   # increasing left  (=  pi   radians)
    @test cos(gphase[1,2*SZ]) -   1.0  < EPS   # increasing right (=   0   radians)
    @test sin(gphase[SZ,1])   -   1.0  < EPS   # increasing up    (=  pi/2 radians)
    @test sin(gphase[2*SZ,1]) - (-1.0) < EPS   # increasing down  (= -pi/2 radians)

    # Test that orientation is perpendicular to gradient
    orient = orientation(gx, gy)
    @test all((cos(gphase).*cos(orient) .+ sin(gphase).*sin(orient) .< EPS) |
                ((gphase .== 0.0) & (orient .== 0.0)))  # this part is where both are
                                                        # zero because there is no gradient

    ## Checkerboard Image with Gray pixels

    (gxg, gyg) = imgradients(cb_image_gray, method)
    magg = magnitude(gxg, gyg)
    gphaseg = phase(gxg, gyg)
    @test (magg, gphaseg) == magnitude_phase(gxg, gyg)

    @test gxg[SZ,1]   < 0.0   # white to black transition
    @test gxg[2*SZ,1] > 0.0   # black to white transition
    @test gyg[1,SZ]   < 0.0   # white to black transition
    @test gyg[1,2*SZ] > 0.0   # black to white transition

    @test cos(gphaseg[SZ,1])   - (-1.0) < EPS   # increasing left  (=  pi   radians)
    @test cos(gphaseg[2*SZ,1]) -   1.0  < EPS   # increasing right (=   0   radians)
    @test sin(gphaseg[1,SZ])   -   1.0  < EPS   # increasing up    (=  pi/2 radians)
    @test sin(gphaseg[1,2*SZ]) - (-1.0) < EPS   # increasing down  (= -pi/2 radians)

    # Test that orientation is perpendicular to gradient
    orientg = orientation(gxg, gyg)
    @test all((cos(gphaseg).*cos(orientg) .+ sin(gphaseg).*sin(orientg) .< EPS) |
                ((gphaseg .== 0.0) & (orientg .== 0.0)))  # this part is where both are
                                                        # zero because there is no gradient

    ## Checkerboard Image with RBG pixels

    (gx_rgb, gy_rgb) = imgradients(cb_image_rgb, method)
    mag_rgb = magnitude(gx_rgb, gy_rgb)
    gphase_rgb = phase(gx_rgb, gy_rgb)
    @test (mag_rgb, gphase_rgb) == magnitude_phase(gx_rgb, gy_rgb)

    @test gx_rgb[1,SZ,1]   < 0.0   # white to black transition
    @test gx_rgb[1,2*SZ,1] > 0.0   # black to white transition
    @test gy_rgb[1,1,SZ]   < 0.0   # white to black transition
    @test gy_rgb[1,1,2*SZ] > 0.0   # black to white transition

    @test cos(gphase_rgb[1,SZ,1])   - (-1.0) < EPS   # increasing left  (=  pi   radians)
    @test cos(gphase_rgb[1,2*SZ,1]) -   1.0  < EPS   # increasing right (=   0   radians)
    @test sin(gphase_rgb[1,1,SZ])   -   1.0  < EPS   # increasing up    (=  pi/2 radians)
    @test sin(gphase_rgb[1,1,2*SZ]) - (-1.0) < EPS   # increasing down  (= -pi/2 radians)

    # Test that orientation is perpendicular to gradient
    orient_rgb = orientation(gx_rgb, gy_rgb)
    @test all((cos(gphase_rgb).*cos(orient_rgb) .+ sin(gphase_rgb).*sin(orient_rgb) .< EPS) |
                ((gphase_rgb .== 0.0) & (orient_rgb .== 0.0)))  # this part is where both are
                                                                # zero because there is no gradient

    ## Checkerboard Image with RBG{Float64} pixels

    (gx_rgb, gy_rgb) = imgradients(cb_image_rgb2, method)
    mag_rgb = magnitude(gx_rgb, gy_rgb)
    gphase_rgb = phase(gx_rgb, gy_rgb)
    @test (mag_rgb, gphase_rgb) == magnitude_phase(gx_rgb, gy_rgb)

    @test gx_rgb[1,SZ,1]   < 0.0   # white to black transition
    @test gx_rgb[1,2*SZ,1] > 0.0   # black to white transition
    @test gy_rgb[1,1,SZ]   < 0.0   # white to black transition
    @test gy_rgb[1,1,2*SZ] > 0.0   # black to white transition

    @test cos(gphase_rgb[1,SZ,1])   - (-1.0) < EPS   # increasing left  (=  pi   radians)
    @test cos(gphase_rgb[1,2*SZ,1]) -   1.0  < EPS   # increasing right (=   0   radians)
    @test sin(gphase_rgb[1,1,SZ])   -   1.0  < EPS   # increasing up    (=  pi/2 radians)
    @test sin(gphase_rgb[1,1,2*SZ]) - (-1.0) < EPS   # increasing down  (= -pi/2 radians)

    # Test that orientation is perpendicular to gradient
    orient_rgb = orientation(gx_rgb, gy_rgb)
    @test all((cos(gphase_rgb).*cos(orient_rgb) .+ sin(gphase_rgb).*sin(orient_rgb) .< EPS) |
                ((gphase_rgb .== 0.0) & (orient_rgb .== 0.0)))  # this part is where both are
                                                                # zero because there is no gradient
end

# Create an image with white along diagonals -2:2 and black elsewhere
m = zeros(Uint8, 20,20)
for i = -2:2; m[diagind(m,i)] = 0xff; end

m_xy = grayim(m')
m_yx = grayim(m)
m_yx["spatialorder"] = ["y","x"]

for method in ["sobel", "prewitt", "ando3", "ando4", "ando5", "ando4_sep", "ando5_sep"]
    ## Diagonal array

    (agx, agy) = imgradients(m, method)
    amag = magnitude(agx, agy)
    agphase = phase(agx, agy)
    @test (amag, agphase) == magnitude_phase(agx, agy)

    @test agx[7,9]  < 0.0   # white to black transition
    @test agx[10,8] > 0.0   # black to white transition
    @test agy[10,8] < 0.0   # white to black transition
    @test agy[7,9]  > 0.0   # black to white transition

    # Test direction of increasing gradient
    @test abs(agphase[10,8] -    pi/4 ) < EPS   # lower edge (increasing up-right  =   pi/4 radians)
    @test abs(agphase[7,9]  - (-3pi/4)) < EPS   # upper edge (increasing down-left = -3pi/4 radians)

    # Test that orientation is perpendicular to gradient
    aorient = orientation(agx, agy)
    @test all((cos(agphase).*cos(aorient) .+ sin(agphase).*sin(aorient) .< EPS) |
                ((agphase .== 0.0) & (aorient .== 0.0)))  # this part is where both are
                                                          # zero because there is no gradient

    ## Diagonal Image, row-major order

    (gx, gy) = imgradients(m_xy, method)
    mag = magnitude(gx, gy)
    gphase = phase(gx, gy)
    @test (mag, gphase) == magnitude_phase(gx, gy)

    @test gx[9,7]  < 0.0   # white to black transition
    @test gx[8,10] > 0.0   # black to white transition
    @test gy[8,10] < 0.0   # white to black transition
    @test gy[9,7]  > 0.0   # black to white transition

    # Test direction of increasing gradient
    @test abs(gphase[8,10] -    pi/4 ) < EPS   # lower edge (increasing up-right  =   pi/4 radians)
    @test abs(gphase[9,7]  - (-3pi/4)) < EPS   # upper edge (increasing down-left = -3pi/4 radians)

    # Test that orientation is perpendicular to gradient
    orient = orientation(gx, gy)
    @test all((cos(gphase).*cos(orient) .+ sin(gphase).*sin(orient) .< EPS) |
                ((gphase .== 0.0) & (orient .== 0.0)))  # this part is where both are
                                                        # zero because there is no gradient

    ## Diagonal Image, column-major order

    (gx, gy) = imgradients(m_yx, method)
    mag = magnitude(gx, gy)
    gphase = phase(gx, gy)
    @test (mag, gphase) == magnitude_phase(gx, gy)

    @test gx[7,9]  < 0.0   # white to black transition
    @test gx[10,8] > 0.0   # black to white transition
    @test gy[10,8] < 0.0   # white to black transition
    @test gy[7,9]  > 0.0   # black to white transition

    # Test direction of increasing gradient
    @test abs(gphase[10,8] -    pi/4 ) < EPS   # lower edge (increasing up-right  =   pi/4 radians)
    @test abs(gphase[7,9]  - (-3pi/4)) < EPS   # upper edge (increasing down-left = -3pi/4 radians)

    # Test that orientation is perpendicular to gradient
    orient = orientation(gx, gy)
    @test all((cos(gphase).*cos(orient) .+ sin(gphase).*sin(orient) .< EPS) |
                ((gphase .== 0.0) & (orient .== 0.0)))  # this part is where both are
                                                        # zero because there is no gradient
end

# Nonmaximal suppression

function thin_edges(img)
    # Get orientation
    gx,gy = imgradients(img)
    orient = phase(gx,gy)

    # Do NMS thinning
    thin_edges_nonmaxsup_subpix(img, orient, radius=1.35)
end

function nms_test_horiz_vert(img, which)
    ## which = :horizontal or :vertical

    # Do NMS thinning
    t,s = thin_edges(img)

    # Calc peak location by hand

    # Interpolate values 1.35 pixels left and right
    # Orientation is zero radians -> to the right
    v1 = 6 - 0.35   # slope on right is 1
    v2 = 5 - 0.35*2 # slope on left is 2
    c = 7.0         # peak value

    # solve v = a*r^2 + b*r + c
    a = (v1 + v2)/2 - c
    b = a + c - v2
    r = -b/2a

    @test abs(r - 1/6) < EPS

    # Location and value at peak
    peakloc = r*1.35 + 3
    peakval = a*r^2 + b*r + c

    transposed = spatialorder(img)[1] == "x"
    horizontal = which == :horizontal

    test_axis1 = transposed $ !horizontal

    @test test_axis1 ? all(t[:,[1,2,4,5]] .== 0) : all(t[[1,2,4,5],:] .== 0)
    @test test_axis1 ? all(t[:,3]   .== peakval) : all(t[3,:]   .== peakval)
    @test test_axis1 ? all(s[:,[1,2,4,5]] .== zero(Graphics.Point)) :
                         all(s[[1,2,4,5],:] .== zero(Graphics.Point))

    if transposed
        if which == :horizontal
            @test     [pt.x for pt in s[:,3]]  == [1:5;]
            @test all([pt.y for pt in s[:,3]] .== peakloc)
        else
            @test all([pt.x for pt in s[3,:]] .== peakloc)
            @test     [pt.y for pt in s[3,:]]  == [1:5;]
        end
    else
        if which == :horizontal
            @test     [pt.x for pt in s[3,:]]  == [1:5;]
            @test all([pt.y for pt in s[3,:]] .== peakloc)
        else
            @test all([pt.x for pt in s[:,3]] .== peakloc)
            @test     [pt.y for pt in s[:,3]]  == [1:5;]
        end
    end
end

# Test image: vertical edge
m = [3.0  5.0  7.0  6.0  5.0
     3.0  5.0  7.0  6.0  5.0
     3.0  5.0  7.0  6.0  5.0
     3.0  5.0  7.0  6.0  5.0
     3.0  5.0  7.0  6.0  5.0]

m_xy = grayim(m')
m_yx = grayim(m)
m_yx["spatialorder"] = ["y","x"]

nms_test_horiz_vert(m, :vertical)
nms_test_horiz_vert(m_xy, :vertical)
nms_test_horiz_vert(m_yx, :vertical)

# Test image: horizontal edge
m = m'
m_xy = grayim(m')
m_yx = grayim(m)
m_yx["spatialorder"] = ["y","x"]

nms_test_horiz_vert(m, :horizontal)
nms_test_horiz_vert(m_xy, :horizontal)
nms_test_horiz_vert(m_yx, :horizontal)

function nms_test_diagonal(img)
    # Do NMS thinning
    t,s = thin_edges(img)

    # Calc peak location by hand

    # Interpolate values 1.35 pixels up and left, down and right
    # using bilinear interpolation
    # Orientation is π/4 radians -> 45 degrees up
    fr = 1.35*cos(π/4)
    lower = (7 + fr*(6-7))
    upper = (6 + fr*(5-6))
    v1 = lower + fr*(upper-lower)

    lower = (7 + fr*(5-7))
    upper = (5 + fr*(3-5))
    v2 = lower + fr*(upper-lower)

    c = 7.0         # peak value

    # solve v = a*r^2 + b*r + c
    a = (v1 + v2)/2 - c
    b = a + c - v2
    r = -b/2a

    @test (r - 1/6) < EPS

    transposed = spatialorder(img)[1] == "x"

    # Location and value at peak

    x_peak_offset, y_peak_offset = r*fr, -r*fr
    peakval = a*r^2 + b*r + c

    @test all(diag(data(t))[2:4] .== peakval)  # Edge pixels aren't interpolated here
    @test all(t - diagm(diag(data(t))) .== 0)

    diag_s = copyproperties(s, diagm(diag(data(s))))
    @test s == diag_s

    @test all([pt.x for pt in diag(data(s))[2:4]] - ((2:4) + x_peak_offset) .< EPS)
    @test all([pt.y for pt in diag(data(s))[2:4]] - ((2:4) + y_peak_offset) .< EPS)

end

# Test image: diagonal edge
m = [7.0  6.0  5.0  0.0  0.0
     5.0  7.0  6.0  5.0  0.0
     3.0  5.0  7.0  6.0  5.0
     0.0  3.0  5.0  7.0  6.0
     0.0  0.0  3.0  5.0  7.0]

m_xy = grayim(m')
m_yx = grayim(m)
m_yx["spatialorder"] = ["y","x"]

nms_test_diagonal(m)
nms_test_diagonal(m_xy)
nms_test_diagonal(m_yx)


using Distributions
using Klara

CURRENTDIR, CURRENTFILE = splitdir(@__FILE__)
ROOTDIR = splitdir(splitdir(CURRENTDIR)[1])[1]
OUTDIR = joinpath(ROOTDIR, "output")

# OUTDIR = "../output"

SUBOUTDIR = "SMMALA"

nchains = 10
nmcmc = 110000
nburnin = 10000

function C(n::Int, c::Float64)
  X = eye(n)
  [(j <= n-i) ? X[i+j, i] = X[i, i+j] = c^j : nothing for i = 1:(n-1), j = 1:(n-1)]
  X
end

n = 20
μ = zeros(n)
Σ = C(n, 0.9)
ν = 30.

Σt = (ν-2)*Σ/ν
Σtinv = inv(Σt)

plogtarget(p::Vector, v::Vector) = logpdf(MvTDist(ν, zeros(n), (ν-2)*Σ/ν), p)

p = BasicContMuvParameter(:p, logtarget=plogtarget, nkeys=1, diffopts=DiffOptions(mode=:forward, order=2))

model = likelihood_model([p], isindexed=false)

sampler = SMMALA(0.02, H -> softabs(H, 1000.))

mcrange = BasicMCRange(nsteps=nmcmc, burnin=nburnin)

mctuner = AcceptanceRateMCTuner(0.7, score=x -> logistic_rate_score(x, 3.), verbose=false)

outopts = Dict{Symbol, Any}(:monitor=>[:value], :diagnostics=>[:accept])

times = Array{Float64}(nchains)
stepsizes = Array{Float64}(nchains)
i = 1

while i <= nchains
  v0 = Dict(:p=>rand(Normal(0, 2), n))

  job = BasicMCJob(model, sampler, mcrange, v0, tuner=mctuner, outopts=outopts)

  tic()
  run(job)
  runtime = toc()

  chain = output(job)
  ratio = acceptance(chain)

  if 0.55 < ratio < 0.85
    writedlm(joinpath(OUTDIR, SUBOUTDIR, "chain"*lpad(string(i), 2, 0)*".csv"), chain.value, ',')
    writedlm(joinpath(OUTDIR, SUBOUTDIR, "diagnostics"*lpad(string(i), 2, 0)*".csv"), vec(chain.diagnosticvalues), ',')

    times[i] = runtime
    stepsizes[i] = job.sstate.tune.step

    println("Iteration ", i, " of ", nchains, " completed with acceptance ratio ", ratio)
    i += 1
  end
end

writedlm(joinpath(OUTDIR, SUBOUTDIR, "times.csv"), times, ',')
writedlm(joinpath(OUTDIR, SUBOUTDIR, "stepsizes.csv"), stepsizes, ',')