# Correlation

#######################################
#
#   Helper functions
#
#######################################

default_laglen(lx::Int) = min(lx-1, int(10log10(lx)))
check_lags(lx::Int, lags::AbstractVector) = (maximum(lags) < lx || error("lags must be less than the sample length."))

function demean_col!{T<:RealFP}(z::Vector{T}, x::Matrix{T}, j::Int, demean::Bool)
    m = size(x, 1)
    @assert m == length(z)
    b = m * (j-1)
    if demean
        s = zero(T)
        for i = 1 : m
            s += x[b + i]
        end
        mv = s / m
        for i = 1 : m
            z[i] = x[b + i] - mv
        end
    else
        copy!(z, 1, x, b+1, m)
    end
    z
end

#######################################
#
#   Auto-correlations
#
#######################################

default_autolags(lx::Int) = 0 : default_laglen(lx)

_autodot{T<:RealFP}(x::Vector{T}, lx::Int, l::Int) = dot(x, 1:lx-l, x, 1+l:lx)

## autocov

function autocov!{T<:RealFP}(r::RealVector, x::Vector{T}, lags::IntegerVector; demean::Bool=true)
    lx = length(x)
    m = length(lags)
    length(r) == m || raise_dimerror()
    check_lags(lx, lags)

    z::Vector{T} = demean ? x - mean(x) : x
    for k = 1 : m  # foreach lag value
        r[k] = _autodot(z, lx, lags[k]) / lx
    end
    return r
end

function autocov!{T<:RealFP}(r::RealMatrix, x::Matrix{T}, lags::IntegerVector; demean::Bool=true)
    lx = size(x, 1)
    ns = size(x, 2)
    m = length(lags)
    size(r) == (m, ns) || raise_dimerror()
    check_lags(lx, lags)

    z = Array(T, lx)
    for j = 1 : ns
        demean_col!(z, x, j, demean)
        for k = 1 : m
            r[k,j] = _autodot(z, lx, lags[k]) / lx
        end
    end
    return r
end

function autocov{T<:Real}(x::Vector{T}, lags::IntegerVector; demean::Bool=true)
    autocov!(Array(fptype(T), length(lags)), float(x), lags; demean=demean)
end

function autocov{T<:Real}(x::Matrix{T}, lags::IntegerVector; demean::Bool=true)
    autocov!(Array(fptype(T), length(lags), size(x,2)), float(x), lags; demean=demean)
end

autocov{T<:Real}(x::VecOrMat{T}; demean::Bool=true) = autocov(x, default_autolags(size(x,1)); demean=demean)

## autocor

function autocor!{T<:RealFP}(r::RealVector, x::Vector{T}, lags::IntegerVector; demean::Bool=true)
    lx = length(x)
    m = length(lags)
    length(r) == m || raise_dimerror()
    check_lags(lx, lags)

    z::Vector{T} = demean ? x - mean(x) : x
    zz = dot(z, z)
    for k = 1 : m  # foreach lag value
        r[k] = _autodot(z, lx, lags[k]) / zz
    end
    return r
end

function autocor!{T<:RealFP}(r::RealMatrix, x::Matrix{T}, lags::IntegerVector; demean::Bool=true)
    lx = size(x, 1)
    ns = size(x, 2)
    m = length(lags)
    size(r) == (m, ns) || raise_dimerror()
    check_lags(lx, lags)

    z = Array(T, lx)
    for j = 1 : ns
        demean_col!(z, x, j, demean)
        zz = dot(z, z)
        for k = 1 : m
            r[k,j] = _autodot(z, lx, lags[k]) / zz
        end
    end
    return r
end

function autocor{T<:Real}(x::Vector{T}, lags::IntegerVector; demean::Bool=true)
    autocor!(Array(fptype(T), length(lags)), float(x), lags; demean=demean)
end

function autocor{T<:Real}(x::Matrix{T}, lags::IntegerVector; demean::Bool=true)
    autocor!(Array(fptype(T), length(lags), size(x,2)), float(x), lags; demean=demean)
end

autocor{T<:Real}(x::VecOrMat{T}; demean::Bool=true) = autocor(x, default_autolags(size(x,1)); demean=demean)

#######################################
#
#   Cross-correlations
#
#######################################

default_crosslags(lx::Int) = (l=default_laglen(lx); -l:l)

_crossdot{T<:RealFP}(x::Vector{T}, y::Vector{T}, lx::Int, l::Int) = (l >= 0 ? dot(x, 1:lx-l, y, 1+l:lx) : dot(x, 1-l:lx, y, 1:lx+l))

## crosscov

function crosscov!{T<:RealFP}(r::RealVector, x::Vector{T}, y::Vector{T}, lags::IntegerVector; demean::Bool=true)
    lx = length(x)
    m = length(lags)
    (length(y) == lx && length(r) == m) || raise_dimerror()
    check_lags(lx, lags)

    zx::Vector{T} = demean ? x - mean(x) : x
    zy::Vector{T} = demean ? y - mean(y) : y
    for k = 1 : m  # foreach lag value
        r[k] = _crossdot(zx, zy, lx, lags[k]) / lx
    end
    return r
end

function crosscov!{T<:RealFP}(r::RealMatrix, x::Matrix{T}, y::Vector{T}, lags::IntegerVector; demean::Bool=true)
    lx = size(x, 1)
    ns = size(x, 2)
    m = length(lags)
    (length(y) == lx && size(r) == (m, ns)) || raise_dimerror()
    check_lags(lx, lags)

    zx = Array(T, lx)
    zy::Vector{T} = demean ? y - mean(y) : y
    for j = 1 : ns
        demean_col!(zx, x, j, demean)
        for k = 1 : m
            r[k,j] = _crossdot(zx, zy, lx, lags[k]) / lx
        end
    end
    return r
end

function crosscov!{T<:RealFP}(r::RealMatrix, x::Vector{T}, y::Matrix{T}, lags::IntegerVector; demean::Bool=true)
    lx = length(x)
    ns = size(y, 2)
    m = length(lags)
    (size(y, 1) == lx && size(r) == (m, ns)) || raise_dimerror()
    check_lags(lx, lags)

    zx::Vector{T} = demean ? x - mean(x) : x
    zy = Array(T, lx)
    for j = 1 : ns
        demean_col!(zy, y, j, demean)
        for k = 1 : m
            r[k,j] = _crossdot(zx, zy, lx, lags[k]) / lx
        end
    end
    return r
end

function crosscov!{T<:RealFP}(r::AbstractArray{T,3}, x::Matrix{T}, y::Matrix{T}, lags::IntegerVector; demean::Bool=true)
    lx = size(x, 1)
    nx = size(x, 2)
    ny = size(y, 2)
    m = length(lags)
    (size(y, 1) == lx && size(r) == (m, nx, ny)) || raise_dimerror()
    check_lags(lx, lags)

    # cached (centered) columns of x
    zxs = Array(Vector{T}, 0)
    sizehint(zxs, nx)
    for j = 1 : nx
        xj = x[:,j]
        if demean
            mv = mean(xj)
            for i = 1 : lx
                xj[i] -= mv
            end
        end
        push!(zxs, xj)
    end

    zx = Array(T, lx)
    zy = Array(T, lx)
    for j = 1 : ny
        demean_col!(zy, y, j, demean)
        for i = 1 : nx
            zx = zxs[i]
            for k = 1 : m
                r[k,i,j] = _crossdot(zx, zy, lx, lags[k]) / lx
            end
        end
    end
    return r
end

function crosscov{T<:Real}(x::Vector{T}, y::Vector{T}, lags::IntegerVector; demean::Bool=true)
    crosscov!(Array(fptype(T), length(lags)), float(x), float(y), lags; demean=demean)
end

function crosscov{T<:Real}(x::Matrix{T}, y::Vector{T}, lags::IntegerVector; demean::Bool=true)
    crosscov!(Array(fptype(T), length(lags), size(x,2)), float(x), float(y), lags; demean=demean)
end

function crosscov{T<:Real}(x::Vector{T}, y::Matrix{T}, lags::IntegerVector; demean::Bool=true)
    crosscov!(Array(fptype(T), length(lags), size(y,2)), float(x), float(y), lags; demean=demean)
end

function crosscov{T<:Real}(x::Matrix{T}, y::Matrix{T}, lags::IntegerVector; demean::Bool=true)
    crosscov!(Array(fptype(T), length(lags), size(x,2), size(y,2)), float(x), float(y), lags; demean=demean)
end

crosscov{T<:Real}(x::VecOrMat{T}, y::VecOrMat{T}; demean::Bool=true) = crosscov(x, y, default_crosslags(size(x,1)); demean=demean)

## crosscor

function crosscor!{T<:RealFP}(r::RealVector, x::Vector{T}, y::Vector{T}, lags::IntegerVector; demean::Bool=true)
    lx = length(x)
    m = length(lags)
    (length(y) == lx && length(r) == m) || raise_dimerror()
    check_lags(lx, lags)

    zx::Vector{T} = demean ? x - mean(x) : x
    zy::Vector{T} = demean ? y - mean(y) : y
    sc = sqrt(dot(zx, zx) * dot(zy, zy))
    for k = 1 : m  # foreach lag value
        r[k] = _crossdot(zx, zy, lx, lags[k]) / sc
    end
    return r
end

function crosscor!{T<:RealFP}(r::RealMatrix, x::Matrix{T}, y::Vector{T}, lags::IntegerVector; demean::Bool=true)
    lx = size(x, 1)
    ns = size(x, 2)
    m = length(lags)
    (length(y) == lx && size(r) == (m, ns)) || raise_dimerror()
    check_lags(lx, lags)

    zx = Array(T, lx)
    zy::Vector{T} = demean ? y - mean(y) : y
    yy = dot(zy, zy)
    for j = 1 : ns
        demean_col!(zx, x, j, demean)
        sc = sqrt(dot(zx, zx) * yy)
        for k = 1 : m
            r[k,j] = _crossdot(zx, zy, lx, lags[k]) / sc
        end
    end
    return r
end

function crosscor!{T<:RealFP}(r::RealMatrix, x::Vector{T}, y::Matrix{T}, lags::IntegerVector; demean::Bool=true)
    lx = length(x)
    ns = size(y, 2)
    m = length(lags)
    (size(y, 1) == lx && size(r) == (m, ns)) || raise_dimerror()
    check_lags(lx, lags)

    zx::Vector{T} = demean ? x - mean(x) : x
    zy = Array(T, lx)
    xx = dot(zx, zx)
    for j = 1 : ns
        demean_col!(zy, y, j, demean)
        sc = sqrt(xx * dot(zy, zy))
        for k = 1 : m
            r[k,j] = _crossdot(zx, zy, lx, lags[k]) / sc
        end
    end
    return r
end

function crosscor!{T<:RealFP}(r::AbstractArray{T,3}, x::Matrix{T}, y::Matrix{T}, lags::IntegerVector; demean::Bool=true)
    lx = size(x, 1)
    nx = size(x, 2)
    ny = size(y, 2)
    m = length(lags)
    (size(y, 1) == lx && size(r) == (m, nx, ny)) || raise_dimerror()
    check_lags(lx, lags)

    # cached (centered) columns of x
    zxs = Array(Vector{T}, 0)
    sizehint(zxs, nx)
    xxs = Array(T, nx)

    for j = 1 : nx
        xj = x[:,j]
        if demean
            mv = mean(xj)
            for i = 1 : lx
                xj[i] -= mv
            end
        end
        push!(zxs, xj)
        xxs[j] = dot(xj, xj)
    end

    zx = Array(T, lx)
    zy = Array(T, lx)
    for j = 1 : ny
        demean_col!(zy, y, j, demean)
        yy = dot(zy, zy)
        for i = 1 : nx
            zx = zxs[i]
            sc = sqrt(xxs[i] * yy)
            for k = 1 : m
                r[k,i,j] = _crossdot(zx, zy, lx, lags[k]) / sc
            end
        end
    end
    return r
end

function crosscor{T<:Real}(x::Vector{T}, y::Vector{T}, lags::IntegerVector; demean::Bool=true)
    crosscor!(Array(fptype(T), length(lags)), float(x), float(y), lags; demean=demean)
end

function crosscor{T<:Real}(x::Matrix{T}, y::Vector{T}, lags::IntegerVector; demean::Bool=true)
    crosscor!(Array(fptype(T), length(lags), size(x,2)), float(x), float(y), lags; demean=demean)
end

function crosscor{T<:Real}(x::Vector{T}, y::Matrix{T}, lags::IntegerVector; demean::Bool=true)
    crosscor!(Array(fptype(T), length(lags), size(y,2)), float(x), float(y), lags; demean=demean)
end

function crosscor{T<:Real}(x::Matrix{T}, y::Matrix{T}, lags::IntegerVector; demean::Bool=true)
    crosscor!(Array(fptype(T), length(lags), size(x,2), size(y,2)), float(x), float(y), lags; demean=demean)
end

crosscor{T<:Real}(x::VecOrMat{T}, y::VecOrMat{T}; demean::Bool=true) = crosscor(x, y, default_crosslags(size(x,1)); demean=demean)

#######################################
#
#   Partial auto-correlations
#
#   TODO: the codes below need cleanup.
#
#######################################

function pacf_regress!{T<:RealFP}(r::RealMatrix, X::AbstractMatrix{T}, lags::IntegerVector, mk::Integer)
    lx = size(X, 1)
    tmpX = ones(T, lx, mk + 1)
    for j = 1 : size(X,2)
        for l = 1 : mk
            for i = 1+l:lx
                tmpX[i,l+1] = X[i-l,j]
            end
        end
        for i = 1 : length(lags)
            l = lags[i]
            sX = sub(tmpX, 1+l:lx, 1:l+1)
            r[i,j] = (cholfact!(sX'sX)\(sX'sub(X, 1+l:lx, j)))[end]
        end
    end
    r
end

function pacf_yulewalker!{T<:RealFP}(r::RealMatrix, X::Matrix{T}, lags::IntegerVector, mk::Integer)
    tmp = Array(T, mk)
    for j = 1 : size(X,2)
        acfs = autocor(X[:,j], 1:mk)
        for i = 1 : length(lags)
            l = lags[i]
            r[i,j] = l == 0 ? one(T) : l == 1 ? acfs[i] : -durbin!(sub(acfs, 1:l), tmp)[l]
        end
    end
