using Distributions, Random
using Test, LinearAlgebra, PDMats

#  dimensions

n = 6
p = 3

#  parameters

v    = 10
M    = randn(n, p)
Σ    = rand(InverseWishart(n + 2, Matrix(1.0I, n, n)))
Ω    = rand(InverseWishart(p + 2, Matrix(1.0I, p, p)))
ΣF32 = Matrix{Float32}(Σ)
ΩBF  = Matrix{BigFloat}(Ω)
PDΣ  = PDMat(Σ)
PDΩ  = PDMat(Ω)
μ    = randn(n)
θ    = randn(p)
σ²   = 2.0rand()

#  distribution instances

D = MatrixTDist(v, M, Σ, Ω)                                   #  n x p
L = MatrixTDist(v, reshape(μ, n, 1), Σ, reshape([σ²], 1, 1))  #  n x 1
H = MatrixTDist(v, reshape(θ, 1, p), reshape([σ²], 1, 1), Ω)  #  1 x p
K = MatrixTDist(1000, M, 1000Σ, Ω)

l  = MvTDist(L)
h  = MvTDist(H)
MN = MatrixNormal(M, Σ, Ω)

#  draws

A = rand(D)

X = rand(L)
Y = rand(H)

x = vec(X)
y = vec(Y)

@testset "Check all MatrixTDist constructors" begin
    @test MatrixTDist(v, M, Σ,        PDΩ.chol) isa MatrixTDist
    @test MatrixTDist(v, M, Σ,        PDΩ)      isa MatrixTDist
    @test MatrixTDist(v, M, PDΣ.chol, Ω)        isa MatrixTDist
    @test MatrixTDist(v, M, PDΣ.chol, PDΩ.chol) isa MatrixTDist
    @test MatrixTDist(v, M, PDΣ.chol, PDΩ)      isa MatrixTDist
    @test MatrixTDist(v, M, PDΣ,      Ω)        isa MatrixTDist
    @test MatrixTDist(v, M, PDΣ,      PDΩ.chol) isa MatrixTDist
    @test MatrixTDist(v, M, PDΣ,      PDΩ)      isa MatrixTDist
end

@testset "MatrixTDist promotion during construction" begin
    R = MatrixTDist(v, M, ΣF32, ΩBF)
    @test partype(R) == BigFloat
end

@testset "MatrixTDist construction errors" begin
    @test_throws ArgumentError MatrixTDist(-1, M, Σ, Ω)
    @test_throws ArgumentError MatrixTDist(Inf, M, Σ, Ω)
    @test_throws ArgumentError MatrixTDist(v, M, Ω, Σ)
    @test_throws ArgumentError MatrixTDist(v, M, Ω, Ω)
    @test_throws ArgumentError MatrixTDist(v, M, Σ, Σ)
end

@testset "MatrixTDist params" begin
    vv, MM, PDΣΣ, PDΩΩ = params(D)
    @test v == vv
    @test M == MM
    @test Σ == PDΣΣ.mat
    @test Ω == PDΩΩ.mat
end

@testset "MatrixTDist size" begin
    @test size(D) == (n, p)
    @test size(L) == (n, 1)
    @test size(H) == (1, p)
end

@testset "MatrixTDist rank" begin
    @test rank(D) == min(n, p)
    @test rank(L) == 1
    @test rank(H) == 1
    @test rank(D) == rank(rand(D))
    @test rank(L) == rank(rand(L))
    @test rank(H) == rank(rand(H))
end