end

function pacf!{T<:RealFP}(r::RealMatrix, X::Matrix{T}, lags::IntegerVector; method::Symbol=:regression)
    lx = size(X, 1)
    m = length(lags)
    minlag, maxlag = minmax(lags)
    (0 <= minlag && 2maxlag < lx) || error("Invalid lag value.")
    size(r) == (m, size(X,2)) || raise_dimerror()

    if method == :regression
        pacf_regress!(r, X, lags, maxlag)
    elseif method == :yulewalker
        pacf_yulewalker!(r, X, lags, maxlag)
    else
        error("Invalid method: $method")
    end
    return r
end

function pacf{T<:Real}(X::Matrix{T}, lags::IntegerVector; method::Symbol=:regression)
    pacf!(Array(fptype(T), length(lags), size(X,2)), float(X), lags; method=method)
end

function pacf{T<:Real}(x::Vector{T}, lags::IntegerVector; method::Symbol=:regression)
    vec(pacf(reshape(x, length(x), 1), lags, method=method))
end


# This file contains functions for calculating spatial lags

"""
    slag(W, x)
Calculate the spatial lag of `x` using the spatial weights `W`.
"""
function slag(W::SpatialWeights, x::AbstractVector{T}) where {T <:Real}
    n = W.n
    nx = length(x)

    n == nx || throw(DimensionMismatch("dimensions must match: W has ($(n)), x has ($nx)"))

    sx = zeros(n)
    @inbounds @simd for i in 1:n
        ni = length(W.neighs[i])
        for j in 1:ni
            sx[i] = sx[i] + W.weights[i][j] * x[W.neighs[i][j]]
        end
    end

    return sx

end

Base.:*(W::SpatialWeights, x::AbstractVector{T}) where {T <:Real} = slag(W, x);


begin
    using CSV
    using Revise
    using BNPRegressionGGA2021
    using DataFrames
    using DataStructures
    using Distributions
    using Random
    using RCall
    using Statistics
    using StatsBase
    using LinearAlgebra
    const BNP = BNPRegressionGGA2021
end;

function preprocess(dy, x0, x1)
    # Simulate the responses
    y0 = @. rand(dy(x0))
    y1 = LinRange(-3, 3, 50) |> collect

    # Expand the grid
    grid = rcopy(R"""expand.grid(y1 = $y1, x1 = $x1)""")
    y1 = grid[!, :y1]
    x1 = grid[!, :x1]

    # Expand x0-x1 using splines
    x0min, x0max = extrema(x0)
    R"""
    X0 <- splines::bs($x0, df = 6, Boundary.knots = c($x0min, $x0max))
    """
    @rget X0

    # Standardise responses and predictors
    mX0, sX0 = mean_and_std(X0, 1)
    for col in 1:size(X0, 2)
        if sX0[col] > 0
            X0[:, col] = (X0[:, col] .- mX0[col]) ./ sX0[col]
        end
    end
    my0, sy0 = mean_and_std(y0, 1)
    y0 = (y0 .- my0) ./ sy0

    # Add a constant term to the design matrices
    X0d = rand([0, 1], size(X0, 1), 4)
    X0d = ((X0d .- mean(X0d, dims = 1)) ./ std(X0d, dims = 1))
    X0 = [ones(size(X0, 1)) X0 X0d]
    X1 = zeros(0, size(X0, 2))

    # Generate the true (standardized) density over the grid
    f1 = pdf.(dy.(x1), my0 .+ y1 .* sy0) .* sy0

    # Set the mapping
    mapping = [[1], collect(2:7), [8], [9], [10], [11]]

    # Return the preprocessed results
    return y0, y1, x1, X0, X1, f1, mapping
end;

function simulate_sample(dy, filename)
    N0, N1, Nrep, Niter = 200, 0, 20, 100
    x0raw = repeat(LinRange(-2, 2, N0 ÷ Nrep), Nrep)
    x1raw = LinRange(-1, 1, N1)
    df = map(1:Niter) do iter
        y0, y1, x1, X0, X1, f1, mapping = preprocess(dy, x0raw, x1raw)
        df = DataFrame(X0, :auto)
        df[!, :iter] .= iter
        df[!, :z0] = x0raw
        df[!, :y0] = y0
    end;
    df = reduce(vcat, df)
    CSV.write(filename, df);
end

# Example 1: Continuous predictor (linear regression, normal distribution)
begin
    Random.seed!(1)
    simulate_sample(Normal(x, 1), "data/2022-03-31-data-ex1.csv")
    simulate_sample(SkewNormal(2x, 4, 4), "data/2022-03-31-data-ex2.csv")
end;


export compute_error

function compute_error(u, v, data::DataSet)

    str3 = "donnee_"
    str5 = ".txt"

    epsilon = data.epsilon

    if (epsilon == 10       )  str4 = "10"        end
    if (epsilon == 5        )  str4 = "5"         end
    if (epsilon == 2.5      )  str4 = "2_5"       end
    if (epsilon == 1        )  str4 = "1"         end
    if (epsilon == 0.5      )  str4 = "0_5"       end
    if (epsilon == 0.2      )  str4 = "0_2"       end
    if (epsilon == 0.25     )  str4 = "0_25"      end
    if (epsilon == 0.1      )  str4 = "0_1"       end
    if (epsilon == 0.05     )  str4 = "0_05"      end
    if (epsilon == 0.025    )  str4 = "0_025"     end
    if (epsilon == 0.01     )  str4 = "0_01"      end
    if (epsilon == 0.005    )  str4 = "0_005"     end
    if (epsilon == 0.0025   )  str4 = "0_0025"    end
    if (epsilon == 0.001    )  str4 = "0_001"     end
    if (epsilon == 0.0005   )  str4 = "0_0005"    end
    if (epsilon == 0.00025  )  str4 = "0_00025"   end
    if (epsilon == 0.0001   )  str4 = "0_0001"    end
    if (epsilon == 0.00005  )  str4 = "0_00005"   end
    if (epsilon == 0.000025 )  str4 = "0_000025"  end
    if (epsilon == 0.00001  )  str4 = "0_00001"   end
    if (epsilon == 0.000005 )  str4 = "0_000005"  end
    if (epsilon == 0.0000025)  str4 = "0_0000025" end
    if (epsilon == 0.000001 )  str4 = "0_000001"  end

    ref_file = joinpath("test", "donnees_data3_128_micmac/", str3 * str4 * str5)

    ndata = 128
    uv    = zeros(Float64, (ndata, 4))

    open(ref_file) do f

        for (j,line) in enumerate(eachline(f))
            for (i, val) in enumerate( [ parse(Float64, val) for val in split(line)])
                uv[j, i] = val
            end
        end

    end

    nx   = data.nx
    xmin = data.xmin
    xmax = data.xmax
    T    = data.T
    x    = data.x
    dx   = (xmax - xmin) / nx
    L    = xmax - xmin
    kx   = zeros(Float64, nx)
    kx   = 2π / (xmax - xmin) * vcat(0:nx÷2-1,-nx÷2:-1)

    ua = zeros(ComplexF64, (nx, 4))
    va = fft(uv, 1) / ndata
    k  = zeros(Float64, ndata)
    k .= 2π / L * vcat(0:ndata÷2-1,-ndata÷2:-1)

    for j in 1:ndata
        vv  = transpose(va[j, :])
	    ua .= ua .+ vv .* exp.(1im * k[j] * (x .- xmin))
    end

    uref = ua[:, 1] .+ 1im * ua[:, 2]
    vref = ua[:, 3] .+ 1im * ua[:, 4]

    refH1 = sqrt(dx * norm(ifft(1im * sqrt.(1 .+ kx.^2) .* fft(uref,1),1))^2
               + dx * norm(ifft(1im * sqrt.(1 .+ kx.^2) .* fft(vref,1),1))^2)

    err  = (sqrt(dx * norm(ifft(1im * sqrt.(1 .+ kx.^2) .* fft(u .- uref,1),1))^2
               + dx * norm(ifft(1im * sqrt.(1 .+ kx.^2) .* fft(v .- vref,1),1))^2)) / refH1

    err

end


using ModelingToolkit, OrdinaryDiffEq
using AnalogMTK
using AnalogMTK.EasyPlot #set, cons
import Printf: @sprintf
j=im

#==Circuit parameters
===============================================================================#
𝑓 = 60 #Hz
Tsim = 1 #sec

#@named Vsrc = VSource(1.0)
@named Vsrc = VSource(Sine(A=1, f=𝑓))
@named ZS = Resistor(R=100)
@named RL = Resistor(R=100)
@named CL = Capacitor(C=10e-9) #10e-6, 10e-9, 47e-9

#==Basic calculations & input validation
===============================================================================#
#=
𝜔 = 2π*𝑓
ZC = 1/(j*𝜔*C)

println()
@info("Simulation parameters:")
@show V, 𝑓, 𝜔
@show R, C, abs(ZC)

@show 𝜏 = R*C
=#

#==Circuit definition
===============================================================================#
#@named D1 = Diode(Ideal())
@named D1 = Diode(Shockley(), IS=1e-3)
#@named D1 = Resistor(R=1)
@named GND = Ground()

rc_eqs = [
	connect(Vsrc.p, ZS.p)
	connect(ZS.n, D1.p)
	connect(D1.n, RL.p, CL.p)
	connect(GND.g, Vsrc.n, RL.n, CL.n)
]

systems = [GND, Vsrc, ZS, D1, RL, CL]
@named rc_model = ODESystem(rc_eqs, AnalogMTK.t; systems)

println()
@info("Initial ODESystem (model definition):"); flush(stdout); flush(stderr)
display(rc_model)

println()
@info("Computing simplified model (ODESystem)..."); flush(stdout); flush(stderr)
#sys = rc_model
sys = structural_simplify(rc_model)
display(sys)
println("\tdone.")

#==Circuit simulation
===============================================================================#
println()
@info("Running simulation..."); flush(stdout); flush(stderr)
u0 = [
      ZS.v => 0.0,
      CL.v => 0.0,
      CL.p.i => 0.0,
      D1.v => 0.0,
     ]
prob = ODEProblem(sys, u0, (0, Float64(Tsim)))
sol = solve(prob, Trapezoid())
	#Trapezoid(), ImplicitEuler()
	#Rodas4(), RadauIIA5(), Tsit5()
@show length(sol.t)
println("\tdone.")

#==Plot results
===============================================================================#
println()
@info("Plotting results..."); flush(stdout); flush(stderr)
ylabels = ["Voltage [V]", "Current [A]"]
plot = plot_transient(sol.t, sol[Vsrc.p.v], id="Vsrc"; ylabels)
plot_transient!(plot, sol.t, sol[D1.v], id="D1.v")
plot_transient!(plot, sol.t, sol[CL.p.v], id="Vout")
plot_transient!(plot, sol.t, sol[RL.p.i], id="Iout", strip=2)
plot_transient!(plot, sol.t, sol[D1.i], id="D1.i", strip=2)
plot_transient!(plot, sol.t, sol[Vsrc.i], id="Vsrc.i", strip=2)
set(plot, ystrip1=set(min=-2, max=2))
displaygui(plot)
println("\tdone.")

:TEST_DONE