@testset "MatrixTDist insupport" begin
    @test insupport(D, rand(D))
    @test insupport(L, rand(L))
    @test insupport(H, rand(H))

    @test !insupport(D, Matrix(rand(D)'))
    @test !insupport(D, M + M * im)
    @test !insupport(L, rand(H))
    @test !insupport(H, rand(L))
end

@testset "MatrixTDist mean" begin
    @test mean(D) == M
    @test mean(L) == reshape(μ, n, 1)
    @test mean(H) == reshape(θ, 1, p)

    @test_throws ArgumentError mean( MatrixTDist(1, M, Σ, Ω) )
end

@testset "MatrixTDist mode" begin
    @test mode(D) == M
    @test mode(L) == reshape(μ, n, 1)
    @test mode(H) == reshape(θ, 1, p)
end

@testset "MatrixTDist cov and var" begin
    @test vec(var(D)) ≈ diag(cov(D))
    @test reshape(cov(D), n, p, n, p) ≈ cov(D, Val(false))
    @test cov(L) ≈ cov(l)
    @test cov(H) ≈ cov(h)

    @test_throws ArgumentError cov( MatrixTDist(1, M, Σ, Ω) )
    @test_throws ArgumentError var( MatrixTDist(1, M, Σ, Ω) )
end

@testset "MatrixTDist logpdf" begin
    #  Check against MvTDist
    @test logpdf(L, X) ≈ logpdf(l, x)
    @test logpdf(H, Y) ≈ logpdf(h, y)

    #  MT(v, M, vΣ, Ω) -> MN(M, Σ, Ω) as v -> ∞
    @test isapprox(logpdf(K, A), logpdf(MN, A), atol = 0.1)
end

@testset "MatrixTDist sample moments" begin
    @test isapprox(mean(rand(D, 100000)), mean(D) , atol = 0.1)
    @test isapprox(cov(hcat(vec.(rand(D, 100000))...)'), cov(D), atol = 0.1)
end

@testset "MatrixTDist conversion" for elty in (Float32, Float64, BigFloat)
    Del1 = convert(MatrixTDist{elty}, D)
    Del2 = convert(MatrixTDist{elty}, v, M, PDΣ, PDΩ, D.logc0)

    @test partype(Del1) == elty
    @test partype(Del2) == elty
end


@testset "Linear Discrete System" begin
    # Create some linear data
    A = [0.9 -0.2; 0.0 0.2]
    y = [[10.; -10.]]
    for i in 1:10
        push!(y, A*y[end])
    end
    X = hcat(y...)
    prob = DiscreteDataDrivenProblem(X, t = 1:11)

    for alg in [DMDPINV(), DMDSVD(), TOTALDMD()]
        # Returns a named
        estimator = solve(prob, alg , operator_only = true)
        @test Matrix(estimator.K) ≈ A
        @test eigvals(estimator.K) ≈ eigvals(A)
        @test estimator.C ≈ diagm(ones(2))
        @test isempty(estimator.B)

        res = solve(prob, alg , operator_only = false)
        m = metrics(res)
        @test all(m[:L₂] ./ size(X, 2) .< 3e-1)
        @test Matrix(result(res)) ≈ Matrix(estimator.K)
    end
end

@testset "Linear Continuous System" begin
    A = [-0.9 0.1; 0.0 -0.2]
    f(u, p, t) = A*u
    u0 = [10.0; -20.0]
    prob = ODEProblem(f, u0, (0.0, 10.0))
    sol = solve(prob, Tsit5(), saveat = 0.001)

    prob = DataDrivenProblem(sol)

    for alg in [DMDPINV(), DMDSVD(), TOTALDMD()]
        estimator = solve(prob, alg , operator_only = true)
        @test isapprox(Matrix(estimator.K), A, atol = 1e-2)
        @test isapprox(eigvals(estimator.K), eigvals(A), atol = 1e-2)
        @test estimator.C ≈ diagm(ones(2))
        @test isempty(estimator.B)
        res = solve(prob, alg , operator_only = false)
        m = metrics(res)
        @test all(m[:L₂] ./ size(X, 2) .< 3e-1)
        @test Matrix(result(res)) ≈ Matrix(estimator.K)
    end
end

@testset "Low Rank Continuous System" begin
    K̃ = -0.5*I + [0 0 -0.2; 0.1 0 -0.1; 0. -0.2 0]
    F = qr(randn(20, 3))
    Q = F.Q[:, 1:3]
    dudt(u, p, t) = K̃*u
    prob = ODEProblem(dudt, [10.0; 0.3; -5.0], (0.0, 10.0))
    sol_ = solve(prob, Tsit5(), saveat = 0.01)

    # True Rank is 3
    X = Q*sol_[:,:] + 1e-3*randn(20, 1001)
    DX = Q*sol_(sol_.t, Val{1})[:,:] + 1e-3*randn(20, 1001)
    ddprob = ContinuousDataDrivenProblem(X, sol_.t, DX = DX)

    for alg in [TOTALDMD(3, DMDPINV()); TOTALDMD(0.01, DMDSVD(3))]
        res = solve(ddprob, alg, digits = 2)
        K = Matrix(result(res))
        m = metrics(res)
        @test Q'*K*Q ≈ K̃ atol = 1e-1
        @test Q*K̃*Q' ≈ K atol = 1e-1
        @test all(m[:L₂] ./ size(X, 2) .< 1e-2)
    end
end

@testset "Big System" begin
    # Creates a big system which would resulting in a segfault otherwise
    X = rand([0, 1], 128, 936);
    T = collect(LinRange(0, 4.367058580858928, 936));
    problem = DiscreteDataDrivenProblem(X, T);
    @test_nowarn res2 = solve(problem, DMDSVD(), operator_only = true)
end


# Implementation of
#  "Table-driven Implementation of the Logarithm Function in IEEE Floating-point Arithmetic"
#  Tang, Ping-Tak Peter
#  ACM Trans. Math. Softw. (1990), 16(4):378--400
#  http://dx.doi.org/10.1145/98267.98294

# Other implementations:
#  Apple OS X libm (APSL v2)
#   http://www.opensource.apple.com/source/Libm/Libm-2026/Source/Intel/xmm_log.c
#  libacml_mv (GPL v2)
#   http://svn.open64.net/svnroot/open64/trunk/osprey/libacml_mv/src/log1p.c
#  libclc (BSD): modified to use different table size
#   https://github.com/llvm-mirror/libclc/blob/master/generic/lib/math/log1p.cl

# At the moment we do not correctly handle floating point flags (inexact, div-by-zero, etc).

import Base.unsafe_trunc

const t_log_64 = Array((Float64,Float64),129)
N=39 # (can be up to N=42, which appears to be what Apple's libm uses).
sN = 2.0^N
isN = 1.0/sN
s7 = 2.0^7
is7 = 1.0/s7

for j=0:128
    l_big = Base.log(big(1.0+j*is7))
    l_hi = isN*float64(round(sN*l_big))
    l_lo = float64(l_big-l_hi)
    t_log_64[j+1] = (l_hi,l_lo)
    # c_invF[j+1] = 1.0+is7*j
end

const t_log_32 = Array((Float32,Float32),129)
N=16
sN = 2f0^N
isN = 1f0/sN
for j=0:128
    l_big = Base.log(big(1.0+j*is7))
    l_hi = isN*float32(round(sN*l_big))
    l_lo = float32(l_big-l_hi)
    t_log_32[j+1] = (l_hi,l_lo)
end

const tt_log_32 = Array(Float64,129)
N=16
sN = 2f0^N
isN = 1f0/sN
for j=0:128
    tt_log_32[j+1] = float64(Base.log(big(1.0+j*is7)))
end

# determine if hardware FMA is available
# should probably check with LLVM, see #9855.
const FMA_NATIVE = muladd(nextfloat(1.0),nextfloat(1.0),-nextfloat(1.0,2)) == -4.930380657631324e-32

# truncate lower order bits (up to 26)
# ideally, this should be able to use ANDPD instructions, see #9868.
@inline function truncbits(x::Float64)
    reinterpret(Float64, reinterpret(UInt64,x) & 0xffff_ffff_f800_0000)
end

# Procedure 1
@inline function log_proc1(y::Float64,mf::Float64,F::Float64,f::Float64,jp::Int)
    ## Steps 1 and 2
    @inbounds hi,lo = libm.t_log_64[jp]
    l_hi = mf* 0x1.62e42fefa4p-1 + hi
    l_lo = mf*-0x1.8432a1b0e2634p-43 + lo

    ## Step 3
    # @inbounds u = f*c_invF[jp]
    # u = f/F
    # q = u*u*@horner(u,
    #                 -0x1.0_0000_0000_0001p-1,
    #                 +0x1.5_5555_5550_9ba5p-2,
    #                 -0x1.f_ffff_ffeb_6526p-3,
    #                 +0x1.9_99b4_dfed_6fe4p-3,
    #                 -0x1.5_5576_6647_2e04p-3)

    ## Step 3' (alternative)
    u = (2.0f)/(y+F)
    v = u*u
    q = u*v*@horner(v,
                    0x1.5_5555_5555_0286p-4,
                    0x1.9_99a0_bc71_2416p-7)

    ## Step 4
    l_hi + (u + (q + l_lo))
end

# Procedure 2
@inline function log_proc2(f::Float64)
    ## Step 1
    g = 1.0/(2.0+f)
    u = 2.0*f*g
    v = u*u

    ## Step 2
    q = u*v*@horner(v,
                    0x1.5_5555_5555_54e6p-4,
                    0x1.9_9999_99ba_c6d4p-7,
                    0x1.2_4923_07f1_519fp-9,
                    0x1.c_8034_c85d_fff0p-12)

    ## Step 3
    # based on:
    #   2(f-u) = 2(f(2+f)-2f)/(2+f) = 2f^2/(2+f) = fu
    #   2(f-u1-u2) - f*(u1+u2) = 0
    #   2(f-u1) - f*u1 = (2+f)u2
    #   u2 = (2(f-u1) - f*u1)/(2+f)
    if FMA_NATIVE
        return u + fma(fma(-u,f,2(f-u)), g, q)
    else
        u1 = truncbits(u) # round to 24 bits
        f1 = truncbits(f)
        f2 = f-f1
        u2 = ((2.0*(f-u1)-u1*f1)-u1*f2)*g
        ## Step 4
        return u1 + (u2 + q)
    end
end

@inline function log_proc1(y::Float32,mf::Float32,F::Float32,f::Float32,jp::Int)
    ## Steps 1 and 2
    @inbounds hi = tt_log_32[jp]
    l = mf*0.6931471805599453 + hi

    ## Step 3
    # @inbounds u = f*c_invF[jp]
    # q = u*u*@horner(u,
    #                 Float32(-0x1.00006p-1),
    #                 Float32(0x1.55546cp-2))

    ## Step 3' (alternative)
    u = (2f0f)/(y+F)
    v = u*u
    q = u*v*Float32(0x1.555584p-4)

    ## Step 4
    float32(l + (u + q))
end

@inline function log_proc2(f::Float32)
    ## Step 1
    # compute in higher precision
    u64 = Float64(2f0*f)/(2.0+f)
    u = Float32(u64)
    v = u*u

    ## Step 2
    q = u*v*@horner(v,
                    Float32(0x1.555552p-4),
                    Float32(0x1.9a012ap-7))

    ## Step 3: not required

    ## Step 4
    Float32(u64 + q)
end

function log_tang(x::Float64)
    if x > 0.0
        x == Inf && return x

        # Step 2
        if 0x1.e_0fab_fbc7_02a3p-1 < x < 0x1.1_082b_577d_34eep0
            f = x-1.0
            return log_proc2(f)
        end

        # Step 3
        xu = reinterpret(UInt64,x)
        m = int(xu >> 52) & 0x07ff
        if m == 0 # x is subnormal
            x *= 0x1p54 # normalise significand
            xu = reinterpret(UInt64,x)
            m = int(xu >> 52) & 0x07ff - 54
        end
        m -= 1023
        y = reinterpret(Float64,(xu & 0x000f_ffff_ffff_ffff) | 0x3ff0_0000_0000_0000)

        mf = Float64(m)
        F = (y + 0x1p45) - 0x1p45 # 0x1p-7*round(0x1p7*y)
        f = y-F
        jp = unsafe_trunc(Int,0x1p7*F)-127

        return log_proc1(y,mf,F,f,jp)
    elseif x == 0.0
        -Inf
    elseif isnan(x)
        NaN
    else
        throw(DomainError())
    end
end

function log_tang(x::Float32)
    if x > 0f0
        x == Inf32 && return x

        # Step 2
        if Float32(0x1.e0fabep-1) < x < Float32(0x1.1082b6p+0)
            f = x-1f0
            return log_proc2(f)
        end

        # Step 3
        xu = reinterpret(UInt32,x)
        m = int(xu >> 23) & 0x00ff
        if m == 0 # x is subnormal
            x *= Float32(0x1p25) # normalise significand
            xu = reinterpret(UInt32,x)
            m = int(xu >> 23) & 0x00ff - 25
        end
        m -= 127
        y = reinterpret(Float32,(xu & 0x007f_ffff) | 0x3f80_0000)

        mf = Float32(m)
        F = (y + Float32(0x1p16)) - Float32(0x1p16) # 0x1p-7*round(0x1p7*y)
        f = y-F
        jp = unsafe_trunc(Int,Float32(0x1p7)*F)-127

        log_proc1(y,mf,F,f,jp)
    elseif x == 0f0
        -Inf32
    elseif isnan(x)
        NaN32
    else
        throw(DomainError())
    end
end

function log1p_tang(x::Float64)
    if x > -1.0
        x == Inf && return x
        if -0x1p-53 < x < 0x1p-53
            return x

        # Step 2
        elseif -0x1.f_0540_438f_d5c4p-5 < x < 0x1.0_82b5_77d3_4ed8p-4
            return log_proc2(x)
        end

        # Step 3
        z = 1.0 + x
        zu = reinterpret(UInt64,z)
        s = reinterpret(Float64,0x7fe0_0000_0000_0000 - (zu & 0xfff0_0000_0000_0000)) # 2^-m
        m = int(zu >> 52) & 0x07ff - 1023 # z cannot be subnormal
        c = m > 0 ? 1.0-(z-x) : x-(z-1.0) # 1+x = z+c exactly
        y = reinterpret(Float64,(zu & 0x000f_ffff_ffff_ffff) | 0x3ff0_0000_0000_0000)

        mf = Float64(m)
        F = (y + 0x1p45) - 0x1p45 # 0x1p-7*round(0x1p7*y)
        f = (y - F) + c*s #2^m(F+f) = 1+x = z+c
        jp = unsafe_trunc(Int,0x1p7*F)-127

        log_proc1(y,mf,F,f,jp)
    elseif x == -1.0
        -Inf
    elseif isnan(x)
        NaN
    else
        throw(DomainError())
    end
end

function log1p_tang(x::Float32)
    if x > -1f0
        x == Inf32 && return x
        if -0x1p-53 < x < 0x1p-53
            return x # Inexact

        # Step 2
        elseif Float32(-0x1.f05406p-5) < x < Float32(0x1.082b58p-4)
            return log_proc2(x)
        end

        # Step 3
        z = 1f0 + x
        zu = reinterpret(UInt32,z)
        s = reinterpret(Float32,0x7f000000 - (zu & 0xff80_0000)) # 2^-m
        m = int(zu >> 23) & 0x00ff - 127 # z cannot be subnormal
        c = m > 0 ? 1f0-(z-x) : x-(z-1f0) # 1+x = z+c
        y = reinterpret(Float32,(zu & 0x007f_ffff) | 0x3f80_0000)

        mf = Float32(m)
        F = (y + Float32(0x1p16)) - Float32(0x1p16) # 0x1p-7*round(0x1p7*y)
        f = (y - F) + s*c #2^m(F+f) = 1+x = z+c
        jp = unsafe_trunc(Int,Float32(0x1p7)*F)-127

        log_proc1(y,mf,F,f,jp)
    elseif x == -1f0
        -Inf32
    elseif isnan(x)
        NaN32
    else
        throw(DomainError())
    end
end

@vectorize_1arg Real log_tang
@vectorize_1arg Real log1p_tang


using DataFrames, RData, LinearAlgebra, Statistics

years = 2000:2014 # vector with years covered by WIOD Rev. 2016
N = 41 # number of countries
S = 35 # number of industries
dir = "C:/Users/u0148308/Desktop/raw/" # location of raw data

# --------------- Import data ---------------------------------------------------------------------------------------------------------------------------------

function import_excel(dir::String, year::Integer, N::Integer, S::Integer)

    # load data for specified year and transform into a DataFrame (automatically)
    path = dir * "WIOT" * string(year)[end-1:end] * "_ROW_Apr12.xlsx"
    df = DataFrame(XLSX.readxlsx(path)["WIOT_$year"][:], :auto)
    df = df[7:end,5:end]

    # Intermediate demand
    Z = df[1:N*S, 1:N*S] # NS×NS
    Z = Matrix(convert.(Float64, Z))

    # Final demand
    F = df[1:N*S, N*S+1:end-1] # NS×5*N

    # Subset final demand for inventory adjustments
    inventory_columns = 5:5:5*N
    IV = F[:, inventory_columns] # NS×N
    F = F[:, Not(inventory_columns)] # NS×4*N

    IV = Matrix(convert.(Float64, IV)) # NS×N
    F = Matrix(convert.(Float64, F)) # NS×4*N

    return Z, F, IV
end

function antras_trans1(Z::Matrix, F::Matrix, IV::Matrix, N::Integer, S::Integer)

    XX = Z
    FF = F
    NN = IV

    NN = [sum(NN[i,:]) for i in 1:N*S] # NS×1

    YY = [sum(XX[i,:]) + sum(FF[i,:]) + NN[i] for i in 1:N*S] # NS×1
    YY = ifelse.(abs.(YY) .< 1e-6, 0.0, YY)

    adj = YY.-NN
    adj = ifelse.(adj .< 0.0, 0.0, adj)

    XXi = XX .* (repeat(YY,1,N*S) ./ repeat(adj,1,N*S)) # NS×NS, inventory adjustment
    XXi = ifelse.(isnan.(XXi), 0.0, XXi)
    XXi = ifelse.(isinf.(XXi), 0.0, XXi)

    FFi = FF .* (repeat(YY,1,N*4) ./ repeat(adj,1,N*4)) # NS×N*4, inventory adjustment
    FFi = ifelse.(isnan.(FFi), 0.0, FFi)
    FFi = ifelse.(isinf.(FFi), 0.0, FFi)
    FFdi = [sum(FFi[i,j:j+3]) for i in 1:N*S, j in 1:4:N*4] # NS×N

    YYi = [sum(XXi[i,:]) + sum(FFdi[i,:]) for i in 1:N*S] # NS×1

    VVi = [YYi[i] - sum(XXi[:,i]) for i in 1:N*S] # NS×1

    return XXi, FFdi, YYi, VVi
end

function antras_trans2(XXi::Matrix, FFdi::Matrix, YYi::Vector, VVi::Vector)

    gammas = zeros(S+1,N*S)

    for n in 1:N
    for s in 1:S
        inputlist = XXi[:,(n-1)*S+s]
        inputlist = reshape(inputlist, S, N)
        inputlist = [sum(inputlist[i,:]) for i in 1:S]
        inputlist = [inputlist; VVi[(n-1)*S+s]]
        inputratio = inputlist ./ sum(inputlist)
        gammas[:,(n-1)*S+s] = inputratio
    end
    end

    VAcoeff = gammas[S+1,:]
    VAcoeff = ifelse.(isnan.(VAcoeff), 1.0, VAcoeff)
    VAcoeff = reshape(VAcoeff, S, N)

    gammas = gammas[1:S, 1:N*S]
    gammas = ifelse.(isnan.(gammas), 0.0, gammas)

    alphas = zeros(S, N)
    CtyCons = [sum(FFdi[:,i]) for i in 1:N]

    for j in 1:N
        for i in 1:S
            alphas[i, j] = sum(FFdi[i:S:(N-1)*S+i, j]) / CtyCons[j]
        end
    end

    VAn = [sum(VVi[i:i+S-1]) for i in 1:S:N*S] # N×1

    Piagg = [sum(XXi[s:S:(N-1)*S+s,j]) for s in 1:S, j in 1:N*S]
    PiFagg = [sum(FFdi[s:S:(N-1)*S+s,j]) for s in 1:S, j in 1:N]

    Pi = XXi ./ repeat(Piagg, N)
    Pi = ifelse.(isnan.(Pi), 0.0, Pi)

    PiF = FFdi ./ repeat(PiFagg, N)
    PiF = ifelse.(isnan.(PiF), 0.0, PiF)

    M = zeros(N)
    E = zeros(N)

    for j in 1:N
    for s in 1:S
        for i in 1:N
            iiss = (i-1)*S+s
            jjss = (j-1)*S+s
            for r in 1:S
                jjrr = (j-1)*S+r
                iirr = (i-1)*S+r
                M[j] = M[j] + Pi[iiss,jjrr]*gammas[s,jjrr]*YYi[jjrr]
                E[j] = E[j] + Pi[jjss,iirr]*gammas[s,iirr]*YYi[iirr]
            end
            M[j] = M[j] + PiF[iiss,j]*alphas[s,j]*CtyCons[j]
            E[j] = E[j] + PiF[jjss,i]*alphas[s,i]*CtyCons[i]
        end
    end
    end

    Surplusn = E .- M

    return gammas, alphas, Surplusn, VAcoeff, VAn, Pi, PiF, CtyCons
end

Z, F, IV = import_excel(dir, 1995, N, S)
##################
XXi, FFdi, YYi, VVi = antras_trans1(Z, F, IV, N, S)
gammas, alphas, Surplusn, VAcoeff, VAn, Pi, PiF, CtyCons = antras_trans2(XXi, FFdi, YYi, VVi)

Z, F, Y, VA = create_matrices(Z, F, IV, N, S)
π_Z, π_F = create_trade_shares(Z, F, N, S)
γ, α, VA_coeff, TB_ctry, VA_ctry, F_ctry = create_expenditure_shares(Z, F, Y, VA, π_Z, π_F, N, S)
#####################

alphas == α
gammas == γ
PiF == π_F
Pi == π_Z
F == FFdi
F_ctry == CtyCons

VA_ctry.- F_ctry
Surplusn
TB_ctry
Surplusn .- TB_ctry

[sum(α[:,j]) for j in 1:N]
[sum(alphas[:,j]) for j in 1:N]

count(round.(gammas,digits=0) .!= round.(γ,digits=0))

count(abs.(alphas .- α) .> 0.0001)

γ2 = zeros(S+1,N*S)
for n in 1:N
    for s in 1:S
        inputlist = Z[:,(n-1)*S+s]
        inputlist = reshape(inputlist, S, N)
        inputlist = [sum(inputlist[i,:]) for i in 1:S]
        inputlist = [inputlist; VA[(n-1)*S+s]]
        inputratio = inputlist ./ sum(inputlist)
        γ2[:,(n-1)*S+s] = inputratio
    end
end

VA2 = γ2[S+1,:]
VA2 = ifelse.(isnan.(VA2), 1.0, VA2)
VA2 = reshape(VA2, S, N)

γ2 = γ2[1:S, 1:N*S]
γ2 = ifelse.(isnan.(γ2), 0.0, γ2)

count(abs.(gammas .- γ2) .> 5)

count(γ.>1)

findall(γ.>=1)
findall(γ.<0)

findall(VA_coeff.>1)

findall(π_Z.<0)
findall(π_F.<0)

findall(Z.<0)
findall(VA_ctry.<0)

E_A = zeros(N)
M_A = zeros(N)
for j in 1:N
    for s in 1:S
        for i in 1:N
            iiss = (i-1)*S+s
            jjss = (j-1)*S+s
            for r in 1:S
                jjrr = (j-1)*S+r
                iirr = (i-1)*S+r
                M_A[j] += π_Z[iiss,jjrr]*γ[s,jjrr]*Y[jjrr]
                E_A[j] += π_Z[jjss,iirr]*γ[s,iirr]*Y[iirr]
            end
            M_A[j] += π_F[iiss,j]*α[s,j]*F_ctry[j]
            E_A[j] += π_F[jjss,i]*α[s,i]*F_ctry[i]
        end
    end
end

TB_ctry2 = E_A .- M_A # N×1

F_prime_ctry = copy(F_ctry)
Y_prime = copy(Y_prime)

LHS = zeros(N)
RHS = zeros(N)
for j in 1:N
    for s in 1:S
        for i in 1:N
            iiss = (i-1)*S+s
            jjss = (j-1)*S+s
            for r in 1:S
                jjrr = (j-1)*S+r
                iirr = (i-1)*S+r
                LHS[j] += π_Z[iiss,jjrr]*γ[s,jjrr]*Y[jjrr]
                RHS[j] += π_Z[jjss,iirr]*γ[s,iirr]*Y[iirr]
            end
            LHS[j] += π_F[iiss,j]*α[s,j]*F_ctry[j]
            RHS[j] += π_F[jjss,i]*α[s,i]*F_ctry[i]
        end
    end
end

ETB_ctry = RHS .- LHS

findall(abs.(b).>1)

Z_prime = π_Z .* repeat(γ, N) .* repeat(transpose(Y), N*S) # NS×NS
F_prime = π_F .* repeat(α, N) .* repeat(transpose(F_ctry), N*S) # NS×N

E_prime = [sum(Y[i:i+S-1]) - sum(Z_prime[i:i+S-1,i:i+S-1]) - sum(F_prime[i:i+S-1,ceil(Int, i/S)]) for i in 1:S:N*S] # N×1

M_prime_Z = [sum(Z_prime[:,j:j+S-1]) - sum(Z_prime[j:j+S-1,j:j+S-1]) for j in 1:S:N*S]
M_prime_F = [sum(F_prime[:,ceil(Int, j/S)]) - sum(F_prime[j:j+S-1,ceil(Int, j/S)]) for j in 1:S:N*S]
M_prime = M_prime_Z .+ M_prime_F # N×1

ETB_ctry_prime = E_prime .- M_prime

findall(abs.(Z.-Z_prime).>0.1)
findall(abs.(ETB_ctry.-ETB_ctry_prime).>1)

TB_ctry

ETB = ETB_ctry .- TB_ctry

b = ETB./VA_ctry

δ = sign.(b) .* abs.(vfactor.*b)
w_hat = ones(N)
w_hat = w_hat .* (1.0 .+ δ ./ w_hat)

a = TB_ctry./VA_ctry
δ = sign.(a) .* abs.(vfactor.*a)
w_hat = ones(N)
w_hat = w_hat .* (1.0 .+ δ ./ w_hat)

findall(abs.(a).>1)

findall(TB_ctry .> VA_ctry)


using NeuralNetworks

net = Serial(Float64, (28 * 28, 1), (10, 1))
add_layer!(net, Linear(Float64, 1, 28 * 28, 300))
add_layer!(net, Sigmoid(Float64, (300, 1)))
add_layer!(net, Linear(Float64, 1, 300, 10))
add_layer!(net, Sigmoid(Float64, (10, 1)))
validate(net)

initWeights!(net)

const num_iters = 1000
const num_images = 60000

using MNIST

target = zeros(Float64, 10, 1)
criterion = MSECriterion(Float64, 1, 10)

sgd_state = ObjectIdDict()
function sgd(x, dfdx)
    local lr = 0.1
    local mom = 0.9
    local cur_dfdx
    if haskey(sgd_state, x)
        cur_dfdx = sgd_state[x]
    else
        cur_dfdx = zeros(x)
    end
    cur_dfdx[:] = mom * cur_dfdx + (1.0 - mom) * dfdx
    x[:] = x - lr * cur_dfdx
end

for iter = 1:num_iters
    local loss
    loss = 0.0
    for img = 1:num_images
        local image_data = reshape(MNIST.trainfeatures(img), 28 * 28, 1)
        local label = Int(MNIST.trainlabel(img)) + 1
        target[label,1] = 1.0

        local out = forward!(net, image_data)
        local l = forward!(criterion, out, target)[1]
        loss += l
        local gi = backward!(criterion, out, target)
        zeroGradParameters!(net)
        backward!(net, image_data, gi)

        target[label,1] = 0.0

        for (x, dfdx) in getParameters(net)
            sgd(x, dfdx)
        end
        println("  ", l)
    end
    loss /= num_images
    println(iter, " ", loss)
end


using LinearAlgebra
using NDTensors
using TBLIS

let
  d = 25

  nthreads = 4

  A = randomTensor(d, d, d, d)
  B = randomTensor(d, d, d, d)
  C_blas = randomTensor(d, d, d, d)

  labelsA = (1, -1, 2, -2)
  labelsB = (-2, 4, 3, -1)
  labelsC = (1, 2, 3, 4)

  println("Contracting (d x d x d x d) tensors A * B -> C, d = ", d)
  @show labelsA
  @show labelsB
  @show labelsC

  #
  # Use BLAS
  #

  NDTensors.disable_tblis()
  BLAS.set_num_threads(nthreads)

  println()
  println("Using BLAS with $nthreads threads")

  time_blas = @belapsed NDTensors.contract!($C_blas, $labelsC, $A, $labelsA, $B, $labelsB) samples =
    100

  println()
  println("Time (BLAS) = ", time_blas, " seconds")

  #
  # Use TBLIS
  #

  NDTensors.enable_tblis()
  TBLIS.set_num_threads(nthreads)

  println()
  println("Using TBLIS with $(TBLIS.get_num_threads()) threads")

  C_tblis = randomTensor(d, d, d, d)

  time_tblis = @belapsed NDTensors.contract!($C_tblis, $labelsC, $A, $labelsA, $B, $labelsB) samples =
    100

  println()
  println("Time (TBLIS) = ", time_tblis, " seconds")

  println()
  @show C_blas ≈ C_tblis

  println()
  println("Time (TBLIS) / Time (BLAS) = ", time_tblis / time_blas)
end


# irflp.m
# IRF Estimates and stdev of IRF by Local Projections
# Luca Brugnolini 2014
# Based on Kilian

function irflp(z, paic, H, A0inv, COVsig)

# Input : data z, lag order paic (Hx1 vector), horizon for IRF H,
#         A0inv- relationship btw sf & rf shocks estimated from VAR
#         COVsig- variance due to A0inv, estimated from VAR
# Output: IRF, Std error of IRF, reduced form IRF

(T,K) = size(z);

# IRF estimate and STdev of IRF for Horizon 0
A1 = vec(A0inv');                                # IRF for horizon 0
A2 = vec(eye(K));                                # Reduced form IRF
COVA1 = zeros(K^2,1);                            # Std(IRF0)=0

# IRF estimate and Stdev of IRF for Horizon 1~H
 for h = 1:H
     ph = paic[h];                     # lag-order for horizon h
     ph = itrunc(ph);
     ys = z[ph+h:T,:];

    # RHS variables: lags
     yt = z[ph:T-h,:];                 # RHS variable of interest: y(t)
     x = ones(T-ph-h+1,1);             # constant term

        for i = 2:ph
            x = [x z[ph+1-i:T-i-h+1,:]];  # other RHS variables: y(t-1)~y(t-p+1)

        end

     Mx = eye(size(x,1))-x/(x'*x)*x';         # annhilation matrix
     B1 = (yt'*Mx*yt)\(yt'*Mx*ys);            # IRF by Local Projection
     A1 = [A1 vec(A0inv'*B1)];                # structural IRF
     A2 = [A2 vec(B1)];                       # reduced form IRF

     # Newey-West HAC Estimator of Stdev (truncation point is h)
     u = Mx*ys - Mx*yt*B1;                                           # residual from LPs
     meanu = [mean(u[:,1]) mean(u[:,2]) mean(u[:,3]) mean(u[:,4])];
     u0 = u-kron(ones(size(u,1),1),meanu);
     R0 = (u0'*u0)./(size(u,1));         # here was length(u)                            # R(0): smpl cov(u)
     SIGu = R0;
     M = h;                                                          # truncation point

        for j = 1:M
            Rj = (u0[1:end-j,:]'*u0[j+1:end,:])./(size(u,1));   # R(j)
            SIGu = SIGu + (1-j/(M+1))*(Rj+Rj');                 # varcov(u)

        end

     invytMxyt  = inv(yt'*Mx*yt);
     SIGB1 = kron(SIGu, invytMxyt);                               # var(vec(B1))
     SIGA1 = kron(eye(K), A0inv')*SIGB1*kron(eye(K), A0inv')';     # var(vec(A0inv*B1))
     COVA1 = [COVA1 reshape(diag(SIGA1), K^2, 1)];

end

    stdA1 = sqrt(COVA1+COVsig);                                       # STD2 is variance due to

     return A1, stdA1, A2

end


#===============================================================================
 Robust NonMonotone Line search
 From DF-Sane paper https://www.ams.org/journals/mcom/2006-75-255/S0025-5718-06-01840-0/S0025-5718-06-01840-0.pdf
 This line search doesn't fit the rest of the API so I need to come up with an
 interface for nonmonotone line searches (it has to accept an fbar (or Q in Hager
 Zhang's notion) and a forcing term)
===============================================================================#
struct RNMS end

function find_steplength(rnms::RNMS, φ, fbar, ηk::T, τmin, τmax) where T
  φ0 = φ.φ0
  α₊ = T(1)
  α₋ = -α₊
  γ = T(1)/4
  for k = 1:100
    φα₊ = φ(α₊)
    if φα₊ ≤ fbar + ηk - γ*α₊^2*φ0
      return α₊, φα₊
    end
    φα₋ = φ(α₋)
    if φα₋ ≤ fbar + ηk - γ*α₋^2*φ0
      return α₋, φα₋
    end
    αt = α₊^2*φ0/(φα₊ + (2*α₊-1)*φ0)
    if αt < τmin*α₊
      α₊ = τmin*α₊
    elseif αt > τmax*α₊
      α₊ = τmax*α₊
    else
      α₊ = αt
    end

    # This needs to be safe guarded!
    αt = α₋^2*φ0/(φα₋ + (2*α₋-1)*φ0)
    if αt < τmin*α₋
      α₋ = τmin*α₋
    elseif αt > τmax*α₋
      α₋ = τmax*α₋
    else
      α₋ = αt
    end
  end
  throw(ErrorException("Line search failed."))
end

function safeguard_σ(σ::T, σmin, σmax, F) where T
  if abs(σ) < σmin || abs(σ) > σmax
    normF = norm(F)
    if normF > T(1)
      σ = T(1)
    elseif T(1)/10^5 ≤ normF ≤ T(1)
      σ = inv(normF)
    elseif normF < T(1)/10^5
      σ = T(10)^5
    end
  end
  return σ
end

struct DFSANE end
function nlsolve(F::OnceDiffed, x0, scheme::DFSANE)
  T = eltype(x0)

  σmin, σmax = 1e-10, 1e10
  τmin, τmax = T(1)/10, T(1)/2

  nexp = 2

  fvals = @SVector(T[])

  x = copy(x0)
  Fx = F(x)
  y = copy(Fx)
  normFx0 = norm(Fx)
  fx = normFx0^nexp
  fvals = push(fvals, fx)
  σ = T(1)
  σ = safeguard_σ(σ, σmin, σmax, Fx)
  for j = 1:10^5
    fbar = maximum(x->fvals[x], max(1, length(fvals)-5+1):length(fvals))
    d = -σ*Fx
    ηk = normFx0/(1+j)^2
    φ = LineObjective(OnceDiffed((_n,x)->norm(F(x))^nexp), x, d, fx, 0*fx)
    α, φα = find_steplength(RNMS(), φ, fbar, ηk, τmin, τmax)
    # LineObjective(x) should return z as well!
    s = α*d
    x .= x.+s
    y .= Fx
    Fx = F(x)
    y .-= Fx
    fx = norm(Fx)^nexp
    fvals = push(fvals, fx)
    σ = norm(s)^2/dot(s, y)
    σ = safeguard_σ(σ, σmin, σmax, Fx)
    if sqrt(φα) < 1e-8
      break
    end
  end
  x, Fx
end


# Generate a random sparse matrix using SuiteSparseGraphBLAS Julia wrapper
# from https://github.com/abhinavmehndiratta/SuiteSparseGraphBLAS.jl

using GraphBLASInterface, SuiteSparseGraphBLAS

"""
    GraphBLAS_sparse_square_matrix!(A::GrB_Matrix{T}, dim::Int) where T

Build a random sparse square matrix of dimension `dim`x`dim` using GraphBLAS.
"""
function GraphBLAS_sparse_square_matrix!(
    A::GrB_Matrix{T},
    dim::Int) where T

    # Build A with at most sqrt(N) elements, N being the number of nonzero matrix elements
    ntuples = dim

    for k in 1:ntuples

        i = ZeroBasedIndex(rand(0:dim-1))
        j = ZeroBasedIndex(rand(0:dim-1))
        x = rand(T)

        GrB_Matrix_setElement(A, x, i, j)
    end

    # Force pending operations to complete and return
    GrB_wait()

    return (GrB_SUCCESS)
end


# # [A classification task when labels are known - determining the country of origin of cars given the cars characteristics](@id classification_tutorial)

# In this exercise we have some car technical characteristics (mpg, horsepower,weight, model year...) and the country of origin and we would like to create a model such that the country of origin can be accurately predicted given the technical characteristics.
# As the information to predict is a multi-class one, this is a _[classification]_(https://en.wikipedia.org/wiki/Statistical_classification) task.
# It is a challenging exercise due to the simultaneous presence of three factors: (1) presence of missing data; (2) unbalanced data - 254 out of 406 cars are US made; (3) small dataset.

#
# Data origin:
# - dataset description: [https://archive.ics.uci.edu/ml/datasets/auto+mpg](https://archive.ics.uci.edu/ml/datasets/auto+mpg)
#src Also useful: https://www.rpubs.com/dksmith01/cars
# - data source we use here: [https://archive.ics.uci.edu/ml/machine-learning-databases/auto-mpg/auto-mpg.data](https://archive.ics.uci.edu/ml/machine-learning-databases/auto-mpg/auto-mpg.data-original)

# Field description:

# 1. mpg:           _continuous_
# 2. cylinders:     _multi-valued discrete_
# 3. displacement:  _continuous_
# 4. horsepower:    _continuous_
# 5. weight:        _continuous_
# 6. acceleration:  _continuous_
# 7. model year:    _multi-valued discrete_
# 8. origin:        _multi-valued discrete_
# 9. car name:      _string (unique for each instance)_ - not used here

cd(@__DIR__)
using Pkg
Pkg.activate(".")
Pkg.add(["HTTP","CSV","DataFrames","Pipe","BenchmarkTools", "BetaML"])
Pkg.instantiate()

# ## Library and data loading

# We load a buch of packages that we'll use during this tutorial..
using Random, HTTP, CSV, DataFrames, BenchmarkTools,BetaML
#import DecisionTree, Flux
import Pipe: @pipe

# To load the data from the internet our workflow is
# (1) Retrieve the data --> (2) Clean it --> (3) Load it --> (4) Output it as a DataFrame.

# For step (1) we use `HTTP.get()`, for step (2) we use `replace!`, for steps (3) and (4) we uses the `CSV` package, and we use the "pip" `|>` operator to chain these operations:

urlDataOriginal = "https://archive.ics.uci.edu/ml/machine-learning-databases/auto-mpg/auto-mpg.data-original"
data = @pipe HTTP.get(urlDataOriginal).body                                                |>
             replace!(_, UInt8('\t') => UInt8(' '))                                        |>
             CSV.File(_, delim=' ', missingstring="NA", ignorerepeated=true, header=false) |>
             DataFrame;

# This results in a table where the rows are the observations (the various cars) and the column the fields. All BetaML models expect this layout.

# As the dataset is ordered, we randomly shuffle the data. Note that we pass to shuffle `copy(FIXEDRNG)` as the random nuber generator in order to obtain reproducible output ( [`FIXEDRNG`](@ref BetaML.Utils.FIXEDRNG) is nothing else than an istance of `StableRNG(123)` defined in the [`BetaML.Utils`](@ref utils_module) sub-module, but you can choose of course your own "fixed" RNG). See the [Dealing with stochasticity](@ref dealing_with_stochasticity) section in the [Getting started](@ref getting_started) tutorial for details.
data[shuffle(copy(FIXEDRNG),axes(data, 1)), :]
describe(data)

# Columns 1 to 7 contain  characteristics of the car, while column 8 encodes the country or origin ("1" -> US, "2" -> EU, "3" -> Japan). That's the variable we want to be able to predict.

# Columns 9 contains the car name, but we are not going to use this information in this tutorial.
# Note also that some fields have missing data.

# Our first step is hence to divide the dataset in features (the x) and the labels (the y) we want to predict. The `x` is then a Julia standard `Matrix` of 406 rows by 7 columns and the `y` is a vector of the 406 observations:
x     = Matrix{Union{Missing,Float64}}(data[:,1:7]);
y     = Vector{Int64}(data[:,8]);

# Some algorithms that we will use today don't work with missing data, so we need to _impute_ them. We use the [`predictMissing`](@ref) function provided by the [`BetaML.Clustering`](@ref clustering_module) sub-module. Internally the function uses a Gaussian Mixture Model to assign to the missing walue of a given record an average of the values of the non-missing records weighted for how close they are to our specific record.
# Note that the same function (`predictMissing`) can be used for Collaborative Filtering / recomendation systems. Using GMM has the advantage over traditional algorithms as k-nearest neighbors (KNN) that GMM can "detect" the hidden structure of the observed data, where some observation can be similar to a certain pool of other observvations for a certain characteristic, but similar to an other pool of observations for other characteristics.

xFull = predictMissing(x,rng=copy(FIXEDRNG)).X̂;

# Further, some models don't work with categorical data as such, so we need to represent our `y` as a matrix with a separate column for each possible categorical value (the so called "one-hot" representation).
# For example, within a three classes field, the individual value `2` (or `"Europe"` for what it matters) would be represented as the vector `[0 1 0]`, while `3` (or `"Japan"`) would become the vector `[0 0 1]`.
# To encode as one-hot we use the function [`oneHotEncoder`](@ref) in [`BetaML.Utils`](@ref utils_module)
y_oh  = oneHotEncoder(y);

# In supervised machine learning it is good practice to partition the available data in a _training_, _validation_, and _test_ subsets, where the first one is used to train the ML algorithm, the second one to train any eventual "hyper-parameters" of the algorithm and the _test_ subset is finally used to evaluate the quality of the algorithm.
# Here, for brevity, we use only the _train_ and the _test_ subsets, implicitly assuming we already know the best hyper-parameters. Please refer to the [regression tutorial](@ref regression_tutorial) for examples of how to use the validation subset to train the hyper-parameters, or even better the [clustering tutorial](@ref clustering_tutorial) for an example of using the [`crossValidation`](@ref) function.

# We use then the [`partition`](@ref) function in [BetaML.Utils](@ref utils_module), where we can specify the different data to partition (each matrix or vector to partition must have the same number of observations) and the shares of observation that we want in each subset. Here we keep 80% of observations for training (`xtrain`, `xTrainFull` and `ytrain`) and we use 20% of them for testing (`xtest`, `xTestFull` and `ytest`):
((xtrain,xtest),(xtrainFull,xtestFull),(ytrain,ytest),(ytrain_oh,ytest_oh)) = partition([x,xFull,y,y_oh],[0.8,1-0.8],rng=copy(FIXEDRNG));

# ## Random Forests

# We are now ready to use our first model, the Random Forests (in the [`BetaML.Trees`](@ref trees_module) sub-module). Random Forests build a "forest" of decision trees models and then average their predictions in order to make an overall prediction out of a feature vector.

# To "build" the forest model (i.e. to "train" it) we need to give the model the training feature matrix and the associated "true" training labels, and we need to specify the number of trees to employ (this is an example of hyper-parameters). Here we use 30 individual decision trees.

# As the labels are encoded using integers,  we need also to specify the parameter `forceClassification=true`, otherwise the model would undergo a _regression_ job instead.

myForest       = buildForest(xtrain,ytrain,30, rng=copy(FIXEDRNG),forceClassification=true);

# To obtain the predicted values, we can simply use the function [`BetaML.Trees.predict`](@ref)
#src [`predict`](@ref BetaML.Trees.predict)  [`predict`](@ref forest_prediction)
# with our `myForest` model and either the training or the testing data.
ŷtrain,ŷtest   = predict.(Ref(myForest), [xtrain,xtest],rng=copy(FIXEDRNG));
# Finally we can measure the _accuracy_ of our predictions with the [`accuracy`](@ref) function, with the sidenote that we need first to "parse" the ŷs as forcing the classification job transformed automatically them to strings (they originally were integers):
trainAccuracy,testAccuracy  = accuracy.([parse.(Int64,mode(ŷtrain,rng=copy(FIXEDRNG))),parse.(Int64,mode(ŷtest,rng=copy(FIXEDRNG)))],[ytrain,ytest])
#src (0.9969230769230769,0.8024691358024691)

@test testAccuracy > 0.8 #src

# The predictions are quite good, for the training set the algoritm predicted almost all cars' origins correctly, while for the testing set (i.e. those records that has **not** been used to train the algorithm), the correct prediction level is still quite high, at 80%

# While accuracy can sometimes suffice, we may often want to better understand which categories our model has trouble to predict correctly.
# We can investigate the output of a multi-class classifier more in-deep with a [`ConfusionMatrix`](@ref) where the true values (`y`) are given in rows and the predicted ones (`ŷ`) in columns, together to some per-class metrics like the _precision_ (true class _i_ over predicted in class _i_), the _recall_ (predicted class _i_ over the true class _i_) and others.

# We fist build the [`ConfusionMatrix`](@ref BetaML.Utils.ConfusionMatrix) object between `ŷ` and `y` and then we print it (we do it here for the test subset):

cm = ConfusionMatrix(parse.(Int64,mode(ŷtest,rng=copy(FIXEDRNG))),ytest,classes=[1,2,3],labels=["US","EU","Japan"])
print(cm;what="all")

# Because the printing of the confusion matrix employs `display`, the confusion matrix is printed on the script building this documentation pages rather than on the documentation pages themselves. Here it is:

#
# ```text
# 4×4 Matrix{Any}:
#  "Labels"    "US"    "EU"   "Japan"
#  "US"      43       3      3
#  "EU"       2      13      1
#  "Japan"    5       2      9
# 4×4 Matrix{Any}:
#  "Labels"   "US"      "EU"       "Japan"
#  "US"      0.877551  0.0612245  0.0612245
#  "EU"      0.125     0.8125     0.0625
#  "Japan"   0.3125    0.125      0.5625
# ```

# From the report we can see that Japanese cars have more trouble in being correctly classified, and in particular many Japanease cars are classified as US ones. This is likely a result of the class imbalance of the data set, and could be solved by balancing the dataset with various sampling tecniques before training the model.

# When we benchmark the resourse used (time and memory) we find that Random Forests remain pretty fast, expecially when we compare them with neural networks (see later)
@btime buildForest(xtrain,ytrain,30, rng=copy(FIXEDRNG),forceClassification=true);
#src   134.096 ms (781027 allocations: 196.30 MiB)

# ### Comparision with DecisionTree.jl

# DecisionTrees.jl random forests are similar in usage: we first "build" (train) the forest and we then make predictions out of the trained model.
# The main difference is that the model requires data with nonmissing values, so we are going to use the `xtrainFull` and `xtestFull` feature labels we created earlier:
## We train the model...
model = DecisionTree.build_forest(ytrain, xtrainFull,-1,30,rng=123)
## ..and we generate predictions and measure their error
(ŷtrain,ŷtest) = DecisionTree.apply_forest.([model],[xtrainFull,xtestFull]);
(trainAccuracy,testAccuracy) = accuracy.([ŷtrain,ŷtest],[ytrain,ytest])
#src (0.9969230769230769, 0.7530864197530864)

#src cm = ConfusionMatrix(ŷtest,ytest,classes=[1,2,3],labels=["US","EU","Japan"])
#src println(cm)
@test testAccuracy > 0.75 #src

# While the accuracy on the training set is exactly the same as for `BetaML` random forets, `DecisionTree.jl` random forests are slighly less accurate in the testing sample.
# Where however `DecisionTrees.jl` excell is in the efficiency: they are extremelly fast and memory thrifty, even if to this benchmark we should add the resources needed to impute the missing values.

# Also, one of the reasons DecisionTrees are such efficient is that internally they sort the data to avoid repeated comparision, but in this way they work only with features that are sortable, while BetaML random forests accept virtually any kind of input without the need of adapt it.
@btime  DecisionTree.build_forest(ytrain, xtrainFull,-1,30,rng=123);
#src 1.431 ms (10875 allocations: 1.52 MiB)

# ### Neural network

# Neural networks (NN) can be very powerfull, but have two "inconvenients" compared with random forests: first, are a bit "picky". We need to do a bit of work to provide data in specific format. Note that this is _not_ feature engineering. One of the advantages on neural network is that for the most this is not needed for neural networks. However we still need to "clean" the data. One issue is that NN don't like missing data. So we need to provide them with the feature matrix "clean" of missing data. Secondly, they work only with numerical data. So we need to use the one-hot encoding we saw earlier.
# Further, they work best if the features are scaled such that each feature has mean zero and standard deviation 1. We can achieve it with the function [`scale`](@ref) or, as in this case, [`getScaleFactors`](@ref).

xScaleFactors   = getScaleFactors(xtrainFull)
D               = size(xtrainFull,2)
classes         = unique(y)
nCl             = length(classes)

# The second "inconvenient" of NN is that, while not requiring feature engineering, they stil lneed a bit of practice on the way to build the network. It's not as simple as `train(model,x,y)`. We need here to specify how we want our layers, _chain_ the layers together and then decide a _loss_ overall function. Only when we done these steps, we have the model ready for training.
# Here we define 2 [`DenseLayer`](@ref) where, for each of them, we specify the number of neurons in input (the first layer being equal to the dimensions of the data), the output layer (for a classification task, the last layer output size beying equal to the number of classes) and an _activation function_ for each layer (default the `identity` function).
ls   = 50
l1   = DenseLayer(D,ls,f=relu,rng=copy(FIXEDRNG))
l2   = DenseLayer(ls,nCl,f=relu,rng=copy(FIXEDRNG))

# For a classification the last layer is a [`VectorFunctionLayer`](@ref) that has no learnable parameters but whose activation function is applied to the ensemble of the neurons, rather than individually on each neuron. In particular, for classification we pass the [`BetaML.Utils.softmax`](@ref) function whose output has the same size as the input (and the number of classes to predict), but we can use the `VectorFunctionLayer` with any function, including the [`pool1d`](@ref) function to create a "pooling" layer (using maximum, mean or whatever other subfunction we pass to `pool1d`)

l3   = VectorFunctionLayer(nCl,f=softmax) ## Add a (parameterless) layer whose activation function (softMax in this case) is defined to all its nodes at once

# Finally we _chain_ the layers and assign a loss function with [`buildNetwork`](@ref):
mynn = buildNetwork([l1,l2,l3],squaredCost,name="Multinomial logistic regression Model Cars") ## Build the NN and use the squared cost (aka MSE) as error function (crossEntropy could also be used)

# Now we can train our network using the function [`train!`](@ref). It has many options, have a look at the documentation for all the possible arguments.
# Note that we train the network based on the scaled feature matrix.
res  = train!(mynn,scale(xtrainFull,xScaleFactors),ytrain_oh,epochs=500,batchSize=8,optAlg=ADAM(),rng=copy(FIXEDRNG)) ## Use optAlg=SGD() to use Stochastic Gradient Descent instead

# Once trained, we can predict the label. As the trained was based on the scaled feature matrix, so must be for the predictions
(ŷtrain,ŷtest)  = predict.(Ref(mynn),[scale(xtrainFull,xScaleFactors),scale(xtestFull,xScaleFactors)])
trainAccuracy   = accuracy(ŷtrain,ytrain,rng=copy(FIXEDRNG))

#-

testAccuracy    = accuracy(ŷtest,ytest,rng=copy(FIXEDRNG))
#src (0.8953846153846153, 0.7654320987654321

@test testAccuracy > 0.76 #src

#src accuracy(mode(ŷtest,rng=copy(FIXEDRNG)),ytest)

cm = ConfusionMatrix(ŷtest,ytest,classes=[1,2,3],labels=["US","EU","Japan"],rng=copy(FIXEDRNG))
print(cm)

# 4×4 Matrix{Any}:
#  "Labels"    "US"    "EU"   "Japan"
#  "US"      44       0      5
#  "EU"       3      10      3
#  "Japan"    6       2      8
# 4×4 Matrix{Any}:
#  "Labels"   "US"      "EU"   "Japan"
#  "US"      0.897959  0.0    0.102041
#  "EU"      0.1875    0.625  0.1875
#  "Japan"   0.375     0.125  0.5

# We see a bit the limits of neural networks in this example. While NN can be extremelly performant in many domains, they also require lot of data and computational power, expecially considering the many possible hyper-parameters and hence its large space in the hyper-parameter tuning.
# In this example we arrive short to the performance of random forests, yet with a significant numberof neurons.

@btime train!(mynn,scale(xtrainFull),ytrain_oh,epochs=300,batchSize=8,rng=copy(FIXEDRNG),verbosity=NONE);
#src  11.841 s (62860672 allocations: 4.21 GiB)

# ### Comparisons with Flux

# In Flux the input must be in the form (fields, observations), so we transpose our original matrices
xtrainT, ytrain_ohT = transpose.([scale(xtrainFull,xScaleFactors), ytrain_oh])
xtestT, ytest_ohT   = transpose.([scale(xtestFull,xScaleFactors), ytest_oh])

# We define the Flux neural network model in a similar way than BetaML and load it with data, we train it, predict and measure the accuracies on the training and the test sets:

#src function poolForFlux(x,wsize=5)
#src     hcat([pool1d(x[:,i],wsize;f=maximum) for i in 1:size(x,2)]...)
#src end
Random.seed!(123)

l1         = Flux.Dense(D,ls,Flux.relu)
#l2         = Flux.Dense(ls,ls,Flux.relu)
l3         = Flux.Dense(ls,nCl,Flux.relu)
Flux_nn    = Flux.Chain(l1,l3)
loss(x, y) = Flux.logitcrossentropy(Flux_nn(x), y)
ps         = Flux.params(Flux_nn)
nndata     = Flux.Data.DataLoader((xtrainT, ytrain_ohT), batchsize=8,shuffle=true)
begin for i in 1:500  Flux.train!(loss, ps, nndata, Flux.ADAM()) end end
ŷtrain     = Flux.onecold(Flux_nn(xtrainT),1:3)
ŷtest      = Flux.onecold(Flux_nn(xtestT),1:3)
trainAccuracy =  accuracy(ŷtrain,ytrain)

#-

testAccuracy  = accuracy(ŷtest,ytest)
#src 0.9384615384615385, 0.7407407407407407
# While the train accuracy is little bit higher that BetaML, the test accuracy remains comparable

@test testAccuracy > 0.74 #src

# However the time is again lower than BetaML, even if here for "just" a factor 2
@btime begin for i in 1:500 Flux.train!(loss, ps, nndata, Flux.ADAM()) end end;
#src  5.665 s (8943640 allocations: 1.07 GiB)

# ## Summary

# This is the summary of the results we had trying to predict the country of origin of the cars, based on their technical characteristics:

# | Model                | Train acc     | Test Acc |  Training time (ms)* | Training mem (MB) *|
# |:-------------------- |:-------------:| --------:| ------------------- | ----------------- |
# | RF                   | 0.9969        | 0.8025   | 134                 | 196               |
# | RF (DecisionTree.jl) | 0.9969        | 0.7531   | 1.43                | 1.5               |
# | NN                   | 0.895         | 0.765    | 11841               | 4311              |
# | NN (Flux.jl)         | 0.938         | 0.741    | 5665                | 1096              |

# * on a Intel Core i5-8350U laptop

# We warn that this table just provides a rought idea of the various algorithms performances. Indeed there is a large amount of stochasticity both in the sampling of the data used for training/testing and in the initial settings of the parameters of the algorithm. For a statistically significant comparision we would have to repeat the analysis with multiple sampling (e.g. by cross-validation, see the [clustering tutorial](@ref clustering_tutorial) for an example) and initial random parameters.

# Neverthless the table above shows that, when we compare BetaML with the algorithm-specific leading packages, we found similar results in terms of accuracy, but often the leading packages are better optimised and run more efficiently (but sometimes at the cost of being less verstatile).
# Also, for this dataset, Random Forests seems to remain marginally more accurate than Neural Network, altought of course this depends on the hyper-parameters and, with a single run of the models, we don't know if this difference is significant.


#
# returns the indexes (i,j) of given the index of the pilled-up columns
#
function setindex2D(ipair,npairs)
  j = ipair%npairs
  if j == 0
    j = 1
  end
  i = round(Int64,(npairs-j)/npairs)+1
  return i, j
end

function setipair2D(n,i,j)
  if j >= i
    ipair = (i-1)*n + j
  else
    ipair = (j-1)*n + i
  end
  ipair = ipair - round(Int64,(i^2+i)/2)
  return ipair
end


abstract type AbstractResponseModel end

"""
    random_betas(gr::GroupedFeatures, αs)

Suppose `gr` consists of ``K`` groups with ``p_1, \\dotsc, p_K`` features each.
Then this returns a random vector of βs of length ``\\sum p_g``,
where for `j` in the `g`-th group
we draw (independent) ``β_j \\sim N(0, α_g^2/p_g)``.
``α_g`` is the `g`-th element of the vectpr `αs`.
"""
function random_betas(gr::GroupedFeatures, αs)
    ps = gr.ps
    βs = zeros(eltype(αs), gr.p)
    for i = 1:gr.num_groups
        βs[group_idx(gr, i)] .= randn(ps[i]) .* sqrt(αs[i]^2 / ps[i])
    end
    βs
end

Base.@kwdef struct RandomLinearResponseModel <: AbstractResponseModel
    αs::Vector{Float64}
    grp::GroupedFeatures
    iid_measure = Normal()
end

function (resp::RandomLinearResponseModel)(X)
    β = random_betas(resp.grp, resp.αs) #todo, allow other noise dbn.
    Xβ = X * β
    Xβ, β
end

Base.@kwdef struct GroupRidgeSimulationSettings{C,R,D}
    groups::GroupedFeatures
    Σ::C
    response_model::R
    response_noise::D = Normal()
    ntest::Int = 10000
    ntrain::Int
    iid_measure = Normal()
end

Base.@kwdef struct GroupRidgeSimulation
    groups::GroupedFeatures
    X::Matrix{Float64}
    Y::Vector{Float64}
    resampling_idx = nothing
    β = nothing
end

function simulate(group_simulation::GroupRidgeSimulationSettings)
    ntrain = group_simulation.ntrain
    ntest = group_simulation.ntest
    ntotal = ntrain + ntest
    @unpack response_model, response_noise = group_simulation

    X = simulate_rotated_design(
        group_simulation.Σ,
        ntotal;
        rotated_measure = group_simulation.iid_measure,
    )

    Xβ, β = response_model(X)
    Y = Xβ .+ rand(response_noise, ntotal)

    resampling_idx = [(1:ntrain, (ntrain+1):ntotal)]

    GroupRidgeSimulation(;
        groups = group_simulation.groups,
        X=X,
        Y=Y,
        resampling_idx = resampling_idx,
        β = β,
    )
end


using SparseArrays

A = [2 3; 1 0; 0 6]

@test @inferred sparsify(A, 0.4, Inf, 200, true) ≈ sparse([2.0 3.0; 1.0 0; 0 6.0])
@test @inferred sparsify(A', 0.4, Inf, 200, true) ≈ sparse([2.0 3.0; 1.0 0; 0 6.0]')
A_id = @inferred sparsity_pattern(A, 0.4, Inf)
A_id_t = @inferred sparsity_pattern(A', 0.4, Inf)
@test A_id ≈ sparse([1 1; 1 0; 0 1])
@test A_id_t ≈ sparse([1 1; 1 0; 0 1]')
@test @inferred bin_sparse_matrix!(A, A_id, 200) ≈ sparse([1 3; 2 0; 0 4])
@test @inferred bin_sparse_matrix!(A', A_id_t, 200) ≈ sparse([1 3 0; 2 0 4])


using Plots, Printf, Statistics
# pyplot()
do_viz = false

function compute_flux!(qx, qy, T, λ, dx, dy, nx, ny)
    # Threads.@threads for iy=2:ny-1
    for iy=2:ny-1
        for ix=1:nx-1
            qx[ix,iy-1] = -λ*(T[ix+1,iy]-T[ix,iy])/dx
        end
    end
    # Threads.@threads for iy=1:ny-1
    for iy=1:ny-1
        for ix=2:nx-1
            qy[ix-1,iy] = -λ*(T[ix,iy+1]-T[ix,iy])/dy
        end
    end
    return
end

function update_T!(T, qx, qy, dt, ρCp, dx, dy, nx, ny)
    # Threads.@threads for iy=2:ny-1
    for iy=2:ny-1
        for ix=2:nx-1
            T[ix,iy] = T[ix,iy] - dt/ρCp*((qx[ix,iy-1]-qx[ix-1,iy-1])/dx + (qy[ix-1,iy]-qy[ix-1,iy-1])/dy)
        end
    end
    return
end

@views function heat_2D()
    @show Threads.nthreads()
    # Physics
    lx, ly = 10.0, 10.0
    λ      = 1.0
    ρCp    = 1.0
    nt     = 200
    # Numerics
    nx, ny = 128, 128
    nout   = 10
    # Derived numerics
    dx, dy = lx/nx, ly/ny
    dt     = min(dx^2,dy^2)/ρCp/λ/4.1
    xc     = LinRange(dx/2, lx-dx/2, nx)
    yc     = LinRange(dy/2, ly-dy/2, ny)
    # Array allocation
    qx     = zeros(nx-1,ny-2)
    qy     = zeros(nx-2,ny-1)
    # Initial condition
    T      = exp.(.-(xc.-lx./2.0).^2 .-(yc.-ly./2.0)'.^2)
    # Time loop
    for it = 1:nt
        if (it==11) global t0 = Base.time() end
        compute_flux!(qx, qy, T, λ, dx, dy, nx, ny)
        update_T!(T, qx, qy, dt, ρCp, dx, dy, nx, ny)
        # Visualise
        if mod(it,nout)==0 && do_viz
            display(heatmap(xc, yc, T', xlabel="lx", ylabel="ly", title="heat diffusion, it=$it", clims=(0.,1.)))
            # sleep(.01)
        end
    end
    time_s = (Base.time()-t0)
    @printf("Time = %1.4e s, T_eff = %1.2f GB/s \n", time_s, round((2/1e9*nx*ny*sizeof(lx))/(time_s/(nt-10)), sigdigits=2))
    return
end

@time heat_2D()


module Data

using DrWatson
using Random, Statistics
using DynamicalSystems, OrdinaryDiffEq, LinearAlgebra
using DelimitedFiles, DrWatson
using StochasticDiffEq

# default constants used if no other is given
const default_N = 10000
const diffeq = (alg = Vern9(), reltol = 1e-12, abstol = 1e-12, maxiters = typemax(Int))
const diffeq_lowacc = (alg = Tsit5(), reltol = 1e-9, abstol = 1e-9, maxiters = typemax(Int))

function kaplanyorke_map(; N = default_N, λ = 0.2, kwargs...)
    f(u, λ, t) = SVector((2u[1]) % 1.0 + 1e-15rand(Random.GLOBAL_RNG), λ*u[2] + cospi(4u[1]))
    ds = DiscreteDynamicalSystem(f, SVector(0.15, 0.2), λ)
    tr = trajectory(ds, N; Ttr = 1000)
    standardize(tr)
end

function roessler_periodic(; N = default_N, Δt = 0.5, kwargs...)
    ds = Systems.roessler([0.1, -0.2, 0.1]; c = 3.0, a = 0.2, b = 0.2)
    tr = trajectory(ds, N*Δt; Δt, Ttr = 1000, diffeq...)
    standardize(tr)
end

function roessler_chaotic(; η=0, correlated=false, N = default_N, Δt = 0.5, kwargs...)
    ds = Systems.roessler([0.1, -0.2, 0.1]; c = 5.7, a = 0.2, b = 0.2)
    tr = trajectory(ds, N*Δt; Δt, Ttr = 1000, diffeq...)
    tr = standardize(tr)
    if η ≠ 0 # add static noise
        if !correlated
            seed = Random.MersenneTwister(5653435)
            for i in 1:length(tr)
                tr[i] += η*randn(seed, SVector{3})
            end
        else correlated
            a = arma_noise(N+1)
            for i in 1:length(tr)
                tr[i] = tr[i] .+ η .* a[i]
            end
        end
    end
    return tr
end

function roessler_embedding(; d = 3, kwargs...)
    x = roessler_chaotic(kwargs...)[:, 1]
    x = standardize(x)
    τ = estimate_delay(x, "mi_min")
    return embed(x, d, τ)
end

function roessler_rounded(; digits = 2, kwargs...)
    X = roessler_chaotic(; kwargs...)
    Z = [round.(x; digits = digits) for x in X]
    return Dataset(Z)
end

function roessler_dynoise(du,u, p, t)
    a, b, c, η = p
    du[1] = -u[2]-u[3]
    du[2] = u[1] + a*u[2]
    du[3] = b + u[3]*(u[1] - c)
end

# SDEProblem wants a specific function for the noise
# (This assumes diagonal noise, with the only noise acting on the
# 2nd variable)
function σ_roessler_dynoise(du, u, p, t)
    a, b, c, η = p
    du[1] = 0
    du[2] = η
    du[3] = 0
end

function roessler_dynamic(; N = default_N, Δt = 0.2, η=0.1, kwargs...)
    # using the SDE solve from DifferentialEquations here
    tspan = (0.0, N*Δt + 100)
    prob_sde_roessler = SDEProblem(
        roessler_dynoise, σ_roessler_dynoise, [0.1, -0.2, 0.1],
        tspan, [0.2, 0.2, 5.7, η]
    )
    saveat = 100.0:Δt:tspan[end]
    # I really don't have much experience with which solver to use
    sol = solve(prob_sde_roessler, SRA(); saveat, maxiters = typemax(Int))
    X = Dataset(hcat(sol.u...)')
    return standardize(X)
end

function roessler_lorenz(; N = default_N, kwargs...)
    tr1 = roessler_periodic(N = N÷2)
    tr2 = lorenz96_chaotic(; D = 4, N = N÷2)[:, 1:3]
    tr2 = trajectory(ds, (N/2)*Δt; Δt, Ttr = 1000, diffeq...)
    append!(standardize(tr1), standardize(tr2))
end

function roessler_sphere(; N = default_N, Δt = 0.5, kwargs...)
    ds = Systems.roessler([0.1, -0.2, 0.1]; c = 3.0, a = 0.2, b = 0.2)
    tr1 = standardize(trajectory(ds, (N/2)*Δt; Δt, Ttr = 1000, diffeq...))
    tr2 = standardize(uniform_sphere(; N = N÷2))
    return append!(tr1, tr2)
end

function roessler_nonstationary_eom(u, p, t)
    @inbounds begin
        a, b, dc = p
        du1 = -u[2]-u[3]
        du2 = u[1] + a*u[2]
        du3 = b + u[3]*(u[1] - u[4])
        du4 = dc
        return SVector{4}(du1, du2, du3, du4)
    end
end

function roessler_nonstationary(; N = default_N, Δt = 0.1, kwargs...)
    u0 = [-5.0, -5.0, 0.0, 3.0] # u0[4] is initial `c` value
    p0 = [0.2, 0.2, (20 - 3)/(N*Δt)]
    ds = ContinuousDynamicalSystem(roessler_nonstationary_eom, u0, p0)
    tr = trajectory(ds, N*Δt; Δt, Ttr = 0, diffeq...)
end

function torus2(; N = default_N, Δt = 0.1, ω = sqrt(3), R = 2.0, r = 1.0, kwargs...)
    function torus(u)
        θ, φ = u
        x = (R + r*cos(θ))*cos(φ)
        y = (R + r*cos(θ))*sin(φ)
        z = r*sin(θ)
        return SVector(x, y, z)
    end
    θs = range(0; step = 0.1, length = N)
    φs = ω .* θs
    return Dataset([torus(u) for u in zip(θs, φs)])
end

function koch(; maxk = 7, kwargs...)
    flakepoints = SVector{2}.([[0.0; 0.0], [0.5; sqrt(3)/2], [1; 0.0], [0.0; 0.0]])
    function innerkoch(points, maxk, α = sqrt(3)/2)
        Q = @SMatrix [0 -1; 1 0]
        for k = 1:maxk
            n = length(points)
            new_points = eltype(points)[]
            for i = 1:n-1
                p1, p2 = points[i], points[i+1]
                v = (p2 - p1) / 3
                q1 = p1 + v
                q2 = p1 + 1.5v + α * Q * v
                q3 = q1 + v
                push!(new_points, p1, q1, q2, q3)
            end
            push!(new_points, points[end])
            points = new_points
        end
        return points
    end
    kochpoints = innerkoch(flakepoints, maxk)
    kochdata = Dataset(unique!(kochpoints))
end

function henonheiles_quasi(; N = default_N, Δt = 0.2, kwargs...)
    u0 = [0.0, 0.1, 0.5, 0.0]
    ds = Systems.henonheiles(u0)
    tr = trajectory(ds, N*Δt; Δt, Ttr = 100, diffeq...)
    standardize(tr)
end

function henonheiles_chaotic(; N = default_N, Δt = 0.5, kwargs...)
    u0 = [0.0, -0.25, 0.42081, 0.0]
    ds = Systems.henonheiles(u0)
    tr = trajectory(ds, N*Δt; Δt, Ttr = 100, diffeq...)
    standardize(tr)
end

function uniform_sphere(; N = default_N, kwargs...)
    A = SVector{3, Float64}[]
    i = 0
    while i < N
        x = rand(SVector{3, Float64}) .- 0.5
        if norm(x) ≤ 1/2√3
            push!(A, x)
            i += 1
        end
    end
    return Dataset(A)
end

function standardmap_chaotic(; N = default_N, k = 64.0, kwargs...)
    ds = Systems.standardmap([0.08152, 0.122717]; k)
    tr = trajectory(ds, N; Ttr = 100)
end

function henon_chaotic(; N = default_N, a = 1.4, b = 0.3, kwargs...)
    ds = Systems.henon([0.08152, 0.122717]; a, b)
    tr = trajectory(ds, N; Ttr = 100)
    standardize(tr)
end

function towel_chaotic(; N = default_N, η = 0, correlated = false, kwargs...)
    ds = Systems.towel([0.085, -0.121, 0.075])
    tr = trajectory(ds, N; Ttr = 100)
    tr = standardize(tr)
    if η ≠ 0 # add static noise
        if !correlated
            seed = Random.MersenneTwister(5653435)
            for i in 1:length(tr)
                tr[i] += η*randn(seed, SVector{3})
            end
        else correlated
            a = arma_noise(N+1)
            for i in 1:length(tr)
                tr[i] = tr[i] .+ η .* a[i]
            end
        end
    end
    return tr
end

function arma_noise(N)
    η = randn(N)
    s = ones(N)
    for n in 4:N
        s[n] = 1.625s[n-1] - 0.284s[n-2] - 0.355s[n-3] + η[n] - 0.96η[n-1]
    end
    s ./= std(s)
    return s
end

function rule_towel_noisy(x, p, n)
    @inbounds x1, x2, x3 = x[1], x[2], x[3]
    η, seed = p
    SVector(
        3.8*x1*(1-x1) - 0.05*(x2+0.35)*(1-2*x3),
        0.1*( (x2+0.35)*(1-2*x3) - 1 + η*randn(seed) )*(1 - 1.9*x1),
        3.78*x3*(1-x3)+0.2*x2
    )
end
function towel_dynamic_noise(; N = default_N, η = 0.1, kwargs...)
    p = [η, MersenneTwister(919009)]
    ds = DiscreteDynamicalSystem(rule_towel_noisy, [0.085, -0.121, 0.075], p)
    tr = trajectory(ds, N; Ttr = 100)
    tr = standardize(tr)
end

function lorenz96_chaotic(; N = default_N, Δt = 0.2, D = 4, F = 24.0, kwargs...)
    ds = Systems.lorenz96(D, range(0; length = D, step = 0.1); F)
    tr = trajectory(ds, N*Δt; Δt, Ttr = 1000.0, diffeq...)
    standardize(tr)
end

function lorenz96_rounded(; digits = 2, kwargs...)
    X = lorenz96_chaotic(; kwargs...)
    Z = [round.(x; digits) for x in X]
    return Dataset(Z)
end

function lorenz96_nonstationary_f!(dx, x, p, t)
    dF = p[1]
    N = length(x)-1
    F = x[end]

    # 3 edge cases
    @inbounds dx[1] = (x[2] - x[N - 1]) * x[N] - x[1] + F
    @inbounds dx[2] = (x[3] - x[N]) * x[1] - x[2] + F
    @inbounds dx[N] = (x[1] - x[N - 2]) * x[N - 1] - x[N] + F
    # then the general case
    for n in 3:(N - 1)
      @inbounds dx[n] = (x[n + 1] - x[n - 2]) * x[n - 1] - x[n] + F
    end
    dx[end] = dF
    return nothing
end

function lorenz96_nonstationary(; N = default_N, Δt = 0.2, D = 6, F2=24.0, F1=1.0, kwargs...)
    u0 = range(0.1; length = D, step = 0.1)
    dF = (F2-F1)/(N*Δt)
    u0 = [u0..., F1]
    ds = ContinuousDynamicalSystem(lorenz96_nonstationary_f!, u0, [dF])
    tr = trajectory(ds, N*Δt; Δt, Ttr = 0,  diffeq...)
    # standardize(tr)
end

function lorenz96_nonstationary_2(; N = default_N, Δt = 0.2, D = 6, F2=40.0, F1=10.0, kwargs...)
    ds = Systems.lorenz96(D, range(0; length = D, step = 0.1); F = F1)
    tr1 = trajectory(ds, (N/2)*Δt; Δt, Ttr = 1000.0, diffeq...)

    ds = Systems.lorenz96(D, range(0; length = D, step = 0.1); F = F2)
    tr2 = trajectory(ds, (N/2)*Δt; Δt, Ttr = 1000.0, diffeq...)

    return standardize(append!(tr1, tr2))
end

logistic(u) = 4*u*(1-u)
function logistics!(un, u, p, t)
    L = length(u)
    k, r, η = p # η is dynamic noise
    for i in 2:L-1
        un[i] = logistic(u[i] + k*(u[i-1] - 2u[i] + u[i+1]) + η*randn())
    end
    # circling around
    un[1] = logistic(u[1] + k*(u[L] - 2u[1] + u[2]) + η*randn())
    un[L] = logistic(u[L] + k*(u[L-1] - 2u[L] + u[1]) + η*randn())
    return
end

function coupled_logistics(; N = default_N, D = 8, k = 0.1, r = 4.0, η = 0, kwargs...)
    u0 = range(0.1, 0.9; length = D) .+ 0.0001
    ds = DiscreteDynamicalSystem(logistics!, u0, [k, r, η])
    tr = trajectory(ds, N; Ttr = 1000)
end

function embed_system(; system, d = 4, kwargs...)
    data_producing_function = getfield(Data, system)
    x = data_producing_function(; kwargs...)[:, 1]
    x = standardize(x)
    τ = estimate_delay(x, "mi_min")
    return embed(x, d, τ)
end

using DelimitedFiles
function experimental_data(; N = nothing, name, kwargs...)
    file = datadir("experimental", name*".txt")
    if name == "electrochemical1"
        τ = (0, 26, 13, 5, 20)
        x = vec(readdlm(file))
        A = genembed(x, τ)
    elseif name == "electrochemical2"
        τ = (0, 25, 16, 148, 138, 87, 60, 105)
        x = vec(readdlm(file))
        A = genembed(x, τ)
    elseif name == "shinriki1"
        τ = (0, 19, 38, 57)
        # js = ones(length(τ)) # perhaps use multivariate here
        x = readdlm(file)[:, 1]
        A = genembed(x, τ)
    elseif name == "shinriki2"
        τ = (0, 20, 40, 60)
        # js = ones(length(τ)) # perhaps use multivariate here
        x = readdlm(file)[:, 1]
        A = genembed(x, τ)
    elseif name == "nifty50"
        τ = [j*43 for j in 0:5]
        x = vec(readdlm(file))
        A = genembed(x, τ)
    elseif name == "vostok"
        τ = [j*50 for j in 0:7]
        x = vec(readdlm(file))
        A = genembed(x, τ)
    elseif name == "roessler_embed"
        τ = [0, 6, 3, 14]      # from PECUZAL
        τ = [j*7 for j in 0:6] # from TRADITIONAL
        x = readdlm(file)[:, 1]
        A = genembed(x, τ)
    elseif name == "roessler_all"
        M = readdlm(file)
        A = Dataset(M) # just take all 28 oscillators
    elseif name == "doublependulum"
        M = readdlm(file)
        τx = [0, 51, 25, 39, 12]
        τy = [0, 60, 30, 46, 14]
        jx = fill(1, length(τx))
        jy = fill(2, length(τy))
        τd = [0, 61, 31, 47, 31, 0]
        jd = [2, 2, 2, 2, 1, 1]
        A = Dataset(M)
        genembed(A, τx, jx)
    else
        error("Experimental dataset $(name) is uknown")
    end
    A = standardize(A)
    if isnothing(N)
        return A
    else
        return Dataset(A[1:min(length(A), N)])
    end
end

function electronic_roessler_mean_field(;R, Y = 1, d = 4, kwargs...)
    tr = mean(readdlm(datadir("exp_pro","R$(R)_ST_10_$Y.dat")); dims = 2)
    x = standardize(vec(tr))
    τ = estimate_delay(x, "mi_min")
    return embed(x, d, τ)
end

function electronic_roessler(; R=1, Y = 1, dms = 1:28, kwargs...)
    tr = readdlm(datadir("exp_pro","R$(R)_ST_10_$Y.dat"))[:,dms]
    standardize(Dataset(tr))
end

end


mutable struct LocalNNFunctionApproximator{NT<:NNTree, V<:AbstractVector{Float64}} <: LocalFunctionApproximator
    nntree::NT
    nnpoints::AbstractVector{V}
    nnvalues::Vector{Float64}
    knnK::Union{Int64,Nothing}
    rnnR::Union{Float64,Nothing}
end
# Two default constructors for k and r
LocalNNFunctionApproximator(nntree::NT, nnpts::AbstractVector{V}, knnK::Int64) where {NT <: NNTree,V<:AbstractVector{Float64}} =
    LocalNNFunctionApproximator(nntree, nnpts, zeros(length(nntree.indices)), knnK, nothing)
LocalNNFunctionApproximator(nntree::NT, nnpts::AbstractVector{V}, rnnR::Float64) where {NT <: NNTree,V<:AbstractVector{Float64}} =
    LocalNNFunctionApproximator(nntree, nnpts, zeros(length(nntree.indices)), nothing, rnnR)

################ INTERFACE FUNCTIONS ################
function n_interpolating_points(nnfa::LocalNNFunctionApproximator)
    return length(nnfa.nntree.indices)
end

function get_all_interpolating_points(nnfa::LocalNNFunctionApproximator)
  return nnfa.nnpoints
end

function get_all_interpolating_values(nnfa::LocalNNFunctionApproximator)
  return nnfa.nnvalues
end

function get_interpolating_nbrs_idxs_wts(nnfa::LocalNNFunctionApproximator, v::AbstractVector{Float64})

    @assert (nnfa.knnK != nothing || nnfa.rnnR != nothing)

    if nnfa.knnK != nothing
        # Do k-NN lookup to get data and distances
        idxs, dists = knn(nnfa.nntree, v, nnfa.knnK)
    else
        # Do inrange lookup to get data
        # Then use metric to get dists between query and each nearest neighbor
        idxs = inrange(nnfa.nntree, v, nnfa.rnnR)
        dists = zeros(length(idxs))
        for (i,idx) in enumerate(idxs)
            dists[i] = Distances.evaluate(nnfa.nntree.metric, v, nnfa.nnpoints[idx])
        end
    end

    weights = zeros(length(dists))

    # If exactly one point, set that probability of that idx to 1
    # and that of all others to 0.
    if minimum(dists) < eps(Float64)
        weights[argmin(dists)] = 1.0
    else
        for (i,d) in enumerate(dists)
            weights[i] = 1.0/d
        end
        weights /= sum(weights)
    end

    return (idxs, weights)
end

function compute_value(nnfa::LocalNNFunctionApproximator, v::AbstractVector{Float64})

    idxs, wts = get_interpolating_nbrs_idxs_wts(nnfa, v)

    # Do a weighted average of values
    value::Float64 = 0.0
    wtsum::Float64 = 0.0
    for (i,idx) in enumerate(idxs)
        value += wts[i]*nnfa.nnvalues[idx]
        wtsum += wts[i]
    end
    value /= wtsum

    return value
end

function compute_value(nnfa::LocalNNFunctionApproximator, v_list::AbstractVector{V}) where V <: AbstractVector{Float64}
    @assert length(v_list) > 0
    vals = [compute_value(nnfa, pt) for pt in v_list]
    return vals
end

function set_all_interpolating_values(nnfa::LocalNNFunctionApproximator, nnvalues::AbstractVector{Float64})
  nnfa.nnvalues = copy(nnvalues)
end


using Test, Random
using StatsBase, Distributions, DependentBootstrap
using ForecastEval

@testset "Diebold-Mariano tests" begin
	Random.seed!(1234)
	x = randn(100)
	y = randn(100)
	z = randn(100)
	ld = (x - z).^2 - (y - z).^2
	Random.seed!(1234)
	bO = dm(ld, :boot)
	hO = dm(ld, :hac)
	#v0.7+
	@test (bO.rejH0, bO.bestinput) == (false, 1)
	@test isapprox(bO.pvalue, 0.52)
	@test isapprox(bO.teststat, -0.23500198562777389)
	@test (hO.rejH0, hO.bestinput) == (false, 1)
	@test isapprox(hO.pvalue, 0.5016201102228858)
	@test isapprox(hO.teststat, -0.6719428022372147)
	#------------------------------------------
	#v0.6
	#@test (bO.rejH0, bO.bestinput) == (false, 1)
	#@test isapprox(bO.pvalue, 0.42599999999999993)
	#@test isapprox(bO.teststat, -0.23500198562777386)
	#@test (hO.rejH0, hO.bestinput) == (false, 1)
	#@test isapprox(hO.pvalue, 0.5016201102228858)
	#@test isapprox(hO.teststat, -0.6719428022372147)
end

@testset "Reality Check tests" begin
	Random.seed!(1234)
	ld = randn(100, 20)
	rcOut = rc(ld)
	#v0.7+
	@test rcOut.rejH0 == false
	@test isapprox(rcOut.pvalue, 0.778)
	#------------------------------------------
	#v0.6
	# @test rcOut.rejH0 == false
	# @test isapprox(rcOut.pvalue, 0.762)
end

@testset "SPA tests" begin
	Random.seed!(1234)
	ld = randn(100, 5)
	spaOut = spa(ld)
	#v0.7+
	@test spaOut.rejH0 == false
	@test isapprox(sum(spaOut.mu_u), 0.026125507727482872)
	@test isapprox(sum(spaOut.mu_c), 0.026125507727482872)
	@test isapprox(sum(spaOut.mu_l), 0.21071619368258687)
	@test isapprox(spaOut.pvalue_u, 0.29)
	@test isapprox(spaOut.pvalue_c, 0.29)
	@test isapprox(spaOut.pvalue_l, 0.181)
	@test isapprox(spaOut.pvalueauto, 0.29)
	@test isapprox(spaOut.teststat, 1.5384635269840627)
	#------------------------------------------
	#v0.6
	# @test spaOut.rejH0 == false
	# @test isapprox(sum(spaOut.mu_u), 0.026125507727482858)
	# @test isapprox(sum(spaOut.mu_c), 0.026125507727482858)
	# @test isapprox(sum(spaOut.mu_l), 0.21071619368258687)
	# @test isapprox(spaOut.pvalue_u, 0.25)
	# @test isapprox(spaOut.pvalue_c, 0.25)
	# @test isapprox(spaOut.pvalue_l, 0.124)
	# @test isapprox(spaOut.pvalueauto, 0.25)
	# @test isapprox(spaOut.teststat, 1.5384635269840627)
end

@testset "MCS tests" begin
	Random.seed!(1234)
	l = randn(100, 8)
	mcsOut = mcs(l)
	#v0.7+
	@test mcsOut.inQF == Int[5, 2, 3, 7, 4, 8, 6, 1]
	@test mcsOut.outQF == Int[]
	@test isapprox(sum(mcsOut.pvalueQF), 6.702999999999999)
	@test mcsOut.inMT == Int[5, 2, 3, 7, 4, 8, 6, 1]
	@test mcsOut.outMT == Int[]
	@test isapprox(sum(mcsOut.pvalueMT), 6.754000000000001)
	#------------------------------------------
	#v0.6
	# @test mcsOut.inQF == Int[5,2,3,7,4,8,6,1]
	# @test mcsOut.outQF == Int[]
	# @test isapprox(sum(mcsOut.pvalueQF), 6.665)
	# @test mcsOut.inMT == Int[5,2,3,7,4,8,6,1]
	# @test mcsOut.outMT == Int[]
	# @test isapprox(sum(mcsOut.pvalueMT), 6.673)
	#------------------------------------------
	Random.seed!(1234)
	mcsLROut = mcs(l, MCSBootLowRAM(l))
	#v0.7+
	@test mcsLROut.inQF == Int[5, 2, 3, 7, 4, 8, 6, 1]
	@test mcsLROut.outQF == Int[]
	@test isapprox(sum(mcsLROut.pvalueQF), 6.67)
	@test mcsLROut.inMT == Int[5, 2, 3, 7, 4, 8, 6, 1]
	@test mcsLROut.outMT == Int[]
	@test isapprox(sum(mcsLROut.pvalueMT), 6.719)
	#------------------------------------------
	#v0.6
	# @test mcsLROut.inQF == Int[5,2,3,7,4,8,6,1]
	# @test mcsLROut.outQF == Int[]
	# @test isapprox(sum(mcsLROut.pvalueQF), 6.604)
	# @test mcsLROut.inMT == Int[5,2,3,7,4,8,6,1]
	# @test mcsLROut.outMT == Int[]
	# @test isapprox(sum(mcsLROut.pvalueMT), 6.563)
end


@doc raw"""
    Direction3

An object `d` of the class [`Direction_2`](@ref) is a vector in the
three-dimensional vector space ``ℝ^3`` where we forget about its length.

They can be viewed as unit vectors, although there is no normalization
internally, since this is error prone. Directions are used whenever the
length of a vector does not matter. They also characterize a set of parallel
oriented lines that have the same orientations. For example, you can ask for
the direction orthogonal to an oriented plane, or the direction of an
oriented line.
"""
Direction3

"""
    Direction3(v::Vector3)

Introduces a direction `d` initialized with the direction of vector `v`.
"""
Direction3(v::Vector3)

"""
    Direction3(l::Line3)

Introduces the direction `d` of line `l`.
"""
Direction3(l::Line3)

"""
    Direction3(r::Ray3)

Introduces the direction `d` of ray `r`.
"""
Direction3(r::Ray3)

"""
    Direction3(s::Segment3)

Introduces the direction `d` of segment `s`.
"""
Direction3(s::Segment3)

"""
    Direction3(x::RT, y::RT, z::RT)
    Direction3(x::Real, y::Real, z::Real)

Introduces a direction `d` initialized with the direction from the origin and
the point with Cartesian coordinates ``(x, y, z)``.
"""
Direction3(x::Real, y::Real, z::Real) =
    Direction3(convert(FT, x), convert(FT, y), convert(FT, z))

@doc raw"""
    delta(d::Direction3, i::Integer)

Returns values, such that
`d == [Direction3](@ref)(delta(d, 0), delta(d, 1), delta(d, 3))`.

!!! info "Precondition"

    ``0 ≤ i ≤ 2``
"""
delta(d::Direction3, i::Integer)

"""
    dx(d::Direction3)

Returns `delta(d, 0)`.
"""
dx(d::Direction3)

"""
    dy(d::Direction3)

returns `delta(d, 1)`.
"""
dy(d::Direction3)

"""
    dz(d::Direction3)

returns `delta(d, 2)`.
"""
dz(d::Direction3)

"""
    -(d::Direction3)

The direction opposite to `d`.
"""
-(d::Direction3)

"""
    vector(d::Direction3)

Returns a vector that has the same direction as `d`.
"""
vector(d::Direction3)

"""
    transform(d::Direction3, t::AffTransformation2)

Returns the direction obtained by applying `t` on `d`.
"""
transform(d::Direction3, t::AffTransformation2)


@testset "Utilities" begin

    @testset "Group actions" begin
        f1 = s -> (s * s,);
        f2 = s-> (2s, -s, 5s);
        f3 = s -> (s + 1,);
        action1 = GroupActions(f1)
        action2 = GroupActions(f1, f2)
        action3 = GroupActions(f1, f2, f3)
        r1 = Int[]
        HC.apply_actions(action1, 3) do s
            push!(r1, s)
            false
        end
        @test r1 == [9]

        r2 = Int[]
        HC.apply_actions(action2, 3) do s
            push!(r2, s)
            false
        end
        @test r2 == [9, 18, -9, 45, 6, -3, 15]

        r3 = Int[]
        HC.apply_actions(action3, 3) do s
            push!(r3, s)
            false
        end
        @test r3 == [9, 18, 19, -9, -8, 45, 46, 10, 6, 7, -3, -2, 15, 16, 4]
        @test action3(3) == [3, 9, 18, 19, -9, -8, 45, 46, 10, 6, 7, -3, -2, 15, 16, 4]

        # also test with arrays
        g1 = s -> [s * s];
        g2 = s-> [2s, -s, 5s];
        g3 = s -> [s + 1];
        action1 = GroupActions(g1)
        action2 = GroupActions(g1, g2)
        action3 = GroupActions(g1, g2, g3)
        r1 = Int[]
        HC.apply_actions(action1, 3) do s
            push!(r1, s)
            false
        end
        @test r1 == [9]

        r2 = Int[]
        HC.apply_actions(action2, 3) do s
            push!(r2, s)
            false
        end
        @test r2 == [9, 18, -9, 45, 6, -3, 15]

        r3 = Int[]
        HC.apply_actions(action3, 3) do s
            push!(r3, s)
            false
        end
        r3
        @test r3 == [9, 18, 19, -9, -8, 45, 46, 10, 6, 7, -3, -2, 15, 16, 4]
    end

    @testset "UniquePoints" begin
        Random.seed!(1234)
        X = [randn(ComplexF64, 10) for _ = 1:2_000]
        indices = HC.unique!(rand(1:2000, 20))
        data = HC.UniquePoints(X)

        test_show_juno(data)
        @test string(data) == "UniquePoints{Array{Complex{Float64},1},Float64,typeof(euclidean_distance),Nothing} with 2000 points"

        @test length(data) == 2_000

        for i ∈ indices
            @test HC.is_contained(data, X[i])
            @test HC.is_contained(data, X[i], Val(true)) == i
            @test HC.is_contained(data, X[i] .+ 1e-4) == false
            @test HC.is_contained(data, X[i] .+ 1e-9, Val(true)) == i
            @test HC.is_contained(data, X[i] .+ 1e-9) == true
            @test HC.add!(data, X[i]) == false
            @test HC.add!(data, X[i], Val(true)) == i
            @test data[i] == X[i]
            @test HC.add!(data, X[i] .+ 1e-4) == true
            @test HC.add!(data, X[i] .- 1e-4, Val(true)) == -1
        end

        empty!(data)
        @test length(points(data)) == 0

        # Test many points with nearly indentical distance to the inserted point
        p = shuffle!([[cis(k/100*2π)] for k=0:99])
        data = HC.UniquePoints(p)
        @test HC.is_contained(data, [0.0im]) == false

        # Test with group action
        x = randn(ComplexF64, 4)
        permutation1(x) = ([x[2]; x[1]; x[3]; x[4]],)
        permutation2(x) = ([x[1]; x[2]; x[4]; x[3]],)
        X = [v for v in GroupActions(permutation1, permutation2)(x)]

        # One group action
        data = UniquePoints(X, group_action = permutation1)
        @test length(data) == 2

        # Two group actions
        data = HC.UniquePoints(X, group_actions = GroupActions(permutation1, permutation2))
        @test length(data) == 1

        # Group action and reality check
        x = randn(4)
        X = [im.*x, x, randn(ComplexF64, 4)]

        data = HC.UniquePoints(X, group_action = x -> (im.*x, (-1).*x, (-im).*x), check_real=false)
        @test HC.is_real_vector(points(data)[1]) == false
        @test length(points(data)) == 2

        data = HC.UniquePoints(X, group_action = x -> (im.*x, (-1).*x, (-im).*x), check_real = true)
        @test HC.is_real_vector(points(data)[1]) == true
        @test length(points(data)) == 2

        data = HC.UniquePoints(randn(ComplexF64, 4), check_real = true)
        k = add!(data, randn(4), Val{true}())
        @test k == -2
        k = add!(data, randn(ComplexF64, 4), Val{true}())
        @test k == -1

        pts = [[1.0,0.5], [1.0,0.5], [0.5,1.0]]
        @test points(UniquePoints(pts)) == [[1.0, 0.5], [0.5, 1.0]]
        @test unique_points(pts) == [[1.0, 0.5], [0.5, 1.0]]
        @test unique_points(pts; group_action = x -> ([x[2],x[1]],)) == [[1.0, 0.5]]
        @test unique_points(pts; group_action = x -> [x[2],x[1]]) == [[1.0, 0.5]]
    end

    @testset "Multiplicities" begin
        # test empty input
        @test isempty(multiplicities(Vector{Vector{ComplexF64}}()))

        V = [randn(3) + im.* randn(3) for i in 1:10]
        W = [map(v -> v + [1e-7 * v[1] ;1.0;1.0], V); V; map(v -> v + [1e-6 * v[1]; 0; 0], V)]
        U = [[cis(rand()) for _=1:3] for _=1:20]
        X = [[0.499064-0.707804im, -0.164223+0.0933025im, -0.00695332-0.139516im, 0.000716617-0.135055im, 0.009846-0.12984im, 0.020204-0.123609im, 0.34326+0.161402im], [-0.795428-0.0975931im, -0.20942-0.382162im, -0.119855-0.0414809im, -0.123105-0.0293171im, -0.128161-0.0148951im, -0.134641+0.00272884im, 0.258239+0.185188im], [0.605937-0.289401im, 0.42308-0.188036im, -0.057496+0.200812im,0.284444+0.0477843im, 0.191189-0.225824im, -0.233732-0.181827im, -0.0939606+0.154403im], [-0.493372+0.512561im, -0.385166+0.204525im, -0.120853-0.142237im, 0.0742939+0.163197im, -0.298456+0.202635im, 0.100108-0.219369im, 0.193075+0.105514im], [0.483174-0.581023im, 0.23432-0.0528133im, -0.111264-0.067633im, -0.292867-0.014453im, 0.24805-0.267034im, 0.141232-0.13631im, 0.168603+0.261794im], [-0.455116+0.4482im, -0.366004+0.264082im, -0.0578486+0.110708im, -0.296085-0.0619238im, 0.216907-0.327015im, 0.207934+0.242906im, -0.121251+0.101268im], [0.526416-0.465825im, 0.39214-0.0908531im, -0.260222+0.00391701im, -0.110348-0.1151im, 0.101057-0.190265im, 0.308379-0.137148im, 0.0488701+0.29665im], [0.147886-0.750864im, -0.235168-0.257487im, 0.0394847+0.194512im, -0.0325228-0.227443im, -0.111526-0.214717im, 0.0245065+0.162303im, 0.304189-0.150105im], [-0.372175-0.538886im, -0.117563-0.341994im, 0.0752294-0.0141672im, 0.143671-0.249302im, -0.374165-0.228281im, -0.0971581+0.346633im, 0.168499+0.0394346im], [-0.503759+0.48im, -0.447123+0.262057im, -0.214529+0.109917im, 0.148493-0.0330798im, -0.0510215-0.262778im, -0.110165+0.198317im, 0.200704+0.0508037im], [0.623685-0.513102im, 0.0522679+0.0223332im, -0.188478-0.10101im, 0.229926-0.157889im, -0.146424-0.0562049im, 0.209112-0.185147im, 0.205229+0.276164im], [0.127017-0.655061im, -0.147053-0.569588im, -0.134863+0.0901295im, 0.152888+0.049255im, -0.124224+0.0624774im, 0.152924+0.0778815im, 0.105575-0.310906im], [0.575711+0.384346im, 0.336129+0.0531313im,0.155698-0.152359im, -0.104555+0.197167im, -0.254208-0.215823im, 0.327345+0.0728999im, -0.0274605+0.288691im], [0.479139-0.593553im, -0.0960412-0.315975im, 0.0438803-0.245038im, 0.000516985+0.193173im, -0.0828034+0.172297im, -0.00479366-0.21881im, 0.353615-0.0168729im], [-0.37103+0.609109im, -0.201984+0.259807im, -0.158043+0.252162im, -0.211405-0.106454im, 0.0274292-0.262351im, 0.333102+0.0455649im, -0.0847849+0.220536im], [-0.200697+0.682536im, -0.247073+0.359358im, -0.172902-0.0702972im, 0.141297+0.110399im, -0.171049+0.317616im, -0.0125869-0.240803im, 0.219982+0.00435721im], [-0.0896611-0.805671im, -0.0393637-0.183565im, 0.224508-0.167715im, 0.0918438-0.180778im, -0.0879493-0.157071im, -0.250607-0.0700599im, 0.0521971+0.291749im], [-0.226054+0.605127im, -0.210133+0.55141im,0.218162-0.0576634im, 0.08772-0.0936im, -0.0820928-0.101495im, -0.22419-0.0507194im, -0.0146449+0.311527im], [0.498516+0.360158im, 0.219394+0.0657585im, -0.0836376+0.0106629im, 0.00702255-0.322935im, 0.481634+0.13367im, -0.133322+0.397338im, -0.171875-0.0532251im], [0.440673-0.62568im, -0.110019-0.330907im, -0.043096+0.193744im, 0.0628576-0.220991im, -0.0144861-0.24152im, -0.0436694+0.158228im, 0.338948-0.0133119im], [0.622749+0.219955im, 0.628838+0.216896im, 0.0331851+0.145177im, -0.123675-0.0636893im, 0.0348129+0.117408im, -0.133898-0.0420959im, 0.155277-0.144857im], [0.652186-0.0569317im, 0.552174-0.169926im, -0.0977967-0.0521929im, -0.161552-0.169677im, 0.00437006+0.237406im, -0.0164646+0.116656im, 0.310812-0.0599935im], [0.448293+0.574542im, 0.247779+0.106134im, -0.0955677-0.215494im, 0.0935008+0.252581im, 0.336093-0.0655203im, -0.242972-0.0256941im, -0.0177608+0.301502im], [0.445422+0.671948im, -0.0726023+0.313409im, -0.044775+0.21726im,0.154346+0.0261763im, -0.0305679+0.188435im, 0.20006+0.0391233im, -0.0810809-0.297159im], [-0.00979382-0.81029im, 0.106391-0.195197im, 0.108622-0.139201im, 0.234362-0.1806im, -0.20136-0.131455im, -0.11823-0.118552im, -0.0308846+0.297292im], [0.318611+0.584183im, 0.105547+0.221636im, 0.248174+0.0583734im, -0.00256072-0.290164im, -0.302039-0.0824955im, -0.114841+0.409322im, 0.25213+0.0736154im], [0.696901+0.0684572im, 0.505426+0.266373im, -0.152493+0.0711982im, 0.197532-0.0412531im, 0.214441+0.0165324im, -0.096866+0.0847148im, -0.103987-0.201388im], [0.115188-0.657244im, -0.157297-0.566844im, -0.133216+0.092546im, 0.153751+0.0464909im, -0.123077+0.0647066im, 0.154303+0.0751121im, 0.0999536-0.312759im], [0.653199-0.277909im, 0.556407-0.00960311im, 0.204776-0.112518im, -0.112199+0.0771487im, -0.0695128+0.163961im, 0.167741-0.0723464im, -0.16148-0.148729im], [0.411538+0.550657im, 0.413916+0.420472im, -0.0555052+0.089203im, -0.0609812+0.227821im, -0.0913546-0.187807im, -0.0607857-0.0922333im, 0.237526+0.0198403im], [0.602049+0.184366im, 0.695878+0.270759im, -0.0420156+0.0359236im, -0.0431941+0.0331572im, -0.0446403+0.0302374im, -0.046326+0.027106im, 0.0979775-0.156785im], [-0.270799+0.567571im, -0.113714+0.377218im, 0.0592799-0.219174im, 0.226587+0.091218im, 0.0269123+0.326853im, -0.402766+0.100383im, -0.018186-0.240642im]]

        M = multiplicities(V, distance=infinity_distance)
        @test length(M) == 0

        N = multiplicities(W, distance=infinity_distance)
        sort!(N, by=first)
        @test length(N) == 10
        @test unique([length(m) for m in N]) == [2]
        @test N == [[i, i+10] for i=11:20]

        O = multiplicities([U;U], distance=infinity_distance)
        @test length(O) == 20

        P = multiplicities(X, distance=(x,y) -> 1-abs(LinearAlgebra.dot(x,y)))
        @test length(P) == 3

        # Test with group action
        x = randn(ComplexF64, 4)
        permutation1(x) = ([x[2]; x[1]; x[3]; x[4]],)
        permutation2(x) = ([x[1]; x[2]; x[4]; x[3]],)
        X = GroupActions(permutation1, permutation2)(x)

        # One group action
        m = multiplicities(X, group_action = permutation1) |> sort
        @test m == [[1, 2], [3, 4]]

        # Two group actions
        m = multiplicities(X, group_actions = GroupActions(permutation1, permutation2))
        @test m == [[1, 2, 3, 4]]

        # docs
        X = [[1, 2, 3, 4], [2,1,3,4], [1,2,4,3], [2,1,4,3]]
        permutation(x) = [x[2], x[1], x[3], x[4]]
        m = sort(multiplicities(X, group_action = permutation))
        @test m == [[1,2], [3,4]]
    end

    @testset "Polynomials" begin
        @polyvar p q a b c x y z
        e = [p + 1]
        f = [a * b * c ]
        g = [x+y, y + z, x + z]
        @test expand(f ∘ g) == [x^2*y + x^2*z + x*y^2 + 2*x*y*z + x*z^2 + y^2*z + y*z^2]
        @test expand(e ∘ f ∘ g) == [x^2*y + x^2*z + x*y^2 + 2*x*y*z + x*z^2 + y^2*z + y*z^2 + 1]
        @test expand(e ∘ (f ∘ g)) == [x^2*y + x^2*z + x*y^2 + 2*x*y*z + x*z^2 + y^2*z + y*z^2 + 1]
        @test expand((e ∘ f) ∘ g) == [x^2*y + x^2*z + x*y^2 + 2*x*y*z + x*z^2 + y^2*z + y*z^2 + 1]
        @test e ∘ f == e ∘ f

        @test validate(f ∘ g) == true
        @test validate(f ∘ g, parameters=[z]) == true
        @test validate(f ∘ g, parameters=[c]) == false
        @test validate(g ∘ f) == false

        @test HC.nvariables(f ∘ g) == 3
        @test HC.nvariables(f ∘ g, parameters=[z]) == 2

        @polyvar x y z
        @test is_homogeneous([x^2+y^2+x*y, x^5])
        @test is_homogeneous([x^2+y^2+x*y, x^4+1]) == false

        #test weighted degree
        @test is_homogeneous(x^3+x*y, [(x, 2), (y, 4)])
        @test homogenize(x+x*y, [(x, 2), (y, 4)], z) == x*z^4+x*y

        @test is_homogeneous(homogenize([x^2+y^2+x*y, x^4+1]))
        @test is_homogeneous([x^2+z^2 + y, x^4+z^4], [x,z]) == false
        @test is_homogeneous(homogenize([x^2+z^2*y, x^4+z^4*y], [x,z]), [x,z]) == true

        @test is_homogeneous(f ∘ g) == true
        h = [a * b * c  + p * c^3] ∘ [x+y, y + z, x + z]
        @test is_homogeneous(h, parameters=[p]) == true
        h2 = [a * b * c  + p] ∘ [x+y, y + z, x + z]
        @test is_homogeneous(h2, parameters=[p]) == false

        #homogenize
        @test homogenize(x^2+y+1, z) == x^2+y*z + z^2
        # This needs to be an array due to a compiler bug
        @test homogenize([x^2+y+p], z, parameters=[p]) == [x^2+y*z + z^2*p]

        @test homogenize([x^3+p, 1+y], z, parameters=[p]) == [x^3+p*z^3, z+y]

        h2 = [a * b * c  + p] ∘ [x+y, y + z, x + q]
        @test validate(homogenize(h2, parameters=[p, q]), parameters=[p, q])
        h3 = [a * b * c  + 1] ∘ [x+y, y + z, x + 1]
        @test validate(homogenize(h3))

        # Weylnorm
        @polyvar x y z
        f = 3.0x^2 + 2x*y - y^2
        g = (-2.5+2im) * x^2 - 3.0*x*y + 4y^2
        @test HC.weyldot(f, g) == 3.0 * conj(-2.5 + 2im) + 2.0 * (-3.0) / 2 + (-1.0) * 4.0
        @test HC.weyldot(f, f) == 9.0 + 4.0  / 2 + 1.0
        @test HC.weylnorm(f)^2 ≈ HC.weyldot(f,f)
        @test HC.weyldot([f, f], [g, g]) == 2 * HC.weyldot(f, g)
        @test HC.weylnorm([f, f]) == √HC.weyldot([f, f], [f, f])

        # linear system
        @polyvar x y z
        f = [2x+3y+z+5, -1x+2y+4z-2]
        A, b = linear_system(f)
        @test A == [2 3 1; -1 2 4]
        @test b == [-5, 2]
        @test A * [x,y,z] - b == f
    end

    @testset "Multihomogeneous" begin
        @polyvar x y
        f = [x*y-2, x^2-4]

        affine_hominfo = HC.HomogenizationInformation(variable_groups = ((x,), (y,)))
        @test HC.is_homogeneous(f, affine_hominfo) == false

        @polyvar v w
        g = [x * y - 2v * w, x^2-4v^2]
        hominfo = HC.HomogenizationInformation(variable_groups = ((x,v), (y,w)))
        @test HC.is_homogeneous(g, hominfo) == true

        hominfo = HC.HomogenizationInformation(homvars=(v,w), variable_groups = ((x,v), (y,w)))
        @test HC.is_homogeneous(g, hominfo) == true

        @test HC.homogenize(f, hominfo) == g
        HC.multidegrees(f, ([x], [y])) == [1 2; 1 0]
    end

    @testset "Linear algebra" begin
        # square LU
        A = rand(Complex{Float64}, 12, 12)
        b = rand(Complex{Float64}, 12)
        x = zeros(ComplexF64, 12)
        jac = HC.Jacobian(A)
        HC.update!(jac, A)
        C, d = copy(A), copy(b)
        HC.solve!(x, jac, b)
        @test norm(x - C \ d) < 1e-10

        # square QR / update infos (full rank)
        HC.update!(jac, A; update_infos=true)
        @test jac.corank == 0
        @test jac.cond > 1
        HC.solve!(x, jac, b; update_digits_lost=false)
        @test norm(x - C \ d) < 1e-10

        x .= 0
        HC.solve!(x, jac, b; update_digits_lost=true)
        @test norm(x - C \ d) < 1e-10

        # square QR / not full rank
        A = rand(Complex{Float64}, 12, 12)
        b = rand(Complex{Float64}, 12)
        x = zeros(ComplexF64, 12)
        jac = HC.Jacobian(A)

        # make a corank 1 matrix
        U, s, V = svd(A)
        s[end] = 1e-16
        Â = U * diagm(0 => s) * V'
        HC.update!(jac, Â; update_infos=true)
        @test jac.corank_proposal == 1
        HC.update_rank!(jac)
        @test jac.corank == 1
        @test jac.cond > 1e14
        HC.solve!(x, jac, copy(b); update_digits_lost=false)
        @test norm(x - qr(Â ./ jac.D, Val(true)) \ (b ./ jac.D)) < 1e-10

        x .= 0
        HC.solve!(x, jac, copy(b); update_digits_lost=true)
        @test norm(x - qr(Â ./ jac.D, Val(true)) \ (b ./ jac.D)) < 1e-10
        @test jac.digits_lost == nothing

        # make a corank 1 matrix with not super clear rank
        U, s, V = svd(A)
        s[end] = 1e-15
        Â = U * diagm(0 => s) * V'
        HC.update!(jac, Â; update_infos=true)
        @test jac.corank_proposal == 1
        HC.update_rank!(jac)
        @test jac.corank == 1
        @test jac.cond > 1e13
        HC.solve!(x, jac, copy(b); update_digits_lost=false)
        @test norm(x - qr(Â ./ jac.D, Val(true)) \ (b ./ jac.D)) < 1e-10

        # make a corank 2 matrix
        U, s, V = svd(A)
        s[end-1:end] .= 1e-16
        Â = U * diagm(0 => s) * V'
        HC.update!(jac, Â; update_infos=true)
        @test jac.corank_proposal == 2
        HC.update_rank!(jac)
        @test jac.corank == 2
        @test jac.cond > 1e14
        HC.solve!(x, jac, copy(b); update_digits_lost=false)
        @test norm(x - qr(Â ./ jac.D, Val(true)) \ (b ./ jac.D)) < 1e-10
        x .= 0
        HC.solve!(x, jac, copy(b); update_digits_lost=true)
        @test norm(x - qr(Â ./ jac.D, Val(true)) \ (b ./ jac.D)) < 1e-10
        @test jac.digits_lost == nothing

        # overdetermined
        A = rand(Complex{Float64}, 15, 12)
        b = rand(Complex{Float64}, 15)
        C, d = copy(A), copy(b)
        jac = HC.Jacobian(A)
        HC.update!(jac, A)
        jac.D .= 1
        x = zeros(ComplexF64, 12)
        HC.solve!(x, jac, b)
        @test norm(x - qr(C, Val(true)) \ d) < 1e-10

        x .= 0
        HC.solve!(x, jac, b; update_digits_lost=true)
        @test norm(x - qr(C, Val(true)) \ d) < 1e-10
        @test jac.digits_lost == nothing

        # make a corank 1 matrix
        U, s, V = svd(A)
        s[end] = 1e-15
        Â = U * diagm(0 => s) * V'
        HC.update!(jac, Â; update_infos=true)
        @test jac.corank_proposal == 1
        HC.update_rank!(jac)
        @test jac.corank == 1
        @test jac.cond > 1e14
        HC.solve!(x, jac, copy(b); update_digits_lost=false)
        @test norm(x - qr(Â ./ jac.D, Val(true)) \ (b ./ jac.D)) < 1e-10

        x .= 0
        HC.solve!(x, jac, copy(b); update_digits_lost=true)
        @test norm(x - qr(Â ./ jac.D, Val(true)) \ (b ./ jac.D)) < 1e-10
        @test jac.digits_lost == nothing
    end

    @testset "Misc" begin
        @test HC.unpack(5, 2) == 5
        @test HC.unpack(nothing, 2) == 2

        x = rand()
        @test HC.nthroot(x, 1) == x
        @test HC.nthroot(x, 0) == one(x)

        x = rand(6)
        @test HC.is_real_vector(x)

        segment = HC.ComplexSegment(2, 4)
        @test length(segment) ≈ 2
        test_show_juno(segment)
        @test string(segment) == "ComplexSegment(2.0 + 0.0im, 4.0 + 0.0im)"
    end

    @testset "Hermite Normal Form" begin
        A = [0 3 1 0; -2 2 -1 2; 1 -1 2 3; -3 3 3 2]
        H, U = HC.hnf(A)
        @test A * U == H

        D = zeros(Int, 1,1)
        D[1,1] = -2
        H, U = HC.hnf(D)
        @test H[1,1] == 2
        @test U[1,1] == -1
    end
end


@testset "Steiner Tree" begin

    function sum_weight(g::AbstractGraph{<:Integer}, distmx::AbstractMatrix{<:Integer} = weights(g))
        sum_wt = zero(eltype(g))
        for e in edges(g)
            sum_wt += distmx[src(e), dst(e)]
        end
        return sum_wt
    end

    approx_factor(t::Integer) = 2-2/t

    g3 = star_graph(5)
    for g in testgraphs(g3)
        g_st = @inferred(steiner_tree(g, [2, 5]))
        @test ne(g_st) == 2 # [Edge(2, 1), Edge(1, 5)]

        g_copy = SimpleGraph(g)
        LightGraphs.filter_non_term_leaves!(g_copy, [2, 5])
        @test ne(g_copy) == 2 # [Edge(2, 1), Edge(1, 5)]
    end

    g4 = path_graph(11)
    for g in testgraphs(g4)
        g_st = @inferred(steiner_tree(g, [4, 8]))
        @test ne(g_st) == 4

        g_copy = SimpleGraph(g)
        LightGraphs.filter_non_term_leaves!(g_copy, [4, 8])
        @test ne(g_copy) == 4
    end

    g5 = grid([5, 5])
    for g in testgraphs(g5)
        g_st = @inferred(steiner_tree(g, [3, 11, 15, 23]))
        @test sum_weight(g_st) == 8
    end

    d = [0   2   3   4   5
         2   0   60   80  1
         3   60   0  120  150
         4   80  120  0  200
         5  1  150  200  0]

    g6 = complete_graph(5)

    for g in testgraphs(g6)
        g_st = @inferred(steiner_tree(g, [2, 4, 5], d))
        @test sum_weight(g_st, d) <= approx_factor(3)*(1+2+4)
    end

    d[2, 5] = d[5, 2] = 100

    for g in testgraphs(g6)
        g_st = @inferred(steiner_tree(g, [2, 4, 5], d))
        @test sum_weight(g_st, d) <= approx_factor(3)*(2+4+5)
    end

end


function test_item_mean()
    println("-- Testing ItemMean recommender")

    data = DataAccessor(sparse([1 2 3; 4 5 6]))
    recommender = ItemMean(data)
    build!(recommender)
    actual = predict(recommender, 1, 1)

    @test actual == 2.5
end

test_item_mean()


using TaylorExponentialMatrix
using Test

@testset " simplest implementation " begin

A = rand( 5, 5)

@test all( expm2(A) .≈ exp(A))

end

@testset " sophisticated implementation " begin

A = rand( 5, 5)

@test all( expm3(A) .≈ exp(A))

end


using Test
using Omega.Space

@testset "Increment" begin
  x = [1,2,3]

  @test increment(x) == [1,2,4]
end

# Vector indices should be a free semigroup
#
# Properties:
# combine(a, combine(b, c)) === combine(combine(a, b), c)
# (combine(a, b) == combine(a, c)) => b == c
# (combine(a, c) == combine(b, c)) => a == b
# append(a, b) === combine(a, base(b))
@testset "combine" begin
  test_elts = [ [1,2], [3], [3, 4]
              , [5, 6], [7], [1,2,3]
              ]
  test_indivs = collect(1:10)

  triples = [(a,b,c) for a in test_elts, b in test_elts, c in test_elts]

  for (a,b,c) in triples
    @test combine(a, combine(b, c)) == combine(combine(a, b), c)
    if combine(a, b) == combine(a, c)
      @test b == c
    end
    if combine(a, c) == combine(b, c)
      @test a == b
    end
  end

  for elt in test_elts
    for x in test_indivs
      @test append(elt, x) == combine(elt, Space.base(Vector{Int}, x))
    end
  end

end