@testset "README examples" begin

    # Don't forget to update README.md! TODO: Use some real tool to handle this instead...

    @testset "First example" begin
        using NaiveNASlib
        using Test
        in1 = inputvertex("in1", 1)
        in2 = inputvertex("in2", 1)

        # Create a new vertex which computes the sum of in1 and in2
        computation = in1 + in2
        @test typeof(computation) == MutationVertex

        # CompGraph helps evaluating the whole graph as a function
        graph = CompGraph([in1, in2], computation);

        # Evaluate the function represented by graph
        @test graph(2,3) == 5
    end

    @testset "Second and third examples" begin

        # First we need something to mutate. Batteries excluded, remember?
        mutable struct SimpleLayer
            W
            SimpleLayer(W) = new(W)
            SimpleLayer(nin, nout) = new(ones(Int, nin,nout))
        end
        (l::SimpleLayer)(x) = x * l.W

        # Helper function which creates a mutable layer.
        layer(in, outsize) = absorbvertex(SimpleLayer(nout(in), outsize), outsize, in, mutation=IoSize)

        invertex = inputvertex("input", 3)
        layer1 = layer(invertex, 4);
        layer2 = layer(layer1, 5);

        @test [nout(layer1)] == nin(layer2) == [4]

        # Lets change the output size of layer1:
        Δnout(layer1, -2);

        @test [nout(layer1)] == nin(layer2) == [2]

        ### Third example ###
        # First a few "normal" layers
        invertex = inputvertex("input", 6);
        start = layer(invertex, 6);
        split = layer(start, div(nout(invertex) , 3));

        # When multiplying with a scalar, the output size is the same as the input size.
        # This vertex type is said to be SizeInvariant (in lack of better words).
        scalarmult(v, f::Integer) = invariantvertex(x -> x .* f, v)

        # Concatenation is a third type of vertex, called SizeStack where the output size is the sum of input sizes
        joined = conc(scalarmult(split, 2), scalarmult(split,3), scalarmult(split,5), dims=2);
        @test trait(joined) == SizeStack()

        # Elementwise addition is SizeInvariant
        out = start + joined;
        @test trait(out) == SizeInvariant()

        @test [nout(invertex)] == nin(start) == nin(split) == [3 * nout(split)] == [sum(nin(joined))] == [nout(out)] == [6]
        @test [nout(start), nout(joined)] == nin(out) == [6, 6]

        graph = CompGraph(invertex, out)
        @test graph((ones(Int, 1,6))) == [78  78  114  114  186  186]

        # Ok, lets try to change the size of the vertex "out".
        # First we need to realize that we can only change it by integer multiples of 3
        # This is because it is connected to "split" through three paths which require nin==nout
        # Therefore, any size change to nout of "split" will result in 3 times the change of nin of "out".
        # Equivalently, nout of "split" is nin of "out" divided by 3 and nin/nout must be integers.

        # We need this information from the layer. Some layers have other requirements
        NaiveNASlib.minΔnoutfactor(::SimpleLayer) = 1
        NaiveNASlib.minΔninfactor(::SimpleLayer) = 1

        @test minΔnoutfactor(out) == minΔninfactor(out) == 3

        # Next, we need to define how to mutate our SimpleLayer
        NaiveNASlib.mutate_inputs(l::SimpleLayer, newInSize) = l.W = ones(Int, newInSize, size(l.W,2))
        NaiveNASlib.mutate_outputs(l::SimpleLayer, newOutSize) = l.W = ones(Int, size(l.W,1), newOutSize)

        # In some cases it is useful to hold on to the old graph before mutating
        # To do so, we need to define the clone operation for our SimpleLayer
        NaiveNASlib.clone(l::SimpleLayer) = SimpleLayer(l.W)
        parentgraph = copy(graph)

        Δnout(out, 3)

        # We didn't touch the input when mutating...
        @test [nout(invertex)] == nin(start) == [6]
        # Start and joined must have the same size due to elementwise op.
        # All three scalarmult vertices are transparent and propagate the size change to split
        @test [nout(start)] == nin(split) == [3 * nout(split)] == [sum(nin(joined))] == [nout(out)] == [9]
        @test [nout(start), nout(joined)] == nin(out) == [9, 9]

        # However, this only updated the mutation metadata, not the actual layer.
        # Some reasons for this are shown in the pruning example below
        @test graph((ones(Int, 1,6))) == [78  78  114  114  186  186]

        # To mutate the graph, we need to apply the mutation:
        apply_mutation(graph);

        @test graph((ones(Int, 1,6))) == [114  114  114  168  168  168  276  276  276]

        # Copy is still intact
        @test parentgraph((ones(Int, 1,6))) == [78  78  114  114  186  186]

        @testset "Add layers example" begin

            invertex = inputvertex("input", 3)
            layer1 = layer(invertex, 5)
            graph = CompGraph(invertex, layer1)

            @test nv(graph) == 2
            @test graph(ones(Int, 1, 3)) == [3 3 3 3 3]

            # Insert a layer between invertex and layer1
            insert!(invertex, vertex -> layer(vertex, nout(vertex)))

            @test nv(graph) == 3
            @test graph(ones(Int, 1, 3)) == [9 9 9 9 9]
        end

        @testset "Remove layers example" begin
            invertex = inputvertex("input", 3)
            layer1 = layer(invertex, 5)
            layer2 = layer(layer1, 4)
            graph = CompGraph(invertex, layer2)

            @test nv(graph) == 3
            @test graph(ones(Int, 1, 3)) == [15 15 15 15]

            # Remove layer1 and change nin of layer2 from 5 to 3
            # Would perhaps have been better to increase nout of invertex, but it is immutable
            remove!(layer1)
            apply_mutation(graph)

            @test nv(graph) == 2
            @test graph(ones(Int, 1, 3)) == [3 3 3 3]
        end

        @testset "Add edge example" begin
            invertices = inputvertex.(["input1", "input2"], [3, 2])
            layer1 = layer(invertices[1], 4)
            layer2 = layer(invertices[2], 4)
            add = layer1 + layer2
            out = layer(add, 5)
            graph = CompGraph(invertices, out)

            @test nin(add) == [4, 4]
            # Two inputs this time, remember?
            @test graph(ones(Int, 1, 3), ones(Int, 1, 2)) == [20 20 20 20 20]

            # This graph is not interesting enough for there to be a good showcase for adding a new edge.
            # Lets create a new layer which has a different output size just to see how things change
            # The only vertex which support more than one input is add
            layer3 = layer(invertices[2], 6)
            create_edge!(layer3, add)
            apply_mutation(graph)

            # By default, NaiveNASlib will try to increase the size in case of a mismatch
            @test nin(add) == [6, 6, 6]
            @test graph(ones(Int, 1, 3), ones(Int, 1, 2)) == [42 42 42 42 42]
        end

        @testset "Remove edge example" begin
            invertex = inputvertex("input", 4)
            layer1 = layer(invertex, 3)
            layer2 = layer(invertex, 5)
            merged = conc(layer1, layer2, layer1, dims=2)
            out = layer(merged, 3)
            graph = CompGraph(invertex, out)

            @test nin(merged) == [3, 5, 3]
            @test graph(ones(Int, 1, 4)) == [44 44 44]

            remove_edge!(layer1, merged)
            apply_mutation(graph)

            @test nin(merged) == [5, 3]
            @test graph(ones(Int, 1, 4)) == [32 32 32]
        end

        @testset "Pruning example" begin
            # Some mockup 'batteries' for this example

            # First, how to select or add rows or columns to a matrix
            # Negative values in selected indicate rows/cols insertion at that index
            function select_params(W, selected, dim)
                Wsize = collect(size(W))
                indskeep = repeat(Any[Colon()], 2)
                newmap = repeat(Any[Colon()], 2)

                # The selected indices
                indskeep[dim] = filter(ind -> ind > 0, selected)
                # Where they are 'placed', others will be zero
                newmap[dim] = selected .> 0
                Wsize[dim] = length(newmap[dim])

                newmat = zeros(Int64, Wsize...)
                newmat[newmap...] = W[indskeep...]
                return newmat
            end

            NaiveNASlib.mutate_inputs(l::SimpleLayer, selected::Vector{<:Integer}) = l.W = select_params(l.W, selected, 1)
            NaiveNASlib.mutate_outputs(l::SimpleLayer, selected::Vector{<:Integer}) = l.W = select_params(l.W, selected, 2)

            # Return layer just so we can easiliy look at it
            function prunablelayer(in, outsize)
                l = SimpleLayer(reshape(1: nout(in) * outsize, nout(in), :))
                return absorbvertex(l, outsize, in), l
            end

            # Ok, now lets get down to business!
            invertices = inputvertex.(["in1", "in2"], [3,4])
            v1, l1 = prunablelayer(invertices[1], 4)
            v2, l2 = prunablelayer(invertices[2], 3)
            merged = conc(v1, v2, dims=2)
            v3, l3 = prunablelayer(merged, 2)
            graph = CompGraph(invertices, v3)

            # These weights are of course complete nonsense from a neural network perspective.
            # They are just to make it easier to spot what has changed after pruning below.
            @test l1.W ==
            [ 1  4  7  10 ;
            2  5  8  11 ;
            3  6  9  12 ]

            @test l2.W ==
            [ 1  5   9 ;
            2  6  10 ;
            3  7  11 ;
            4  8  12 ]

            @test l3.W ==
            [ 1   8 ;
            2   9 ;
            3  10 ;
            4  11 ;
            5  12 ;
            6  13 ;
            7  14 ]

            # A limitation in current implementation is that one must change the size before pruning
            # See https://github.com/DrChainsaw/NaiveNASlib.jl/issues/40
            Δnin(v3, -3)

            # What did that do?
            @test nout(v1) == 2
            @test nout(v2) == 2

            # Doing this however makes it possible to do several mutations without throwing away
            # more information than needed.
            # For example, if we had first applied the previous mutation we would have thrown away
            # weights for v2 which would then just be replaced by 0s when doing this:
            Δnout(v2, 2)

            # What did that do?
            @test nout(v1) == 2
            @test nout(v2) == 4
            @test nin(v3) == [6]
            # Net result is that v1 shall decrease output size by 1 and v2 shall increase its output size by 1

            # Now, we need a utility/value metric per neuron in order to determine which neurons to keep
            # Give high utility to neurons 1 and 3 of v1, same for all others...
            utility(v) = v == v1 ? [10, 1, 10, 1] : ones(nout_org(v))
            # Then select the neurons.
            Δoutputs(graph, utility)
            # And apply it to the actual weights
            apply_mutation(graph)

            @test l1.W ==
            [ 1  7 ;
            2  8 ;
            3  9 ]

            # Note how column 3 was not replaced by zeros: We increased the target size before pruning
            @test l2.W ==
            [ 1  5   9  0;
            2  6  10  0;
            3  7  11  0;
            4  8  12  0]

            @test l3.W ==
            [ 1   8 ;
            3  10 ;
            5  12 ;
            6  13 ;
            7  14 ;
            0   0]
        end

        @testset "Advanced usage" begin

            @testset "Strategies" begin
                # A simple graph where one vertex has a constraint for changing the size.
                invertex = inputvertex("in", 3)
                layer1 = layer(invertex, 4)
                joined = conc(scalarmult(layer1, 2), scalarmult(layer1, 3), dims=2)

                # joined can only change in steps of 2
                @test minΔnoutfactor(joined) == 2

                # all_in_graph finds all vertices in the same graph as the given vertex
                verts = all_in_graph(joined)

                # Strategy to try to change it by one and throw an error when not successful
                exactOrFail = ΔNout{Exact}(joined, 1, ΔSizeFailError("Size change failed!!"))

                # Note that we now call Δsize instead of Δnout as the direction is given by the strategy
                @test_throws ErrorException Δsize(exactOrFail, verts)

                # No change was made
                @test nout(joined) == 2*nout(layer1) == 8

                # Try to change by one and fail silently when not successful
                exactOrNoop = ΔNout{Exact}(joined, 1, ΔSizeFailNoOp())

                Δsize(exactOrNoop, verts)

                # No change was made
                @test nout(joined) == 2*nout(layer1) == 8

                # In many cases it is ok to not get the exact change which was requested
                relaxedOrFail = ΔNout{Relaxed}(joined, 1, ΔSizeFailError("This should not happen!!"))

                Δsize(relaxedOrFail, verts)

                # Changed by two as this was the smallest possible change
                @test nout(joined) == 2*nout(layer1) == 10

                # Logging when fallback is applied is also possible
                using Logging
                # Yeah, this is not easy on the eyes, but it gets the job done...
                exactOrLogThenRelax = ΔNout{Exact}(joined, 1, LogΔSizeExec(Logging.Info, "Exact failed, relaxing", ΔNout{Relaxed}(joined, 1, ΔSizeFailError("This should not happen!!"))))

                @test_logs (:info, "Exact failed, relaxing") Δsize(exactOrLogThenRelax, verts)

                @test nout(joined) == 2*nout(layer1) == 12
            end

            @testset "Traits" begin
                noname = layer(inputvertex("in", 2), 2)
                @test name(noname) == "MutationVertex::SizeAbsorb"

                # Naming vertices is so useful for logging and debugging I almost made it mandatory
                named = absorbvertex(SimpleLayer(2, 3), 3, inputvertex("in", 2), traitdecoration = t -> NamedTrait(t, "named layer"))
                @test name(named) == "named layer"

                # Speaking of logging...
                layer1 = absorbvertex(SimpleLayer(2, 3), 3, inputvertex("in", 2), traitdecoration = t -> SizeChangeLogger(NamedTrait(t, "layer1")))

                # What info is shown can be controlled by supplying an extra argument to SizeChangeLogger
                nameonly = NameInfoStr()
                layer2 = absorbvertex(SimpleLayer(nout(layer1), 4), 4, layer1, traitdecoration = t -> SizeChangeLogger(nameonly, NamedTrait(t, "layer2")))

                @test_logs(
                (:info, "Change nout of layer1, inputs=[in], outputs=[layer2], nin=[2], nout=[3], SizeAbsorb() by 1"),
                (:info, "Change nin of layer2 by 1"), # Note: less verbose compared to layer1 due to NameInfoStr
                 Δnout(layer1, 1))

                 # traitdecoration works exactly the same for conc and invariantvertex as well, no need for an example

                 # Use the >> operator when creating SizeInvariant vertices using arithmetic operators:
                 add = "addvertex" >> inputvertex("in1", 1) + inputvertex("in2", 1)
                 @test name(add) == "addvertex"

                 # For more elaborate traits one can use traitconf
                 add2 = traitconf(t -> SizeChangeLogger(NamedTrait(t, "layer1 + layer2"))) >> layer1 + layer2
                 @test name(add2) == "layer1 + layer2"

                 @test_logs(
                 (:info, "Change nout of layer1, inputs=[in], outputs=[layer2, layer1 + layer2], nin=[2], nout=[4], SizeAbsorb() by 1"),
                 (:info, "Change nin of layer2 by 1"),
                 (:info, "Change nout of layer2 by 1"),
                 (:info, "Change nin of layer1 + layer2, inputs=[layer1, layer2], outputs=[], nin=[4, 4], nout=[4], SizeInvariant() by 1, 1"),
                 (:info, "Change nout of layer1 + layer2, inputs=[layer1, layer2], outputs=[], nin=[5, 5], nout=[4], SizeInvariant() by 1"),
                  Δnout(layer1, 1))

                  # When creating own trait wrappers, remember to subtype DecoratingTrait or else there will be pain!

                  # Wrong!! Not a subtype of DecoratingTrait
                  struct PainfulTrait{T<:MutationTrait} <: MutationTrait
                      base::T
                  end
                  painlayer = absorbvertex(SimpleLayer(2, 3), 3, inputvertex("in", 2), traitdecoration = PainfulTrait)

                  # Now one must implement a lot of methods for PainfulTrait...
                  @test_throws MethodError Δnout(painlayer, 1)

                  # Right! Is a subtype of DecoratingTrait
                  struct SmoothSailingTrait{T<:MutationTrait} <: DecoratingTrait
                      base::T
                  end
                  # Just implement base and all will be fine
                  NaiveNASlib.base(t::SmoothSailingTrait) = t.base

                  smoothlayer = absorbvertex(SimpleLayer(2, 3), 3, inputvertex("in", 2), traitdecoration = SmoothSailingTrait)

                  Δnout(smoothlayer, 1)
                  @test nout(smoothlayer) == 4
            end

            @testset "Graph instrumentation and modification" begin
                invertex = inputvertex("in", 2)
                layer1 = layer(invertex, 3)
                layer2 = layer(layer1, 4)

                graph = CompGraph(invertex, layer2)

                @test name.(vertices(graph)) == ["in", "MutationVertex::SizeAbsorb", "MutationVertex::SizeAbsorb"]

                # Ok, lets add names to layer1 and layer2 and change the name of invertex

                # Lets first define the default: Fallback to "clone"
                # clone is the existing function to copy things in this manner as I did not want to override Base.copy
                copyfun(args...;cf) = clone(args...;cf=cf) # Keyword argument cf is the function to use for copying all fields of the input

                # Add a name to layer1 and layer2
                function copyfun(v::MutationVertex,args...;cf)
                    # This is probably not practical to do in a real graph, so make sure you have names when first creating it...
                    name = v == layer1 ? "layer1" : "layer2"
                    addname(args...;cf) = clone(args...;cf=cf)
                    addname(t::SizeAbsorb;cf) = NamedTrait(t, name) # SizeAbsorb has no fields, otherwise we would have had to run cf for each one of them...
                    clone(v, args...;cf=addname)
                end

                # Change name of invertex
                # Here we can assume that invertex name is unique in the whole graph or else we would have had to use the above way
                copyfun(s::String; cf) = s == name(invertex) ? "in changed" : s

                # Now supply copyfun when copying the graph.
                # I must admit that thinking about what this does makes me a bit dizzy...
                namedgraph = copy(graph, copyfun)

                @test name.(vertices(namedgraph)) == ["in changed", "layer1", "layer2"]
            end

        end
    end

end


using ModelingToolkit
using Test

@parameters t
@variables x(t) y(t) z(t)

null_op = 0*t
@test isequal(simplify(null_op), 0)

one_op = 1*t
@test isequal(simplify(one_op), t)

identity_op = Num(Term(identity,[x.val]))
@test isequal(simplify(identity_op), x)

minus_op = -x
@test isequal(simplify(minus_op), -1x)
simplify(minus_op)

@variables x

@test toexpr(expand_derivatives(Differential(x)((x-2)^2))) == :($(*)(2, $(+)(-2, x)))
@test toexpr(expand_derivatives(Differential(x)((x-2)^3))) == :($(*)(3, $(^)($(+)(-2, x), 2)))
@test toexpr(simplify(x+2+3)) == :($(+)(5, x))

d1 = Differential(x)((-2 + x)^2)
d2 = Differential(x)(d1)
d3 = Differential(x)(d2)

@test toexpr(expand_derivatives(d3)) == :(0)
@test toexpr(simplify(x^0)) == :(1)

@test ModelingToolkit.substitute(2x + y == 1, Dict(x => 0.0, y => 0.0)) === false
@test ModelingToolkit.substitute(2x + y == 1, Dict(x => 0.0, y => 1.0)) === true


function _approx_ceil(val::Float64, ϵ::Float64)::Float64
    if (val - floor(val)) < ϵ
        return floor(val)
    end
    return ceil(val)
end

function _approx_floor(val::Float64, ϵ::Float64)::Float64
    if (ceil(val) - val) < ϵ
        return ceil(val)
    end
    return floor(val)
end

# Follows Gurobi convention
function _optimality_gap(primal::Float64, dual::Float64)
    primal == dual == 0 && return 0.0
    primal == 0 && return Inf
    return 100 * abs(primal - dual) / abs(primal)
end

function _is_root_node(node::Cerberus.Node)
    return isempty(node.lt_bounds) &&
           isempty(node.gt_bounds) &&
           isempty(node.lt_general_constrs) &&
           isempty(node.gt_general_constrs) &&
           node.depth == 0
end


using NewtonMethods
include("ChebyPoly.jl")

fSol = open("nleq1.dat","w")
fRest = open("nleq1.out","w")
dimMax = 9
for dim = 2:dimMax
    n1 = dim + 1

    # Initialize the options
    opt = OptionsNLEQ(OPT_MODE              => 1,
                      OPT_JACGEN            => 1,
                      OPT_JACFCN            => chebyQuadJac,
                      OPT_MSTOR             => 0,
                      OPT_NOROWSCAL         => 0,
                      OPT_PRINTWARNING      => 1,
                      OPT_PRINTITERATION    => 3,
                      OPT_PRINTSOLUTION     => 2,
                      OPT_PRINTIOWARN       => fRest,
                      OPT_PRINTIOMON        => fRest,
                      OPT_PRINTIOSOL        => fSol,
                      OPT_NITMAX            => 10,
                      OPT_RTOL              => 1e-5)

    x0    = collect(1:dim)./n1
    xScal = zeros(x0)

    retCode = -1

    i = 1

    println("Calling the while loop for solving the Cheby Quad equation of dimension $dim")
    while retCode == -1
        (x0, _, retCode) = nleq1(chebyQuad, x0, xScal, opt)
        write(fRest, @sprintf("Returned from call %4i of NLEQ1\n",i))
		flush(fSol)
	    flush(fRest)
        i += 1
    end
    println("Solution = $x0")

    println("DONE")
end
close(fSol)
close(fRest)


using IterativeSolvers
using Base.Test

@testset "SVD Lanczos" begin

srand(1234567)

#Thick restart methods
@testset "Thick restart with method=$method" for method in (:ritz, :harmonic)
    for T in (Float32, Float64)
        @testset "Diagonal Matrix{$T}" begin
            n = 30
            ns = 5
            tol = 1e-5

            A = full(Diagonal(T[1.0 : n;]))
            q = ones(T, n) / √n
            σ, L, history = svdl(A, nsv=ns, v0=q, tol=tol, reltol=tol, maxiter=n, method=method, vecs=:none, log=true)
            @test isa(history, ConvergenceHistory)
            @test norm(σ - [n : -1.0 : n - 4;]) < 5^2 * 1e-5
            @test_throws ArgumentError svdl(A, nsv=ns, v0=q, tol=tol, reltol=tol, maxiter=n, method=:fakemethod, vecs=:none)

            #Check the singular vectors also
            Σ, L = svdl(A, nsv=ns, v0=q, tol=tol, reltol=tol, maxiter=n, method=method, vecs=:both)

            #The vectors should have the structure
            # [ 0  0 ...  0 ]
            # [    ...      ]
            # [ 0  0 ... ±1 ]
            # [ 0   ...     ]
            # [ 0 ±1 ...  0 ]
            # [±1  0 ...  0 ]
            # and furthermore the signs should be aligned across Σ[:U] and Σ[:V]
            for i = 1 : 5
                Σ[:U][end + 1 - i, i] -= sign(Σ[:U][end + 1 - i, i])
            end
            @test vecnorm(Σ[:U]) < σ[1] * √tol
            for i = 1 : 5
                Σ[:Vt][i, end + 1 - i] -= sign(Σ[:Vt][i, end + 1 - i])
            end
            @test vecnorm(Σ[:U]) < σ[1] * √tol
            @test norm(σ - Σ[:S]) < 2max(tol * ns * σ[1], tol)

            #Issue #55
            let
                σ1, _ = svdl(A, nsv=1, tol=tol, reltol=tol)
                @test abs(σ[1] - σ1[1]) < 2max(tol * σ[1], tol)
            end
        end

        @testset "Rectangular Matrix{$T}" begin
            srand(1)
            m = 300
            n = 200
            k = 5
            l = 10

            A = randn(T, m, n)
            q = randn(T, n) |> x -> x / norm(x)
            σ, L = svdl(A, nsv=k, k=l, v0=q, tol=1e-5, maxiter=30, method=method)
            @test norm(σ - svdvals(A)[1 : k]) < k^2 * 1e-5
        end
    end
end
end #svdl

@testset "BrokenArrowBidiagonal" begin
    B = IterativeSolvers.BrokenArrowBidiagonal([1, 2, 3], [1, 2], Int[])
    @test full(B) == [1 0 1; 0 2 2; 0 0 3]
    @test B[3,3] == 3
    @test B[2,3] == 2
    @test B[3,2] == 0
    @test B[1,3] == 1
    @test size(B) == (3,3)
    @test_throws ArgumentError size(B,3)
    @test_throws BoundsError B[1,5]
end


using OpenQuantumTools, LaTeXStrings
using OrdinaryDiffEq, Plots, StatsBase

# calculate the mean and standard deviation of the mean estimator from a sample
function mean_std(sample)
    m, v = mean_and_std(sample, corrected=true)
    m, v/sqrt(length(sample))
end

# All values calculated below are in angular frequency units
num = 10
bvec = 0.2 * ones(num)
γvec = log_uniform(0.01, 1, num)
fluctuator_ensemble = EnsembleFluctuator(bvec, γvec);

plot(fluctuator_ensemble, :spectrum, 2*π*range(0.01, 4, length=100), xscale=:log10, yscale=:log10, linewidth=2, label="")
xlabel!(L"\omega")
ylabel!(L"S(\omega)")

H = DenseHamiltonian([(s)->1], [σz], unit=:ħ)
coupling = ConstantCouplings([0.5*σz], unit=:ħ)
u0 = PauliVec[1][1]
annealing = Annealing(H, u0, coupling = coupling, bath=fluctuator_ensemble)
tf = 10
# create object for parallel simulation
prob = build_ensembles(annealing, tf, :stochastic)
# we run each trajectory serially for this example
sol = solve(prob, Tsit5(), EnsembleSerial(), trajectories=1000, reltol=1e-6, saveat=range(0,tf,length=100))

t_axis = range(0,tf,length=100)
es = []
err = []
for (i, s) in enumerate(t_axis)
    sample = [abs2(u0'*so[i]) for so in sol]
    pop, pop_std = mean_std(sample)
    push!(es, pop)
    push!(err, 2*pop_std)
end

plot(t_axis, es, ribbon=err, linewidth=2, label="")
xlabel!("ns")
ylabel!("<+>")

cb = InstPulseCallback([0.25, 0.5, 0.75] .* tf, (c, x) -> c .= σx * c)
sol = solve(prob, Tsit5(), EnsembleSerial(), trajectories=1000, reltol=1e-6, saveat=range(0,tf,length=100), callback=cb)

s_axis = range(0,1,length=100)
es = []
err = []
for (i, s) in enumerate(s_axis)
    sample = [abs2(u0'*so[i]) for so in sol]
    pop, pop_std = mean_std(sample)
    push!(es, pop)
    push!(err, 2*pop_std)
end

plot(tf*s_axis, es, ribbon=err, linewidth=2, label="")
xlabel!("ns")
ylabel!("<+>")

using HOQSTTutorials
HOQSTTutorials.tutorial_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])


immutable Logistic <: ContinuousUnivariateDistribution
    location::Real
    scale::Real
    function Logistic(l::Real, s::Real)
    	s > zero(s) || error("scale must be positive")
    	new(float64(l), float64(s))
    end
    Logistic(l::Real) = new(float64(l), 1.0)
    Logistic() = new(0.0, 1.0)
end

@_jl_dist_2p Logistic logis

entropy(d::Logistic) = log(d.scale) + 2.0

insupport(::Logistic, x::Real) = isfinite(x)
insupport(::Type{Logistic}, x::Real) = isfinite(x)

kurtosis(d::Logistic) = 1.2

mean(d::Logistic) = d.location

median(d::Logistic) = d.location

mode(d::Logistic) = d.location
modes(d::Logistic) = [d.location]

skewness(d::Logistic) = 0.0

std(d::Logistic) = pi * d.scale / sqrt(3.0)

var(d::Logistic) = (pi * d.scale)^2 / 3.0


### A Pluto.jl notebook ###
# v0.14.4

using Markdown
using InteractiveUtils

# ╔═╡ 56a604d4-ec32-404f-91bc-86fa24b9bb47
using Pkg, DrWatson

# ╔═╡ 286f9602-6773-4fe0-ac75-36b19d8333f0
begin
	@quickactivate "NumericalMethodsforEngineers"
	using Symbolics, SymbolicUtils, Latexify, LaTeXStrings, ModelingToolkit
	using NumericalMethodsforEngineers, Test
end

# ╔═╡ 987efa36-893d-45cc-88b5-234b80071ab2
using NamedTupleTools

# ╔═╡ 278c20ed-066d-436d-836a-10c34661c3eb
md"
!!! tip
	Output from `versioninfo()` is shown in terminal.
"

# ╔═╡ 8e5d2539-81c6-4396-81ff-a6908b582096
versioninfo()

# ╔═╡ f419c6fc-7b77-4ec6-80c0-8672b064c7f4
begin
	@variables t x y
	D = Differential(t)
end;

# ╔═╡ 303f5542-8e24-4994-a701-2bb664560699
md"
!!! note
	Helper functions such as `LATEXIFY()` and `L()` are defined in pluto-helpers.jl
"

# ╔═╡ 0eb5fb16-765e-4bfc-9300-e7a867037eb5
begin
	a = [x^2+y 0 2x;
		 0 0 2y;
		 y^2 + x 0 0]
	LATEXIFY("a", a)
end

# ╔═╡ 62588b63-213e-44b4-b3a0-eec4b1c930c3
L("a", a)

# ╔═╡ 2292c3af-6ca4-415f-b637-610dbbe42cab
begin
	z = t*t^4 + 3t^5
	D(z) # symbolic representation of derivative(t*t^4 + 15t^4, t)
	L(expand_derivatives(D(z))) # 10t^4
end

# ╔═╡ c899af71-1ae1-4b17-ba3a-d3ba989e49b6
L(z)

# ╔═╡ 5e617b9a-eedc-4a0a-82df-57dad2b91b3e
md"b0 = $(latexify(z))"

# ╔═╡ 340e4b31-4de7-40e3-816e-6e705df711de
begin
	b0 = Symbolics.jacobian([x + x*y, x^2 + y],[x, y])
	md"b0 = $(latexify(b0))"
end

# ╔═╡ 0d87482f-7a1a-495f-94fb-a6a625534801
begin
	b = [
		t^2 + t + t^2  2t + 4t
		x + y + y + 2t  x^2 - x^2 + y^2
	]
	b1 = simplify.(b)
	Text(sprint(show, "text/plain", b1))
end

# ╔═╡ ede004a7-e53c-4102-a4dd-6529dcd8f855
md"b1 = $(latexify(b1))"

# ╔═╡ a9376da1-b58a-4134-ab77-3db8ef42695e
begin
	b2 = simplify.(substitute.(b, (Dict(x => y^2),)))
	md"b2 = $(latexify(b2))"
end

# ╔═╡ 65237ef6-7c92-4687-96cb-766016f5ed73
md"
!!! note
	Element b3[2,1] == 16
"

# ╔═╡ c0df2952-f0f0-414a-92a3-3b1c6c622206
begin
	#2×2 Matrix{Num}:
	# 36.0  24.0
	# 16.0   9.0

	b3 = substitute.(b, (Dict(x => 2.0, y => 3.0, t => 4.0),))
	LATEXIFY("b3", b3)
end

# ╔═╡ e4998080-32ef-4095-bad0-92cc27bf5170
LATEXIFY(b3)

# ╔═╡ 9b7a5d1d-1a3e-40bc-86c2-6cf7cedc57a0
begin
	#2×2 Matrix{Num}:
	#   t + 2(t^2)   6t
	# 2t + 2y  y^2

	b4 = simplify.(substitute.(b, (Dict(x => 2y),)))
	Text(sprint(show, "text/plain", b4))
end

# ╔═╡ 412ef249-a7ac-42ef-b30e-4e77cddcd9a6
begin
	b5 = substitute.(b4, (Dict(x => 2.0, y => 3.0, t => 4.0),))
	Text(sprint(show, "text/plain", b5))
end

# ╔═╡ 1e48e58c-509b-42fd-b7e8-a943ff48d7f9
struct MyStruct
           tally::Int
           team::String
       end

# ╔═╡ 170bf340-b629-4cce-831f-570b1f7f4494
mystruct = MyStruct(5, "hometeam")

# ╔═╡ 7f3c0230-0619-498f-85de-9ab5982f72c7
mynamedtuple = ntfromstruct(mystruct)

# ╔═╡ 93ff9a77-8419-465a-8f7b-b565179364bf
ntstruct = structfromnt(MyStruct, mynamedtuple)

# ╔═╡ 2aeebb4c-379e-4876-9d1e-99e7bc8622f2
mystruct == ntstruct

# ╔═╡ Cell order:
# ╠═56a604d4-ec32-404f-91bc-86fa24b9bb47
# ╠═286f9602-6773-4fe0-ac75-36b19d8333f0
# ╟─278c20ed-066d-436d-836a-10c34661c3eb
# ╠═8e5d2539-81c6-4396-81ff-a6908b582096
# ╠═f419c6fc-7b77-4ec6-80c0-8672b064c7f4
# ╟─303f5542-8e24-4994-a701-2bb664560699
# ╠═0eb5fb16-765e-4bfc-9300-e7a867037eb5
# ╠═62588b63-213e-44b4-b3a0-eec4b1c930c3
# ╠═2292c3af-6ca4-415f-b637-610dbbe42cab
# ╠═c899af71-1ae1-4b17-ba3a-d3ba989e49b6
# ╠═5e617b9a-eedc-4a0a-82df-57dad2b91b3e
# ╠═340e4b31-4de7-40e3-816e-6e705df711de
# ╠═0d87482f-7a1a-495f-94fb-a6a625534801
# ╠═ede004a7-e53c-4102-a4dd-6529dcd8f855
# ╠═a9376da1-b58a-4134-ab77-3db8ef42695e
# ╠═e4998080-32ef-4095-bad0-92cc27bf5170
# ╟─65237ef6-7c92-4687-96cb-766016f5ed73
# ╠═c0df2952-f0f0-414a-92a3-3b1c6c622206
# ╠═9b7a5d1d-1a3e-40bc-86c2-6cf7cedc57a0
# ╠═412ef249-a7ac-42ef-b30e-4e77cddcd9a6
# ╠═987efa36-893d-45cc-88b5-234b80071ab2
# ╠═1e48e58c-509b-42fd-b7e8-a943ff48d7f9
# ╠═170bf340-b629-4cce-831f-570b1f7f4494
# ╠═7f3c0230-0619-498f-85de-9ab5982f72c7
# ╠═93ff9a77-8419-465a-8f7b-b565179364bf
# ╠═2aeebb4c-379e-4876-9d1e-99e7bc8622f2


"""
    KernelSum <: Kernel

Create a sum of kernels. One can also use the operator `+`.

There are various ways in which you create a `KernelSum`:

The simplest way to specify a `KernelSum` would be to use the overloaded `+` operator. This is
equivalent to creating a `KernelSum` by specifying the kernels as the arguments to the constructor.
```jldoctest kernelsum
julia> k1 = SqExponentialKernel(); k2 = LinearKernel(); X = rand(5);

julia> (k = k1 + k2) == KernelSum(k1, k2)
true

julia> kernelmatrix(k1 + k2, X) == kernelmatrix(k1, X) .+ kernelmatrix(k2, X)
true

julia> kernelmatrix(k, X) == kernelmatrix(k1 + k2, X)
true
```

You could also specify a `KernelSum` by providing a `Tuple` or a `Vector` of the
kernels to be summed. We suggest you to use a `Tuple` when you have fewer components
and a `Vector` when dealing with a large number of components.
```jldoctest kernelsum
julia> KernelSum((k1, k2)) == k1 + k2
true

julia> KernelSum([k1, k2]) == KernelSum((k1, k2)) == k1 + k2
true
```
"""
struct KernelSum{Tk} <: Kernel
    kernels::Tk
end

function KernelSum(kernel::Kernel, kernels::Kernel...)
    return KernelSum((kernel, kernels...))
end

@functor KernelSum

Base.length(k::KernelSum) = length(k.kernels)

(κ::KernelSum)(x, y) = sum(k(x, y) for k in κ.kernels)

function kernelmatrix(κ::KernelSum, x::AbstractVector)
    return sum(kernelmatrix(k, x) for k in κ.kernels)
end

function kernelmatrix(κ::KernelSum, x::AbstractVector, y::AbstractVector)
    return sum(kernelmatrix(k, x, y) for k in κ.kernels)
end

function kernelmatrix_diag(κ::KernelSum, x::AbstractVector)
    return sum(kernelmatrix_diag(k, x) for k in κ.kernels)
end

function Base.show(io::IO, κ::KernelSum)
    return printshifted(io, κ, 0)
end

function Base.:(==)(x::KernelSum, y::KernelSum)
    return (
        length(x.kernels) == length(y.kernels) &&
        all(kx == ky for (kx, ky) in zip(x.kernels, y.kernels))
    )
end

function printshifted(io::IO, κ::KernelSum, shift::Int)
    print(io, "Sum of $(length(κ)) kernels:")
    for k in κ.kernels
        print(io, "\n")
        for _ in 1:(shift + 1)
            print(io, "\t")
        end
        printshifted(io, k, shift + 2)
    end
end


using NearestNeighbors, StaticArrays, LinearAlgebra
using Distances: Metric, Cityblock, Euclidean

export broomhead_king
export numericallyapunov
export Cityblock, Euclidean

#####################################################################################
#                    Numerical Lyapunov (from reconstruction)                       #
#####################################################################################
# Everything in this section is based on Ulrich Parlitz [1]

"""
    numericallyapunov(R::Dataset, ks;  refstates, w, distance, ntype)
Return `E = [E(k) for k ∈ ks]`, where `E(k)` is the average logarithmic distance
between states of a [`neighborhood`](@ref)
that are evolved in time for `k` steps (`k` must be integer).

## Keyword Arguments

* `refstates = 1:(length(R) - ks[end])` : Vector of indices
  that notes which
  states of the reconstruction should be used as "reference states", which means
  that the algorithm is applied for all state indices contained in `refstates`.
* `w::Int = 1` : The Theiler window, which determines
  whether points are separated enough in time to be considered separate trajectories
  (see [1] and [`neighborhood`](@ref)).
* `ntype::AbstractNeighborhood = FixedMassNeighborhood(1)` : The method to
  be used when evaluating the neighborhood of each reference state. See
  [`AbstractNeighborhood`](@ref) or [`neighborhood`](@ref) for more info.
* `distance::Metric = Cityblock()` : The distance function used in the
  logarithmic distance of nearby states. The allowed distances are `Cityblock()`
  and `Euclidean()`. See below for more info.

## Description
If the dataset/reconstruction
exhibits exponential divergence of nearby states, then it should clearly hold
```math
E(k) \\approx \\lambda\\Delta t k + E(0)
```
for a *well defined region* in the `k` axis, where ``\\lambda`` is
the approximated
maximum Lyapunov exponent. ``\\Delta t`` is the time between samples in the
original timeseries.
You can use [`linear_region`](@ref) with arguments `(ks .* Δt, E)` to
identify the slope
(= ``\\lambda``)
immediatelly, assuming you
have choosen sufficiently good `ks` such that the linear scaling region is bigger
than the saturated region.

The algorithm used in this function is due to Parlitz [1], which itself
expands upon Kantz [2]. In sort, for
each reference state a neighborhood is evaluated. Then, for each point in this
neighborhood, the logarithmic distance between reference state and neighborhood
state is
calculated as the "time" index `k` increases. The average of the above over
all neighborhood states over all reference states is the returned result.

If the `Metric` is `Euclidean()` then use the Euclidean distance of the
full `D`-dimensional points (distance ``d_E`` in ref. [1]).
If however the `Metric` is `Cityblock()`, calculate
the absolute distance of *only the first elements* of the `m+k` and `n+k` points
of the reconstruction `R` (distance ``d_F`` in ref. [1]).

## References

[1] : Skokos, C. H. *et al.*, *Chaos Detection and Predictability* - Chapter 1
(section 1.3.2), Lecture Notes in Physics **915**, Springer (2016)

[2] : Kantz, H., Phys. Lett. A **185**, pp 77–87 (1994)
"""
function numericallyapunov(R::AbstractDataset{D, T}, ks;
                           refstates = 1:(length(R) - ks[end]),
                           w = 1,
                           distance = Cityblock(),
                           ntype = FixedMassNeighborhood(1)) where {D, T}
    Ek = numericallyapunov(R, ks, refstates, w, distance, ntype)
end

function numericallyapunov(R::AbstractDataset{D, T},
                           ks::AbstractVector{Int},
                           ℜ::AbstractVector{Int},
                           w::Int,
                           distance::Metric,
                           ntype::AbstractNeighborhood) where {D, T}

    # ℜ = \Re<tab> = set of indices that have the points that one finds neighbors.
    # n belongs in ℜ and R[n] is the "reference state".
    # Thus, ℜ contains all the reference states the algorithm will iterate over.
    # ℜ is not estimated. it is given by the user. Most common choice:
    # ℜ = 1:(length(R) - ks[end])

    # ⩅(n) = \Cup<tab> = neighborhood of reference state n
    # which is evaluated for each n and for the given neighborhood ntype

    # Initialize:
    timethres = length(R) - ks[end]
    if maximum(ℜ) > timethres
        erstr = "Maximum index of reference states is > length(R) - ks[end] "
        erstr*= "and the algorithm cannot be performed on it. You have to choose "
        erstr*= "reference state indices of at most up to length(R) - ks[end]."
        throw(ArgumentError(erstr))
    end
    E = zeros(T, length(ks))
    E_n = copy(E); E_m = copy(E)
    data = R.data #tree data
    tree = KDTree(data, Euclidean()) # this creates a copy of `data`
    skippedm = 0; skippedn = 0

    for n in ℜ
        # The ⋓(n) can be evaluated on the spot instead of being pre-calculated
        # for all reference states. (it would take too much memory)
        # Since ⋓[n] doesn't depend on `k` one can then interchange the loops:
        # Instead of k being the outermost loop, it becomes the innermost loop!
        point = data[n]
        ⋓ = neighborhood(point, tree, ntype, n, w)
        for m in ⋓
            # If `m` is nearer to the end of the timeseries than k allows
            # is it completely skipped (and length(⋓) reduced).
            # If m is closer to n than the Theiler window allows, also skip.
            # The Theiler window defaults to τ
            if m > timethres || abs(m - n) <= w
                skippedm += 1
                continue
            end
            for (j, k) in enumerate(ks) #ks should be small (of order 10 to 100 MAX)
                E_m[j] = log(delay_distance(distance, R, m, n, k))
            end
            E_n .+= E_m # no need to reset E_m
        end
        if skippedm >= length(⋓)
            skippedn += 1
            skippedm = 0
            continue # be sure to continue if no valid point!
        end
        E .+= E_n ./ (length(⋓) - skippedm)
        skippedm = 0
        fill!(E_n, zero(T)) #reset distances for n reference state
    end

    if skippedn >= length(ℜ)
        ers = "skippedn ≥ length(R)\n"
        ers*= "Could happen because all the neighbors fall within the Theiler "
        ers*= "window. Fix: increase neighborhood size."
        error(ers)
    end
    E ./= length(ℜ) - skippedn
end

@inline function delay_distance(di::Cityblock, R, m, n, k)
    @inbounds abs(R[m+k][1] - R[n+k][1])
end

@inline function delay_distance(di::Euclidean, R, m, n, k)
    @inbounds norm(R[m+k] - R[n+k])
end

#####################################################################################
#                                  Broomhead-King                                   #
#####################################################################################
"""
    broomhead_king(s::AbstractVector, d::Int) -> U, S, Vtr
Return the Broomhead-King coordinates of a timeseries `s`
by performing `svd` on the so-called trajectory matrix with dimension `d`.

## Description
Broomhead and King coordinates is an approach proposed in [1] that applies the
Karhunen–Loève theorem to delay coordinates embedding with smallest possible delay.

The function performs singular value decomposition
on the `d`-dimensional trajectory matrix ``X`` of ``s``,
```math
X = \\frac{1}{\\sqrt{N}}\\left(
\\begin{array}{cccc}
x_1 & x_2 & \\ldots & x_d \\\\
x_2 & x_3 & \\ldots & x_{d+1}\\\\
\\vdots & \\vdots & \\vdots & \\vdots \\\\
x_{N-d+1} & x_{N-d+2} &\\ldots & x_N
\\end{array}
\\right) = U\\cdot S \\cdot V^{tr}.
```
where ``x := s - \\bar{s}``.
The columns of ``U`` can then be used as a new coordinate system, and by
considering the values of the singular values ``S`` you can decide how many
columns of ``U`` are "important". See the documentation page for example application.

## References
[1] :  D. S. Broomhead, R. Jones and G. P. King, J. Phys. A **20**, 9, pp L563 (1987)
"""
function broomhead_king(x::AbstractArray, d::Int)
    X = trajectory_matrix(x, d)
    F = LinearAlgebra.svd(X)
    return F.U, F.S, F.Vt
end

function trajectory_matrix(x::AbstractArray, d::Int)
    xdash = mean(x)
    N = length(x); sqN = √N
    X = zeros(N-d+1, d)
    for j in 1:d
        for i in 0:(N-d)
            @inbounds X[i+1, j] = (x[j+i] - xdash)/sqN
        end
    end
    return X
end


using Revise,Zygote,LinearAlgebra,Random, ChainRules,BenchmarkTools
Random.seed!(1111)

#value of the function-used for testing purposes

#D(sum(A)) = [matrix of ones]
f_sA(A) = sum(A)
D_sA(A::TM, ::Type{T}=eltype(TM)) where {TM<:AbstractMatrix,T} = ones(T, size(A))

#D(At) = (noop- doesn't affect the input)
f_At(A) = A'
D_sAt(A::TM, ::Type{T}=eltype(TM)) where {TM,T} = D_A(A,T)

#D(AtA): let b=ones(K). Then D(AtA)=2A*bb'
f_AtA(A) = A'*A
function D_AtA(A::TM, ::Type{T}=eltype(TM)) where {TM,T}
  N,K = size(A)
  postmult = fill(T(2.0), K,K)
  #DA = D_sA(A)
  return A*postmult
end

#why doesn't this work?
function ChainRulesCore.rrule(::typeof(f_AtA), A::TM, ::Type{T}) where{TM,T}
  K::Int = size(A,2)
  f_AtA(A), (A)->A*fill(T(2.0), K,K)
end

function C_sAtA(A::TM, ::Type{T}=eltype(TM)) where {TM,T}
  N,K = size(A)
  AtA,pAtA  = rrule(*, A',A)
  sAtA,psAtA = rrule(sum, AtA)

  _, _psAtA = psAtA(T(1))
  _,_,_pAtA2 = pAtA(_psAtA)
  #_,_pAt = pAt(extern(_pAtA1))

  return T(2) .* _pAtA2
end

f_Aty(A,y) = A'*y
D_Aty(A::TM, y::AbstractVector{T}, ::Type{T}=eltype(A)) where {TM,T} =(
  reduce(hcat,[y for i in 1:size(A,2)]))

#D(AtAinv) = -AtAinv * D(AtA) * AtAinv
AtAinv(A) = cholesky(A'*A)\I
function D_AtAinv(A::TM, AtAinv::TM=cholesky(A'*A)\I, ::Type{T}=eltype(TM)) where {TM,T}
  N,K = size(A)
  return T(-1) .* A*AtAinv*fill(T(2.0),K,K)*AtAinv
  #return T(-1) .* DA*AtAinv*DAtA*AtAinv
end

#function rrule(::typeof(AtAinv), A::TM where TM, ::Type{T})
#  AtAinv(A),

f_AtAinvAty(A,y) = cholesky(X'*X)\y
#=D((X'X)^-1*X'y) wrt X:
  =D((X'X)^-1)*X'y+(X'X)^-1*D(X'y)
  =#

#=function D_AtAinvAty(A::TM, y::AbstractVector{T}, ::Type{T}=eltype(A)) where {TM,T}

  #build up least squares
  At, pAt = rrule(transpose, A)
  AtA,pAtA = rrule(*, At,A)
  cholAtA, pcholAtA = rrule(cholesky, AtA)
  Aty, pAty

  AtAinv = cholesky(A'*A)\I
  DAtAinv = D_AtAinv(A)
  Aty =A'*y
  DAty = D_Aty(A,y)

  term1 = DAtAinv*Aty
  term2 = AtAinv*DAty
  return  term1 .+ term2
end=#

function testD(f,D,args...; testname="", erroronfail::Bool = true)
  testval = D(args...)

  ∇(args) = gradient((args)->sum(f(args...)), args)
  verval = ∇(args)[1][1]

  local ε::Float64
  try
    ε =sum(abs.(testval .- verval))
  catch err
    msg = "Test $testname: failed probably due to dimensional mismatch: $err
      #args: $args
      testval: $testval
      verval: $verval"
    error(msg)
  end
  if args==verval
    @assert εrms==0
    @info "Test $testname: passed with a difference of 0"
  elseif testval≈verval
    @info "Test $testname: passed with absolute error of $(ε)"
  else
    msg = "Test $testname: failed with absolute error of $(ε)
      #args: $args
      testval: $testval
      verval: $verval"

    erroronfail && error(msg)
    @warn msg
  end

  return nothing
end

function testderivs(N=3,K=2)
  A= rand(N,K)
  y=rand(N)
  testD(f_sA, D_sA, A, testname="noop")
  testD(f_At, D_At, A, testname="transpose")
  testD(f_AtA, D_AtA, A, testname="AtA")
  testD(f_AtA, C_sAtA, A, testname="CAtA")
  #testD(f_AtA, C2_sAtA, A, testname="CAtA")

  testD(f_Aty, D_Aty, A, y, testname="Aty")

  g(A) = gradient((A)->sum(f_AtA(A)), A)
  @btime $g($A)[1][1]
  @btime C_sAtA($A)
  #@btime D_AtA($A)
  #testD(f_AtAinv, D_AtAinv, A, testname="AtAinv")
  #testD(f_AtAinvAty, D_AtAinvAty, A, y, testname="AtAinvAty")

end

testderivs(3,2)


using ModelingToolkit
using LightGraphs
using DiffEqBase
using Test
using UnPack

# Define some variables
@parameters t L g
@variables x(t) y(t) w(t) z(t) T(t) xˍt(t) yˍt(t) xˍˍt(t) yˍˍt(t)
D = Differential(t)

eqs2 = [D(D(x)) ~ T*x,
        D(D(y)) ~ T*y - g,
        0 ~ x^2 + y^2 - L^2]
pendulum2 = ODESystem(eqs2, t, [x, y, T], [L, g], name=:pendulum)
lowered_sys = ModelingToolkit.ode_order_lowering(pendulum2)

lowered_eqs = [D(xˍt) ~ T*x,
               D(yˍt) ~ T*y - g,
               D(x) ~ xˍt,
               D(y) ~ yˍt,
               0 ~ x^2 + y^2 - L^2,]
@test ODESystem(lowered_eqs, t, [xˍt, yˍt, x, y, T], [L, g]) == lowered_sys
@test isequal(equations(lowered_sys), lowered_eqs)

# Simple pendulum in cartesian coordinates
eqs = [D(x) ~ w,
       D(y) ~ z,
       D(w) ~ T*x,
       D(z) ~ T*y - g,
       0 ~ x^2 + y^2 - L^2]
pendulum = ODESystem(eqs, t, [x, y, w, z, T], [L, g], name=:pendulum)

pendulum = initialize_system_structure(pendulum)
sss = structure(pendulum)
@unpack graph, fullvars, varassoc = sss
@test StructuralTransformations.matching(sss, varassoc .== 0) == map(x -> x == 0 ? StructuralTransformations.UNASSIGNED : x, [0, 0, 0, 0, 1, 2, 3, 4, 0])

sys, assign, eqassoc = StructuralTransformations.pantelides!(pendulum)
sss = structure(sys)
@unpack graph, fullvars, varassoc = sss
scc = StructuralTransformations.find_scc(graph, assign)
@test sort(sort.(scc)) == [
                           [1],
                           [2],
                           [3, 4, 7, 8, 9],
                           [5],
                           [6],
                          ]

@test graph.fadjlist == sort.([
 [5, 3],               # 1
 [6, 4],               # 2
 [7, 9, 1],            # 3
 [8, 9, 2],            # 4
 [2, 1],               # 5
 [2, 1, 6, 5],         # 6
 [5, 3, 10, 7],        # 7
 [6, 4, 11, 8],        # 8
 [2, 1, 6, 5, 11, 10], # 9
])
#                  [1, 2, 3, 4,   5,   6,  7,  8, 9,   10,   11]
#                  [x, y, w, z, xˍt, yˍt, w', z', T, xˍt', yˍt']
@test varassoc  == [5, 6, 7, 8,  10,  11,    0,    0, 0,      0,      0]
#1: D(x) ~ w
#2: D(y) ~ z
#3: D(w) ~ T*x
#4: D(z) ~ T*y - g
#5: 0 ~ x^2 + y^2 - L^2
# ----
#6: D(eq:5) -> 0 ~ 2xx'+ 2yy'
#7: D(eq:1) -> D(D(x)) ~ D(w) -> D(xˍt) ~ D(w) -> D(xˍt) ~ T*x
#8: D(eq:2) -> D(D(y)) ~ D(z) -> D(y_t) ~ T*y - g
#9: D(eq:6) -> 0 ~ 2xx'' + 2x'x' + 2yy'' + 2y'y'
#                 [1, 2, 3, 4, 5, 6, 7, 8, 9]
@test eqassoc == [7, 8, 0, 0, 6, 9, 0, 0, 0]

using ModelingToolkit
@parameters t L g
@variables x(t) y(t) w(t) z(t) T(t) xˍt(t) yˍt(t)
D = Differential(t)
idx1_pendulum = [D(x) ~ w,
                 D(y) ~ z,
                 #0 ~ x^2 + y^2 - L^2,
                 D(w) ~ T*x,
                 D(z) ~ T*y - g,
                 # intermediate 1: 0 ~ 2x*D(x) + 2y*D(y) - 0,
                 # intermediate 2(a): 0 ~ 2x*w + 2y*z - 0, (substitute D(x) and D(y))
                 #0 ~ 2x*w + 2y*z,
                 # D(D(x)) ~ D(w) and substitute the rhs
                 D(xˍt) ~ T*x,
                 # D(D(y)) ~ D(z) and substitute the rhs
                 D(yˍt) ~ T*y - g,
                 # 2x*D(D(x)) + 2*D(x)*D(x) + 2y*D(D(y)) + 2*D(y)*D(y) and
                 # substitute the rhs
                 0 ~ 2x*(T*x) + 2*xˍt*xˍt + 2y*(T*y - g) + 2*yˍt*yˍt]
idx1_pendulum = ODESystem(idx1_pendulum, t, [x, y, w, z, xˍt, yˍt, T], [L, g])
first_order_idx1_pendulum = ode_order_lowering(idx1_pendulum)

using OrdinaryDiffEq
using LinearAlgebra
prob = ODEProblem(ODEFunction(first_order_idx1_pendulum),
        #  [x, y, w, z, xˍt, yˍt, T]
           [1, 0, 0, 0,   0,   0, 0.0],# 0, 0, 0, 0],
           (0, 100.0),
           [1, 9.8],
           mass_matrix=calculate_massmatrix(first_order_idx1_pendulum))
sol = solve(prob, Rodas5());
#plot(sol, vars=(1, 2))

new_sys = dae_index_lowering(ModelingToolkit.ode_order_lowering(pendulum2))

prob_auto = ODEProblem(ODEFunction(new_sys),
        #  [xˍt, yˍt, x, y,   T]
           [0,     0, 1, 0, 0.0],# 0, 0, 0, 0],
           (0, 100.0),
           [1, 9.8],
           mass_matrix=calculate_massmatrix(new_sys))
sol = solve(prob_auto, Rodas5());
#plot(sol, vars=(3, 4))

# Define some variables
@parameters t L g
@variables x(t) y(t) T(t)
D = Differential(t)

eqs2 = [D(D(x)) ~ T*x,
        D(D(y)) ~ T*y - g,
        0 ~ x^2 + y^2 - L^2]
pendulum2 = ODESystem(eqs2, t, [x, y, T], [L, g], name=:pendulum)

# Turn into a first order differential equation system
first_order_sys = ModelingToolkit.ode_order_lowering(pendulum2)

# Perform index reduction to get an Index 1 DAE
new_sys = dae_index_lowering(first_order_sys)

u0 = [
  D(x)    => 0.0,
  D(y)    => 0.0,
  x       => 1.0,
  y       => 0.0,
  T       => 0.0
]

p = [
    L => 1.0,
    g => 9.8
]

prob_auto = ODEProblem(new_sys,u0,(0.0,100.0),p)
sol = solve(prob_auto, Rodas5());
#plot(sol, vars=(D(x), y))

###
### More BLT/SCC tests
###

# Test Tarjan (1972) Fig. 3
g = [
     [2],
     [3,8],
     [4,7],
     [5],
     [3,6],
     Int[],
     [4,6],
     [1,7],
    ]
graph = StructuralTransformations.BipartiteGraph(8, 8)
for (eq, vars) in enumerate(g), var in vars
    add_edge!(graph, eq, var)
end
scc = StructuralTransformations.find_scc(graph)
@test scc == [
              [6],
              [3, 4, 5, 7],
              [1, 2, 8],
             ]


# blocking analysis of single time series
# `blocking_analysis()`
# `mean_and_se()`

"""
    BlockingResult(mean, err, err_err, p_cov, k, blocks)
Result of [`blocking_analysis()`](@ref).

### Fields:
- `mean`: sample mean
- `err`: standard error (estimated standard deviation of the mean)
- `err_err`: estimated uncertainty of `err`
- `p_cov`: estimated pseudo covariance of `mean`, relevant for complex time series
- `k::Int`: k-1 blocking steps were used to uncorrelate time series
- `blocks::Int`: number of uncorrelated values after blocking

Has methods for [`NamedTuple`](@ref), [`val_and_errs`](@ref), [`val`](@ref), [`errs`](@ref),
[`mean_and_se`](@ref), `Measurements.:±`,
`MonteCarloMeasurements.Particles`, and
`Statistics.cov` for `Complex` data.

**Example:**
```jldoctest; setup = :(Random.seed!(1234))
julia> blocking_analysis(smoothen(randn(2^10), 2^5))
BlockingResult{Float64}
  mean = -0.025 ± 0.025
  with uncertainty of ± 0.00311966837382259
  from 32 blocks after 5 transformations (k = 6).
```
"""
struct BlockingResult{T}
    mean::T
    err::Float64
    err_err::Float64
    p_cov::T # pseudo covariance for complex normal distribution
    k::Int
    blocks::Int
end
# constructor from NamedTuple of vectors
function BlockingResult(nt, k)
    T = eltype(nt.mean)
    if k < 0 # blocking failed
        return BlockingResult{T}(nt.mean[1], NaN, NaN, T(NaN), k, 0)
    end
    return BlockingResult{T}(
        nt.mean[1], nt.std_err[k], nt.std_err_err[k], nt.p_cov[k], k, nt.blocks[k]
    )
end

function Base.show(io::IO, r::BlockingResult{T}) where T
    println(io, "BlockingResult{$T}")
    println(io, "  mean = $(Measurements.measurement(r))")
    println(io, "  with uncertainty of ± $(r.err_err)")
    T<:Complex && println(io, "  cov = $(cov(r))")
    if r.k > 0
        println(io,"  from $(r.blocks) blocks after $(r.k-1) transformations (k = $(r.k)).")
    else
        println(io, "  Blocking unsuccessful. k = $(r.k). Try using more time steps!")
    end
end

"""
    cov(r::BlockingResult{<:Complex})
Return the covariance matrix of the multivariate normal distribution approximating
the uncertainty of the blocking result `r` of a complex time series.
See (https://en.wikipedia.org/wiki/Complex_normal_distribution).
"""
function Statistics.cov(r::BlockingResult{<:Complex})
    v_xx = real(r.err^2 + r.p_cov)/2
    v_xy = imag(r.p_cov)/2
    v_yy = real(r.err^2 - r.p_cov)/2
    return [v_xx v_xy; v_xy v_yy]
end

"""
    measurement(r::BlockingResult)
    Measurements.±(r::BlockingResult)
Convert a `BlockingResult` into a `Measurement` for linear error propagation with
[`Measurements`](@ref).

Limitation: Does not account for covariance in complex `BlockingResult`.
Consider using [`MonteCarloMeasurements.Particles(r)`](@ref)!
"""
function Measurements.measurement(r::BlockingResult{<:Real})
    return measurement(r.mean, r.err)
end
function Measurements.measurement(r::BlockingResult{<:Complex})
    Σ = cov(r) # real valued covariance matrix
    cm = complex(measurement(real(r.mean), √Σ[1,1]), measurement(imag(r.mean), √Σ[2,2]))
    return cm
end

"""
    MonteCarloMeasurements.Particles(r::BlockingResult; mc_samples = 2000)
    MonteCarloMeasurements.±(r::BlockingResult)
Convert a `BlockingResult` into a `Particles` object for nonlinear error propagation with
[`MonteCarloMeasurements`](@ref).
"""
function MonteCarloMeasurements.Particles(r::BlockingResult{<:Real}; mc_samples = 2000)
    return Particles(mc_samples, Normal(r.mean, r.err))
end
function MonteCarloMeasurements.Particles(r::BlockingResult{<:Complex}; mc_samples = 2000)
    Σ = cov(r) # real valued covariance matrix
    ps = Particles(mc_samples, MvNormal([real(r.mean), imag(r.mean)], Σ))
    return complex(ps[1], ps[2])
end
MonteCarloMeasurements.:±(r::BlockingResult) = Particles(r)

"""
    blocking_analysis(v::AbstractVector; α = 0.01, corrected = true, skip=0)
    -> BlockingResult(mean, err, err_err, p_cov, k, blocks)
Compute the sample mean `mean` and estimate the standard deviation of the mean
(standard error) `err` of a correlated time series using the blocking algorithm from
Flyvberg and Peterson [JCP (1989)](http://aip.scitation.org/doi/10.1063/1.457480)
and the M test of Jonsson
[PRE (2018)](https://link.aps.org/doi/10.1103/PhysRevE.98.043304) at significance level
``1-α``. `k` is the number of
blocking transformations required to pass the hypothesis test for an uncorrelated time
series and `err_err` the estimated standard error or `err`. Use `skip` to skip the first
`skip` elements in `v`.

If decorrelating the
time series fails according to the M test, `NaN` is returned as the standard error and `-1`
for `k`.
`corrected` controls whether
bias correction for variances is used.

See [`BlockingResult`](@ref), [`shift_estimator`](@ref), [`ratio_of_means`](ref).
"""
function blocking_analysis(v::AbstractVector; α = 0.01, corrected::Bool=true, skip=0)
    T = float(eltype(v))
    v = @view v[skip+1:end]
    if length(v) == 0
        @error "Attempted blocking on an empty vector"
        return BlockingResult(zero(T), NaN, NaN, T(NaN), -1, 0)
    elseif length(v) == 1 # treat like failed M test
        return BlockingResult(T(v[1]), NaN, NaN, T(NaN), -1, 1)
    end
    nt = blocks_with_m(v; corrected)
    k = mtest(nt.mj; α, warn=false)
    return BlockingResult(nt, k)
end

"""
    mean_and_se(v::AbstractVector; α = 0.01, corrected::Bool=true, skip=0) -> mean, err
    mean_and_se(r::BlockingResult) -> mean, err
Return the mean and standard error (as a tuple) of a time series obtained from
[`blocking_analysis`](@ref). See also [`BlockingResult`](@ref).
"""
mean_and_se(r::BlockingResult) = r.mean, r.err
function mean_and_se(v::AbstractVector; kwargs...)
    return mean_and_se(blocking_analysis(v; kwargs...))
end

"""
    blocker(v::Vector) -> new_v::Vector
Reblock the data by successively taking the mean of two adjacent data points to
form a new vector with a half of the `length(v)`. The last data point will be
discarded if `length(v)` is odd.
"""
function blocker(v::AbstractVector{T}) where T
    P = typeof(zero(T)/1)
    new_v = Array{P}(undef,(length(v)÷2))
    for i  in 1:length(v)÷2
        @inbounds new_v[i] = (v[2i-1]+v[2i])/2
    end
    return new_v
end

"""
    mtest(mj::AbstractVector; α = 0.01, warn = true) -> k
    mtest(table; α = 0.01, warn = true) -> k
Hypothesis test for decorrelation of a time series after blocking transformations
with significance level ``1-α`` after Jonson
[PRE (2018)](https://doi.org/10.1103/PhysRevE.98.043304).
`mj` or `table.mj` is expected to be a vector with relevant ``M_j`` values from a blocking analysis as
obtained from [`blocks_with_m()`](@ref).
Returns the row number `k` where the M-test is passed.
If the M-test has failed `mtest()` returns the value `-1` and optionally prints
a warning message.
"""
function mtest(mj::AbstractVector; α = 0.01, warn = true)
    m = reverse(cumsum(reverse(mj)))
    k = 1
    while k <= length(m)-1
       if m[k] < cquantile(Chisq(k), α) # compare to χ^2 distribution at quantile 1-α
           return k
       else
           k += 1
       end
    end
    if warn
        @warn "M test failed, more data needed"
    end
    return -1 # indicating the the M-test has failed
end
mtest(table) = mtest(table.mj)

"""
    blocks_with_m(v; corrected = true) -> (;blocks, mean, std_err, std_err_err, p_cov, mj)
Perform the blocking algorithm from Flyvberg and Peterson
[JCP (1989)](http://aip.scitation.org/doi/10.1063/1.457480).
Returns named tuple with the results from all blocking steps. See [`mtest()`](@ref).
"""
@inline function blocks_with_m(v; corrected::Bool=true)
    T = float(eltype(v))
    R = real(T)
    n_steps = floor(Int,log2(length(v)))

    # initialise arrays to be returned
    blocks = Vector{Int}(undef,n_steps)
    mean_arr = Vector{T}(undef,n_steps)
    std_err = Vector{R}(undef,n_steps)
    std_err_err = Vector{R}(undef,n_steps)
    p_cov = Vector{T}(undef,n_steps)
    mj = Vector{R}(undef,n_steps)

    for i in 1:n_steps
        n = length(v)
        blocks[i] = n
        mean_v = mean(v)
        mean_arr[i] = mean_v
        variance = var(v; corrected, mean=mean_v) # variance
        # sample covariance ŷ(1) Eq. (6) [Jonsson]
        gamma = real(autocovariance(v,1; corrected, mean=mean_v))
        # the M value Eq. (12) [Jonsson]
        mj[i] = n*((n-1)*variance/(n^2)+gamma)^2/(variance^2)
        stderr_v = sqrt(variance/n) # standard error
        std_err[i] = stderr_v
        std_err_err[i] = stderr_v/sqrt(2*(n-1)) # error on standard error Eq. (28) [F&P]
        p_cov[i] = pseudo_cov(v,v; xmean=mean_v, ymean=mean_v, corrected)/n
        v = blocker(v) # re-blocking the dataset
    end
    (length(v)≤ 0 || length(v)>2) && @error "Something went wrong in `blocks_with_m`."
    return (;blocks, mean=mean_arr, std_err, std_err_err, p_cov, mj)
end


#md # Needed packages
using DataFrames
using GraphRecipes
using Plots
using AbstractTrees
gr()

#md # this file holds the needed function to build a desition tree base on
#md # [Iterative Dichotomiser 3 method](http://hunch.net/~coms-4771/quinlan.pdf)
#md # this file documentation is base on [Weave.jl](https://github.com/JunoLab/Weave.jl) and [Literate](https://github.com/fredrikekre/Literate.jl)

#md # # ENTROPY

#md # this function calculate the entropy, at Base 2, of the input data for a given
#md # category.
#md #+
#md # base 2 entropy is good for binary class that has only two values.
#md # for more information see: [Entropy](https://en.wikipedia.org/wiki/Entropy)

#src """
#src    entropy2(data::DataFrame, categorycol::Symbol)
#src
#src this function calculate the category entropy using base 2.
#src this is good for binary class that can have only two values.
#src for more information see: [Entropy](https://en.wikipedia.org/wiki/Entropy)
#src the function output the category entropy
#src """

function entropy2(data::DataFrame, categorycol::Symbol)
    #src # find the total number of samples
    totalnumofsamples = length(data[!,categorycol])
    #src # get the category sets (there should be only 2)
    categorysets = Set(data[!,categorycol])
    #src # define the Dictionary for the Probability of each dataset to have each categoryset
    Pcategory = Dict{String,Real}()
    for categoryset in categorysets
        #src # find the number of samples of the choosen dataset
        Pcategory[string(categoryset)] = count(i->i == categoryset, data[!,categorycol])/totalnumofsamples
    end
    return -sum([v == 0 ? 0 : v*log2(v) for v in values(Pcategory)])
end # function entropy2

#md # # Best feature to split
#md # this function run on all the features in the DataFrame and find the best
#md # feature to split about, base on 'entropy2' function

#src """
#src    best_feature_to_split(data::DataFrame, features::Tuple{Symbol}, categorycol::Symbol)
#src
#src this function run on all the features in the DataFrame and find the best
#src feature to split about, base on 'entropy2' function
#src """
function best_feature_to_split(data::DataFrame, features::Tuple, categorycol::Symbol)
    baseline = entropy2(data,categorycol) # the base line is the entropy of the total data
    feature_entropy =  Dict{String,Real}()
    #src # this for loop run on each feature column in the data and calculate his entropy
    for feature in features
        feature_entropy[string(feature)] = 0.0
        # categorysets = Set(data[!,categorycol])
        featuresets = Set(data[!,feature])
        # for categoryset in categorysets
        # partitioned_data = filter(row -> row[categorycol] == categoryset,data)
        # porportion = length(partitioned_data[!,feature])/length(data[!,feature])
        # feature_entropy[string(feature)] += porportion * entropy2(partitioned_data, feature)
        for featureset in featuresets
            partitioned_data = filter(row -> row[feature] == featureset,data)
            porportion = length(partitioned_data[!,feature])/length(data[!,feature])
            feature_entropy[string(feature)] += porportion * entropy2(partitioned_data, categorycol)
        end
    end
    information_gain = Dict{String,Real}()
    for feature in features
        information_gain[string(feature)] = baseline - feature_entropy[string(feature)]
    end
    @info "information_gain: $information_gain"
    return findmax(information_gain)
end # function best_feature_to_split

#md # Potential leaf node
#md # this function find the most common category

#src """
#src    potential_leaf_node(data::DataFrame, categorycol::Symbol)
#src
#src this function find the most common category
#src """
function potential_leaf_node(data::DataFrame, categorycol::Symbol)
    counter = Dict{String,Real}()
    categorysets = Set(data[!,categorycol])
    for categoryset in categorysets
        counter[string(categoryset)] = count(i->i == categoryset, data[!,categorycol])
    end

    return findmax(counter)
end # function potential_leaf_node

#md # # Create Tree
#md # the most common is define as the category which has larger percentage then the 'persent_threshold'
#md # the default value is set to 100%

#src """
#src    create_tree(data::DataFrame, features::Tuple, categirycol::Symbol, persent_threshold::Int = 100)
#src
#src the most common is define as the category which has larger percentage then the 'persent_threshold'
#src the default value is set to 100%
#src """
function create_tree(data::DataFrame, features::Tuple, categorycol::Symbol, persent_threshold::Int = 100)
    counter , category = potential_leaf_node(data, categorycol)
    counter == length(data[!,categorycol]) && return category
    node = Dict{String,Dict}()
    feature_gain, feature_lable = best_feature_to_split(data, features, categorycol)
    node[feature_lable] = Dict()
    classes = Set(data[!,Symbol(feature_lable)])
    for c in classes
        partitioned_data = data[data[!,Symbol(feature_lable)] .== c,:]
        node[feature_lable][string(c)] = create_tree(partitioned_data, features, categorycol)
    end
    return node
end # function create_tree

# AbstractTrees.children(d::Dict) = [p for p in d]
# AbstractTrees.children(p::Pair) = AbstractTrees.children(p[2])
# function AbstractTrees.printnode(io::IO, p::Pair)
#     str = isempty(AbstractTrees.children(p[2])) ? string(p[1], ": ", p[2]) : string(p[1], ": ")
#     print(io, str)
# end
"""
    plottree(tree)

this function plot the tree
"""
function plottree(tree)
    default(size=(1000, 1000))
    plot(TreePlot(tree), method=:tree, fontsize=10, nodeshape=:rectangle, nodesize = 0.05)
end # function plottree

"""
    classify(tree::Dict, data)

this function classify data using the a tree
"""
function classify(tree::Dict,data::DataFrame)
    root = collect(keys(tree))[1]
    node = tree[root]
    res = nothing
    for k in keys(node)
        if string(data[!,Symbol(root)][1]) == k
            if node[k] isa Dict
                res = classify(node[k],data)
            else
                return node[k]
            end
        end
    end
    return res
end # function classify

#====== exmaples for building desition tree and calssify new data =============#
# desition_tree_data_weather = DataFrame( outlook = ["sunny","sunny","overcast","rain","rain","rain","overcast","sunny","sunny","rain","sunny","overcast","overcast","rain"],
#                                         temperature = ["hot","hot","hot","mild","cool","cool","cool","mild","cool","mild","mild","mild","hot","mild"],
#                                         humidity = ["high","high","high","high","normal","normal","normal","high","normal","normal","normal","high","normal","high"],
#                                         windy = [0,1,0,0,0,1,1,0,0,0,1,1,0,1],
#                                         class = [0,0,1,1,1,0,1,0,1,1,1,1,1,0])
# desition_tree_data_weather_features = (:outlook,:temperature,:humidity,:windy)
# desition_tree_data_weather_categorycol = :class
#
# tree = create_tree(desition_tree_data_weather,desition_tree_data_weather_features,desition_tree_data_weather_categorycol)
#
# plottree(tree)
#
# data = DataFrame( outlook = "rain", temperature = "hot", humidity = "normal", windy = 0)
#
# @show classify(tree,data)
#
# desition_tree_data_box2 = DataFrame(x = [0,-1,1,0,0], y = [0,0,0,-1,1], out = [0,1,1,1,1])
# desition_tree_data_box_features = (:x,:y)
# desition_tree_data_box_categorycol = :out
#
# tree = create_tree(desition_tree_data_box2,desition_tree_data_box_features,desition_tree_data_box_categorycol)
#
# plottree(tree)
#
# data = DataFrame( y = 1, x = 1)
#
# @show classify(tree,data)


using Base: @propagate_inbounds, OneTo

import LinearAlgebra: lmul!, rmul!
import Base: Matrix

@propagate_inbounds is_offdiagonal_small(H::AbstractMatrix{T}, i::Int, tol = eps(real(T))) where {T} =
    abs(H[i+1,i]) ≤ tol*(abs(H[i,i]) + abs(H[i+1,i+1]))

function double_shift_schur!(H::AbstractMatrix{Tv}, from::Int, to::Int,
                             μ::Complex, Q = NotWanted()) where {Tv<:Real}
    m, n = size(H)

    # Compute the nonzero entries of p = (H - μI)(H - μ'I)e₁.
    # Because of the Hessenberg structure we only need H[min:min+2,min:min+1] to form p.
    @inbounds H₁₁ = H[from+0,from+0]
    @inbounds H₂₁ = H[from+1,from+0]

    @inbounds H₁₂ = H[from+0,from+1]
    @inbounds H₂₂ = H[from+1,from+1]
    @inbounds H₃₂ = H[from+2,from+1]

    p₁ = abs2(μ) - 2real(μ) * H₁₁ + H₁₁ * H₁₁ + H₁₂ * H₂₁
    p₂ = -2real(μ) * H₂₁ + H₂₁ * H₁₁ + H₂₂ * H₂₁
    p₃ = H₃₂ * H₂₁

    # Map that column to a mulitiple of e₁ via two Given's rotations
    G₁, nrm = get_rotation(p₁, p₂, p₃, from)

    # Apply the Given's rotations
    lmul!(G₁, view(H, :, from:n))
    rmul!(view(H, OneTo(min(from+3, m)), :), G₁)
    rmul!(Q, G₁)

    # Bulge chasing. First step of the for-loop below looks like:
    #  from           to
    #     ↓           ↓
    #     x x x x x x x     x x x x x x x     x + + + x x x
    # i → x x x x x x x     + + + + + + +     x + + + x x x
    #     x x x x x x x     o + + + + + +       + + + x x x
    #     x x x x x x x  ⇒  o + + + + + +  ⇒    + + + x x x
    #       |   x x x x           x x x x       + + + x x x
    #       |     x x x             x x x             x x x
    #       |       x x               x x               x x
    #       ↑
    #       i
    #
    # Last iterations looks like:
    #  from           to
    #     ↓           ↓
    #     x x x x x x x     x x x x x x x     x x x x + + +
    #     x x x x x x x     x x x x x x x     x x x x + + +
    #       x x x x x x       x x x x x x       x x x + + +
    #         x x x x x  ⇒    x x x x x x  ⇒      x x + + +
    # i → ----- x x x x           + + + +           x + + +
    #           x x x x           o + + +             + + +
    #           x x x x           o + + +             + + +
    #             ↑
    #             i

    @inbounds for i = from + 1 : to - 2
        p₁ = H[i+0,i-1]
        p₂ = H[i+1,i-1]
        p₃ = H[i+2,i-1]

        G, nrm = get_rotation(p₁, p₂, p₃, i)

        # First column is done by hand
        H[i+0,i-1] = nrm
        H[i+1,i-1] = zero(Tv)
        H[i+2,i-1] = zero(Tv)

        # Rotate remaining columns
        lmul!(G, view(H, :, i:n))

        # Create a new bulge
        rmul!(view(H, OneTo(min(i + 3, m)), :), G)
        rmul!(Q, G)
    end

    # Last bulge is just one Given's rotation
    #     from          to
    #       ↓           ↓
    #from → x x x x x x x    x x x x x x x    x x x x x + +
    #       x x x x x x x    x x x x x x x    x x x x x + +
    #         x x x x x x      x x x x x x      x x x x + +
    #           x x x x x  ⇒     x x x x x  ⇒     x x x + +
    #             x x x x          x x x x          x x + +
    #               x x x            + + +            x + +
    # to  → ------- x x x            o + +              + +

    @inbounds Gₙ, nrm = get_rotation(H[to-1,to-2], H[to,to-2], to-1)
    @inbounds H[to-1,to-2] = nrm
    @inbounds H[to,to-2] = zero(Tv)

    lmul!(Gₙ, view(H, :, to-1:n))
    rmul!(view(H, OneTo(to), :), Gₙ)
    rmul!(Q, Gₙ)

    H
end

function single_shift_schur!(H::AbstractMatrix{Tv}, from::Int, to::Int,
                             μ::Number, Q = NotWanted()) where {Tv<:Number}
    m, n = size(H)

    # Compute the nonzero entries of p = (H - μI)e₁.
    @inbounds H₁₁ = H[from+0,from+0]
    @inbounds H₂₁ = H[from+1,from+0]

    p₁ = H₁₁ - μ
    p₂ = H₂₁

    # Map that column to a mulitiple of e₁ via two Given's rotations
    G₁, nrm = get_rotation(p₁, p₂, from)

    # Apply the Given's rotations
    lmul!(G₁, view(H, :, from:n))
    rmul!(view(H, OneTo(min(from + 2, m)), :), G₁)
    rmul!(Q, G₁)

    # Bulge chasing. First step of the for-loop below looks like:
    #  from           to
    #     ↓           ↓
    #     x x x x x x x     x x x x x x x     x + + x x x x
    # i → x x x x x x x     + + + + + + +     x + + x x x x
    #     x x x x x x x     o + + + + + +       + + x x x x
    #         x x x x x  ⇒      x x x x x  ⇒    + + x x x x
    #       |   x x x x           x x x x           x x x x
    #       |     x x x             x x x             x x x
    #       |       x x               x x               x x
    #       ↑
    #       i
    #
    # Last iterations looks like:
    #  from           to
    #     ↓           ↓
    #     x x x x x x x     x x x x x x x     x x x x x + +
    #     x x x x x x x     x x x x x x x     x x x x x + +
    #       x x x x x x       x x x x x x       x x x x + +
    #         x x x x x  ⇒      x x x x x  ⇒      x x x + +
    #           x x x x           x x x x           x x + +
    # i → ------- x x x             + + +             x + +
    #             x x x             0 + +               + +
    #               ↑
    #               i

    @inbounds for i = from + 1 : to - 1
        p₁ = H[i+0,i-1]
        p₂ = H[i+1,i-1]

        G, nrm = get_rotation(p₁, p₂, i)

        # First column is done by hand
        H[i+0,i-1] = nrm
        H[i+1,i-1] = zero(Tv)

        # Rotate remaining columns
        lmul!(G, view(H, :, i:n))

        # Create a new bulge
        rmul!(view(H, OneTo(min(i + 2, m)), :), G)
        rmul!(Q, G)
    end

    H
end

###
### Real arithmetic
###
function local_schurfact!(H::AbstractMatrix{T}, start::Int, to::Int,
                          Q = NotWanted(), tol = eps(T),
                          maxiter = 100*size(H, 1)) where {T<:Real}
    # iteration count
    iter = 0

    @inbounds while true
        iter += 1

        if iter > maxiter
            throw("QR algorithm did not converge")
        end

        # Indexing
        # `to` points to the column where the off-diagonal value was last zero.
        # while `from` points to the smallest index such that there is no small off-diagonal
        # value in columns from:end-1. Sometimes `from` is just 1. Cartoon of a split
        # with from != 1:
        #
        #  + + | | | | + +
        #  + + | | | | + +
        #    o X X X X = =
        #      X X X X = =
        #      . X X X = =
        #      .   X X = =
        #      .     o + +
        #      .     . + +
        #      ^     ^
        #   from   to
        # The X's form the unreduced Hessenberg matrix we are applying QR iterations to,
        # the | and = values get updated by Given's rotations, while the + values remain
        # untouched! The o's are zeros -- or numerically considered zeros.

        # We keep `from` one column past the zero off-diagonal value, so we check whether
        # the `from - 1` column has a small off-diagonal value.
        from = to
        while from > start && !is_offdiagonal_small(H, from - 1, tol)
            from -= 1
        end

        if from == to
            # This just means H[to, to-1] == 0, so one eigenvalue converged at the end
            H[from,from-1] = zero(T)
            to -= 1
        else
            # Now we are sure we can work with a 2×2 block H[to-1:to,to-1:to]
            # We check if this block has a conjugate eigenpair, which might mean we have
            # converged w.r.t. this block if from + 1 == to.
            # Otherwise, if from + 1 < to, we do either a single or double shift, based on
            # whether the H[to-1:to,to-1:to] part has real eigenvalues or a conjugate pair.

            H₁₁, H₁₂ = H[to-1,to-1], H[to-1,to]
            H₂₁, H₂₂ = H[to  ,to-1], H[to  ,to]

            # Scaling to avoid losing precision in the case where we have nearly
            # repeated eigenvalues.
            scale = abs(H₁₁) + abs(H₁₂) + abs(H₂₁) + abs(H₂₂)
            H₁₁ /= scale; H₁₂ /= scale; H₂₁ /= scale; H₂₂ /= scale

            # Trace and discriminant of small eigenvalue problem.
            t = (H₁₁ + H₂₂) / 2
            d = (H₁₁ - t) * (H₂₂ - t) - H₁₂ * H₂₁
            sqrt_discr = sqrt(abs(d))

            # Very important to have a strict comparison here!
            if d < zero(T)
                # Real eigenvalues.
                # Note that if from + 1 == to in this case, then just one additional
                # iteration is necessary, since the Wilkinson shift will do an exact shift.

                # Determine the Wilkinson shift -- the closest eigenvalue of the 2x2 block
                # near H[to,to]

                λ₁ = t + sqrt_discr
                λ₂ = t - sqrt_discr
                λ = abs(H₂₂ - λ₁) < abs(H₂₂ - λ₂) ? λ₁ : λ₂
                λ *= scale

                # Run a bulge chase
                single_shift_schur!(H, from, to, λ, Q)
            else
                # Conjugate pair
                if from + 1 == to
                    # A conjugate pair has converged apparently!
                    if from != 1
                        H[from,from-1] = zero(T)
                    end
                    to -= 2
                else
                    # Otherwise we do a double shift!
                    complex_shift = scale * (t + sqrt_discr * im)
                    double_shift_schur!(H, from, to, complex_shift, Q)
                end
            end
        end

        # Converged!
        to ≤ start && break
    end

    return true
end

###
### Generic implementation
###
function local_schurfact!(H::AbstractMatrix{T}, start::Int, to::Int,
                          Q = NotWanted(), tol = eps(real(T)),
                          maxiter = 100*size(H, 1)) where {T}
    # iteration count
    iter = 0

    @inbounds while true
        iter += 1
        iter > maxiter && return false

        # Indexing, see the real arithmetic version!

        from = to
        while from > start && !is_offdiagonal_small(H, from - 1, tol)
            from -= 1
        end

        if from == to
            # This just means H[to, to-1] == 0, so one eigenvalue converged at the end
            H[from,from-1] = zero(T)
            to -= 1
        else
            # Compute Wilkinson shift
            H₁₁, H₁₂ = H[to-1,to-1], H[to-1,to]
            H₂₁, H₂₂ = H[to  ,to-1], H[to  ,to]
            d = H₁₁ * H₂₂ - H₂₁ * H₁₂
            t = H₁₁ + H₂₂
            sqr = sqrt(t * t - 4d)
            λ₁ = (t + sqr) / 2
            λ₂ = (t - sqr) / 2
            λ = abs(H₂₂ - λ₁) < abs(H₂₂ - λ₂) ? λ₁ : λ₂

            # Run a bulge chase
            single_shift_schur!(H, from, to, λ, Q)
        end

        # Converged!
        to ≤ start && break
    end

    return true
end

local_schurfact!(H::AbstractMatrix{T}, Q = NotWanted(), tol = eps(real(T)), maxiter = 100*size(H, 1)) where {T} =
    local_schurfact!(H, 1, size(H, 2), Q, tol, maxiter)