using BenchmarkTools

list_test_str = split(
"7,4,9,5,11,17,23,2,0,14,21,24,10,16,13,6,15,25,12,22,18,20,8,19,3,26,1

22 13 17 11  0
 8  2 23  4 24
21  9 14 16  7
 6 10  3 18  5
 1 12 20 15 19

 3 15  0  2 22
 9 18 13 17  5
19  8  7 25 23
20 11 10 24  4
14 21 16 12  6

14 21 17 24  4
10 16 15  9 19
18  8 23 26 20
22 11 13  6  5
 2  0 12  3  7",
'\n'
)

list_str = readlines("day4_input.txt")

# Part 1

list = copy(list_test_str)

function part1(list)
    numbers = parse.(Int, split(list[1], ','))
    num_boards = div(length(list) - 1, 6)
    boards = [Matrix{Int}(undef, 5, 5) for _ in 1:num_boards]
    masks = [BitMatrix(zeros(Bool, 5, 5)) for _ in 1:num_boards]

    for k in 1:num_boards
        for i in 1:5
            boards[k][i, :] .= parse.(Int, split(list[2 + 6 * (k - 1) + i]))
        end
    end

    round = 0
    winner = 0
    while winner == 0
        round += 1
        n = numbers[round]
        for k in 1:num_boards
            masks[k] .|= boards[k] .== n
            for i in 1:5
                if all(masks[k][i,:]) || all(masks[k][:,i])
                    winner = k
                end
            end
        end
    end

    return numbers[round] * sum((!).(masks[winner]) .* boards[winner])
end

#end

@info "Part 1"
@info "test:"
@show part1(list_test_str) == 4512
@info "challenge:"
@show part1(list_str) == 10374

# Part 2

function part2(list)
    numbers = parse.(Int, split(list[1], ','))
    num_boards = div(length(list) - 1, 6)
    boards = [Matrix{Int}(undef, 5, 5) for _ in 1:num_boards]
    masks = [BitMatrix(zeros(Bool, 5, 5)) for _ in 1:num_boards]

    for k in 1:num_boards
        for i in 1:5
            boards[k][i, :] .= parse.(Int, split(list[2 + 6 * (k - 1) + i]))
        end
    end

    round = 0
    winners = Int[]
    while length(winners) < num_boards
        round += 1
        n = numbers[round]
        for k in 1:num_boards
            masks[k] .|= boards[k] .== n
            if k ∉ winners
                for i in 1:5
                    if all(masks[k][i,:]) || all(masks[k][:,i])
                        push!(winners, k)
                        break
                    end
                end
            end
        end
    end

    return numbers[round] * sum((!).(masks[winners[end]]) .* boards[winners[end]])
end

@info "Part 2"
@info "test:"
@show part2(list_test_str) == 1924
@info "challenge:"
@show part2(list_str) == 24742

# Splitting up things

function init(list)
    numbers = parse.(Int, split(list[1], ','))
    num_boards = div(length(list) - 1, 6)
    boards = [Matrix{Int}(undef, 5, 5) for _ in 1:num_boards]
    masks = [BitMatrix(zeros(Bool, 5, 5)) for _ in 1:num_boards]

    for k in 1:num_boards
        for i in 1:5
            boards[k][i, :] .= parse.(Int, split(list[2 + 6 * (k - 1) + i]))
        end
    end
    return numbers, boards, masks
end

function part1!(numbers, boards, masks)
    num_boards = length(boards)
    round = 0
    winner = 0
    @inbounds while winner == 0
        round += 1
        n = numbers[round]
        for k in 1:num_boards
            masks[k] .|= boards[k] .== n
            for i in 1:5
                if all(@view(masks[k][i,:])) || all(@view(masks[k][:,i]))
                    winner = k
                end
            end
        end
    end

    return numbers[round] * sum((!).(masks[winner]) .* boards[winner])
end

function part2!(numbers, boards, masks)
    num_boards = length(boards)
    round = 0
    winners = Int[]
    @inbounds while length(winners) < num_boards
        round += 1
        n = numbers[round]
        for k in 1:num_boards
            masks[k] .|= boards[k] .== n
            if k ∉ winners
                for i in 1:5
                    if all(@view(masks[k][i,:])) || all(@view(masks[k][:,i]))
                        push!(winners, k)
                        break
                    end
                end
            end
        end
    end

    return numbers[round] * sum((!).(masks[winners[end]]) .* boards[winners[end]])
end

@info "Part 1 - splitted"
@info "test:"
numbers, boards, masks = init(list_test_str)
@show part1!(numbers, boards, masks) == 4512
@info "challenge:"
numbers, boards, masks = init(list_str)
@show part1!(numbers, boards, masks) == 10374

@info "Part 2 - splitted"
@info "test:"
numbers, boards, masks = init(list_test_str)
@show part2!(numbers, boards, masks) == 1924
@info "challenge:"
numbers, boards, masks = init(list_str)
@show part2!(numbers, boards, masks) == 24742

@info "Benchmarks"
@info "Init"
@btime init($list_str);
@info "Part 1"
@btime part1!($numbers, b, m) setup=(b=boards; m=masks) evals=1
@info "Part 2"
@btime part2!($numbers, b, m) setup=(b=boards; m=masks) evals=1


export draw_xaxis, draw_yaxis, draw_axes, draw_xtick, draw_ytick

"""
`draw_xaxis(x)` draws a thin arrow from the origin to `(x,0)`.

`draw_xaxis(x1,x2)` is equivalent to `draw_xaxis(x1); draw_xaxis(x2)`.

If no x-value is given, draw a pair of arrows from the origin
to values 10% larger than those returned by `xlims()`.
"""
function draw_xaxis(x::Real; opts...)
    draw_vector(x, 0; color = :black, linewidth = 0.5, legend = false, opts...)
end

function draw_xaxis(x1::Real, x2::Real; opts...)
    draw_xaxis(x1; opts...)
    draw_xaxis(x2; opts...)
end

function draw_xaxis(; opts...)
    x1, x2 = xlims()
    draw_xaxis(x1; opts...)
    draw_xaxis(x2; opts...)
end

"""
`draw_yaxis(y)` draws a thin arrow from the origin to `(0,y)`.

`draw_yaxis(y1,y2)` is equivalent to `draw_yaxis(x1); draw_yaxis(x2)`.

If no y-value is given, draw a pair of arrows from the origin
to values 10% larger than those returned by `ylims()`.
"""
function draw_yaxis(y::Real; opts...)
    draw_vector(0, y; color = :black, linewidth = 0.5, opts...)
end

function draw_yaxis(y1::Real, y2::Real; opts...)
    draw_yaxis(y1; opts...)
    draw_yaxis(y2; opts...)
end

function draw_yaxis(; opts...)
    y1, y2 = ylims()
    draw_yaxis(y1; opts...)
    draw_yaxis(y2; opts...)
end

"""
`draw_axes()` invokes `draw_xaxis()` and `draw_yaxis()`.
"""
function draw_axes(; opts...)
    draw_xaxis(; opts...)
    draw_yaxis(; opts...)
end

const _DEFAULT_TICK_LEN = 0.2

"""
`draw_xtick(x::Real,len)` draws a tick mark on the x-axis with total length `len`.
If `len` is omitted, use `SimpleDrawing._DEFAULT_TICK_LEN`.

`draw_xtick(xlist,len)` calls `draw_xtick` for each value in `xlist`.
"""
function draw_xtick(x::Real, len = _DEFAULT_TICK_LEN; opts...)
    draw_segment(x, -len / 2, x, len / 2, color = :black, linewidth = 0.5, opts...)
end

function draw_xtick(xlist, len = _DEFAULT_TICK_LEN; opts...)
    for x in xlist
        draw_xtick(x, len, opts...)
    end
    plot!()
end

"""
`draw_ytick(y::Real,len)` draws a tick mark on the y-axis with total length `len`.
If `len` is omitted, use `SimpleDrawing._DEFAULT_TICK_LEN`.

`draw_xtick(ylist,len)` calls `draw_ytick` for each value in `xyist`.
"""
function draw_ytick(y::Real, len = _DEFAULT_TICK_LEN; opts...)
    draw_segment(-len / 2, y, len / 2, y, color = :black, linewidth = 0.5, opts...)
end

function draw_ytick(ylist, len = _DEFAULT_TICK_LEN; opts...)
    for y in ylist
        draw_ytick(y, len, opts...)
    end
    plot!()
end


@testset "ui and its derivative" begin

    um1 = ui(-1, Float64[], Float64[], Float64[])
    x = randn()
    @test um1(x)==0.0

    @test D(um1)(x)==0.0

    a0 = randn()
    u0 = ui(0, Float64[], Float64[], [a0])
    @test norm(u0(x) - a0*x)<1e-10

    @test norm(D(u0)(x) - ForwardDiff.derivative(u0, x))<1e-10

    @test ui(0, Float64[], Float64[], [1.0])(x)==x

    u1 = ui(1, [-2.0], [1.0] , [3.0,2.0])
    @test abs(u1(x) - (3.0*x + 2.0*rbf(-2.0,1.0)(x)))<1e-10

    u4 = ui(2, zeros(2), ones(2), [1.0,0.0,0.0])
    @test u4.p == 2

    @test u4(2.0) ==2.0

    u4.coeffi[1] = 0.0
    u4.coeffi[2] = 1.0

    @test u4.coeffi == [0.0, 1.0, 0.0]

    @test u4(2.0) == rbf(0.0,1.0)(2.0)

    u4.coeffi[1] = 2.0
    u4.coeffi[3] = 2.0
    u4.ξi[2] = 2.0
    u4.σi[2] = 5.0

    @test abs(u4(x)- (Lin(x)*2 + rbf(0.0, 1.0)(x) +
                2.0*rbf(2.0, 5.0)(x))) <1e-10

    @test abs(D(u4)(x) - ForwardDiff.derivative(u4, x))<1e-10

    u5 = ui(4, randn(4), rand(4), randn(5))

    @test abs(D(u5)(x) - ForwardDiff.derivative(u5, x))<1e-10
end

@testset "uk and its derivative (first and second)" begin
    um1 = uk(-1, Float64[], Float64[],Float64[])
    x = randn()
    @test um1(x) == x
    @test D!(um1, x)==1.0
    @test D(um1)(x)==1.0
    @test H!(um1, x)==0.0
    @test H(um1)(x)==0.0

    @test uk(0, Float64[], Float64[],[5.0;2.0])(x)== 5.0 + 2.0*x

    a = randn(2)
    u0 = uk(0, Float64[], Float64[],a)
    @test norm(D!(u0, x) - ForwardDiff.derivative(u0,x))<1e-10

    u4 = uk(2)

    @test u4.p == 2

    @test u4(x) == 0.0
    #

    u4.coeffk[1] = 3.0

    @test u4.coeffk == [3.0, 0.0, 0.0, 0.0, 0.0]

    @test u4(x) == 3.0

    u4.coeffk[1] = 0.0
    u4.coeffk[2] = 2.0

    @test u4(x) == 2.0ψ₀(0.0,1.0)(x)

    u4.coeffk .= ones(5)

    @test abs(u4(x) -( 1.0 + ψ₀(0.0,1.0)(x) + 2.0*ψj(0.0,1.0)(x) +
                        ψpp1(0.0,1.0)(x))) <1e-10

    u4.coeffk[1] = 6.0
    u4.σk .= [1.0,2.0,3.0,4.0]

    @test abs(u4(x) -( 6.0 + ψ₀(0.0,1.0)(x) + ψj(0.0,2.0)(x) + ψj(0.0,3.0)(x) +
                        ψpp1(0.0,4.0)(x))) <1e-10
    @test norm(D!(u4, x) - ForwardDiff.derivative(u4, x))<1e-10

    u5 = uk(5, randn(7), rand(7), randn(8))

    @test norm(D(u5)(x) - ForwardDiff.derivative(u5, x))<1e-10

    @test norm(H!(u5, x) - ForwardDiff.derivative(D(u5), x))<1e-10
    @test norm(H(u5)(x) - ForwardDiff.derivative(D(u5), x))<1e-10

end


const ReComp = Union{Real,Complex}

struct Dual{T<:ReComp} <: Number
    value::T
    epsilon::T
end
Dual(x::S, y::T) where {S<:ReComp,T<:ReComp} = Dual(promote(x,y)...)
Dual(x::ReComp) = Dual(x, zero(x))
Dual{T}(x::ReComp) where T<:ReComp = Dual{T}(T(x), zero(T))

const ɛ = Dual(false, true)
const imɛ = Dual(Complex(false, false), Complex(false, true))

const Dual128 = Dual{Float64}
const Dual64  = Dual{Float32}
const Dual32  = Dual{Float16}
const DualComplex256 = Dual{ComplexF64}
const DualComplex128 = Dual{ComplexF32}
const DualComplex64  = Dual{ComplexF16}

Base.convert(::Type{Dual{T}}, z::Dual{T}) where {T<:ReComp} = z
Base.convert(::Type{Dual{T}}, z::Dual) where {T<:ReComp} = Dual{T}(convert(T, value(z)), convert(T, epsilon(z)))
Base.convert(::Type{Dual{T}}, x::Number) where {T<:ReComp} = Dual{T}(convert(T, x), convert(T, 0))
Base.convert(::Type{T}, z::Dual) where {T<:ReComp} = (epsilon(z)==0 ? convert(T, value(z)) : throw(InexactError()))

Base.promote_rule(::Type{Dual{T}}, ::Type{Dual{S}}) where {T<:ReComp,S<:ReComp} = Dual{promote_type(T, S)}
Base.promote_rule(::Type{Dual{T}}, ::Type{S}) where {T<:ReComp,S<:ReComp} = Dual{promote_type(T, S)}
Base.promote_rule(::Type{Dual{T}}, ::Type{T}) where {T<:ReComp} = Dual{T}

Base.widen(::Type{Dual{T}}) where {T} = Dual{widen(T)}

value(z::Dual) = z.value
epsilon(z::Dual) = z.epsilon

value(x::Number) = x
epsilon(x::Number) = zero(typeof(x))

dual(x::ReComp, y::ReComp) = Dual(x, y)
dual(x::ReComp) = Dual(x)
dual(z::Dual) = z

const realpart = value
const dualpart = epsilon

Base.isnan(z::Dual) = isnan(value(z))
Base.isinf(z::Dual) = isinf(value(z))
Base.isfinite(z::Dual) = isfinite(value(z))
isdual(x::Dual) = true
isdual(x::Number) = false
Base.eps(z::Dual) = eps(value(z))
Base.eps(::Type{Dual{T}}) where {T} = eps(T)

function dual_show(io::IO, z::Dual{T}, compact::Bool) where T<:Real
    x, y = value(z), epsilon(z)
    if isnan(x) || isfinite(y)
        compact ? show(IOContext(io, :compact=>true), x) : show(io, x)
        if signbit(y)==1 && !isnan(y)
            y = -y
            print(io, compact ? "-" : " - ")
        else
            print(io, compact ? "+" : " + ")
        end
        compact ? show(IOContext(io, :compact=>true), y) : show(io, y)
        printtimes(io, y)
        print(io, "ɛ")
    else
        print(io, "Dual{",T,"}(", x, ",", y, ")")
    end
end

function dual_show(io::IO, z::Dual{T}, compact::Bool) where T<:Complex
    x, y = value(z), epsilon(z)
    xr, xi = reim(x)
    yr, yi = reim(y)
    if isnan(x) || isfinite(y)
        compact ? show(IOContext(io, :compact=>true), x) : show(io, x)
        if signbit(yr)==1 && !isnan(y)
            yr = -yr
            print(io, " - ")
        else
            print(io, " + ")
        end
        if compact
            if signbit(yi)==1 && !isnan(y)
                yi = -yi
                show(IOContext(io, :compact=>true), yr)
                printtimes(io, yr)
                print(io, "ɛ-")
                show(IOContext(io, :compact=>true), yi)
            else
                show(IOContext(io, :compact=>true), yr)
                print(io, "ɛ+")
                show(IOContext(io, :compact=>true), yi)
            end
        else
            if signbit(yi)==1 && !isnan(y)
                yi = -yi
                show(io, yr)
                printtimes(io, yr)
                print(io, "ɛ - ")
                show(io, yi)
            else
                show(io, yr)
                print(io, "ɛ + ")
                show(io, yi)
            end
        end
        printtimes(io, yi)
        print(io, "imɛ")
    else
        print(io, "Dual{",T,"}(", x, ",", y, ")")
    end
end

function dual_show(io::IO, z::Dual{T}, compact::Bool) where T<:Bool
    x, y = value(z), epsilon(z)
    if !value(z) && epsilon(z)
        print(io, "ɛ")
    else
        print(io, "Dual{",T,"}(", x, ",", y, ")")
    end
end

function dual_show(io::IO, z::Dual{Complex{T}}, compact::Bool) where T<:Bool
    x, y = value(z), epsilon(z)
    xr, xi = reim(x)
    yr, yi = reim(y)
    if !xr
        if xi*!yr*!yi
            print(io, "im")
        elseif !xi*yr*!yi
            print(io, "ɛ")
        elseif !xi*!yr*yi
            print(io, "imɛ")
        end
    else
        print(io, "Dual{",T,"}(", x, ",", y, ")")
    end
end

function printtimes(io::IO, x::Real)
    if !(isa(x,Integer) || isa(x,Rational) ||
         isa(x,AbstractFloat) && isfinite(x))
        print(io, "*")
    end
end

Base.show(io::IO, z::Dual) = dual_show(io, z, get(IOContext(io), :compact, false))

function Base.read(s::IO, ::Type{Dual{T}}) where T<:ReComp
    x = read(s, T)
    y = read(s, T)
    Dual{T}(x, y)
end
function Base.write(s::IO, z::Dual)
    write(s, value(z))
    write(s, epsilon(z))
end

## Generic functions of dual numbers ##

Base.convert(::Type{Dual}, z::Dual) = z
Base.convert(::Type{Dual}, x::Number) = Dual(x)

Base.:(==)(z::Dual, w::Dual) = value(z) == value(w)
Base.:(==)(z::Dual, x::Number) = value(z) == x
Base.:(==)(x::Number, z::Dual) = value(z) == x

Base.isequal(z::Dual, w::Dual) = isequal(value(z),value(w)) && isequal(epsilon(z), epsilon(w))
Base.isequal(z::Dual, x::Number) = isequal(value(z), x) && isequal(epsilon(z), zero(x))
Base.isequal(x::Number, z::Dual) = isequal(z, x)

Base.isless(z::Dual{<:Real},w::Dual{<:Real}) = value(z) < value(w)
Base.isless(z::Real,w::Dual{<:Real}) = z < value(w)
Base.isless(z::Dual{<:Real},w::Real) = value(z) < w

Base.hash(z::Dual) = (x = hash(value(z)); epsilon(z)==0 ? x : bitmix(x,hash(epsilon(z))))

Base.float(z::Union{Dual{T}, Dual{Complex{T}}}) where {T<:AbstractFloat} = z
Base.complex(z::Dual{<:Complex}) = z

Base.floor(z::Dual) = floor(value(z))
Base.ceil(z::Dual)  = ceil(value(z))
Base.trunc(z::Dual) = trunc(value(z))
Base.round(z::Dual) = round(value(z))
Base.floor(::Type{T}, z::Dual) where {T<:Real} = floor(T, value(z))
Base.ceil( ::Type{T}, z::Dual) where {T<:Real} = ceil( T, value(z))
Base.trunc(::Type{T}, z::Dual) where {T<:Real} = trunc(T, value(z))
Base.round(::Type{T}, z::Dual) where {T<:Real} = round(T, value(z))

for op in (:real, :imag, :conj, :float, :complex)
    @eval Base.$op(z::Dual) = Dual($op(value(z)), $op(epsilon(z)))
end

Base.abs(z::Dual) = sqrt(abs2(z))
Base.abs2(z::Dual) = real(conj(z)*z)

Base.real(z::Dual{<:Real}) = z
Base.abs(z::Dual{<:Real}) = z ≥ 0 ? z : -z

Base.angle(z::Dual{<:Real}) = z ≥ 0 ? zero(z) : one(z)*π
function Base.angle(z::Dual{Complex{T}}) where T<:Real
    if z == 0
        if imag(epsilon(z)) == 0
            Dual(zero(T), zero(T))
        else
            Dual(zero(T), convert(T, Inf))
        end
    else
        real(log(sign(z)) / im)
    end
end

Base.flipsign(x::Dual,y::Dual) = y == 0 ? flipsign(x, epsilon(y)) : flipsign(x, value(y))
Base.flipsign(x, y::Dual) = y == 0 ? flipsign(x, epsilon(y)) : flipsign(x, value(y))
Base.flipsign(x::Dual, y) = dual(flipsign(value(x), y), flipsign(epsilon(x), y))

# algebraic definitions
conjdual(z::Dual) = Dual(value(z),-epsilon(z))
absdual(z::Dual) = abs(value(z))
abs2dual(z::Dual) = abs2(value(z))

# algebra

Base.:+(z::Dual, w::Dual) = Dual(value(z)+value(w), epsilon(z)+epsilon(w))
Base.:+(z::Number, w::Dual) = Dual(z+value(w), epsilon(w))
Base.:+(z::Dual, w::Number) = Dual(value(z)+w, epsilon(z))

Base.:-(z::Dual) = Dual(-value(z), -epsilon(z))
Base.:-(z::Dual, w::Dual) = Dual(value(z)-value(w), epsilon(z)-epsilon(w))
Base.:-(z::Number, w::Dual) = Dual(z-value(w), -epsilon(w))
Base.:-(z::Dual, w::Number) = Dual(value(z)-w, epsilon(z))

# avoid ambiguous definition with Bool*Number
Base.:*(x::Bool, z::Dual) = ifelse(x, z, ifelse(signbit(real(value(z)))==0, zero(z), -zero(z)))
Base.:*(x::Dual, z::Bool) = z*x

Base.:*(z::Dual, w::Dual) = Dual(value(z)*value(w), epsilon(z)*value(w)+value(z)*epsilon(w))
Base.:*(x::Number, z::Dual) = Dual(x*value(z), x*epsilon(z))
Base.:*(z::Dual, x::Number) = Dual(x*value(z), x*epsilon(z))

Base.:/(z::Dual, w::Dual) = Dual(value(z)/value(w), (epsilon(z)*value(w)-value(z)*epsilon(w))/(value(w)*value(w)))
Base.:/(z::Number, w::Dual) = Dual(z/value(w), -z*epsilon(w)/value(w)^2)
Base.:/(z::Dual, x::Number) = Dual(value(z)/x, epsilon(z)/x)

for f in [:(Base.:^), :(NaNMath.pow)]
    @eval function ($f)(z::Dual{T1}, w::Dual{T2}) where {T1, T2}
        T = promote_type(T1, T2) # for type stability in ? : statements
        val = $f(value(z), value(w))

        ezvw = epsilon(z) * value(w) # for using in ? : statement
        du1 = iszero(ezvw) ? zero(T) : ezvw * $f(value(z), value(w) - 1)
        ew = epsilon(w) # for using in ? : statement
        # the float is for type stability because log promotes to floats
        du2 = iszero(ew) ? zero(float(T)) : ew * val * log(value(z))
        du = du1 + du2

        Dual(val, du)
    end
end

Base.mod(z::Dual, n::Number) = Dual(mod(value(z), n), epsilon(z))

# introduce a boolean !iszero(n) for hard zero behaviour to combat NaNs
function pow(z::Dual, n::AbstractFloat)
    return Dual(value(z)^n, !iszero(n) * (epsilon(z) * n * value(z)^(n - 1)))
end
function pow(z::Dual{T}, n::Integer) where T
    iszero(n) && return Dual(one(T), zero(T)) # avoid DomainError Int^(negative Int)
    isone(z) && return Dual(one(T), epsilon(z) * n)
    return Dual(value(z)^n, epsilon(z) * n * value(z)^(n - 1))
end
# these first two definitions are needed to fix ambiguity warnings
for T1 ∈ (:Integer, :Rational, :Number)
    @eval Base.:^(z::Dual{T}, n::$T1) where T = pow(z, n)
end

NaNMath.pow(z::Dual{T}, n::Number) where T = Dual(NaNMath.pow(value(z),n), epsilon(z)*n*NaNMath.pow(value(z),n-1))
NaNMath.pow(z::Number, w::Dual{T}) where T = Dual(NaNMath.pow(z,value(w)), epsilon(w)*NaNMath.pow(z,value(w))*log(z))

Base.inv(z::Dual) = dual(inv(value(z)),-epsilon(z)/value(z)^2)

# force use of NaNMath functions in derivative calculations
function to_nanmath(x::Expr)
    if x.head == :call
        funsym = Expr(:.,:NaNMath,Base.Meta.quot(x.args[1]))
        return Expr(:call,funsym,[to_nanmath(z) for z in x.args[2:end]]...)
    else
        return Expr(:call,[to_nanmath(z) for z in x.args]...)
    end
end
to_nanmath(x) = x

for (funsym, exp) in Calculus.symbolic_derivatives_1arg()
    funsym == :exp && continue
    funsym == :abs2 && continue
    funsym == :inv && continue
    if isdefined(SpecialFunctions, funsym)
        @eval function SpecialFunctions.$(funsym)(z::Dual)
            x = value(z)
            xp = epsilon(z)
            Dual($(funsym)(x),xp*$exp)
        end
    elseif isdefined(Base, funsym)
        @eval function Base.$(funsym)(z::Dual)
            x = value(z)
            xp = epsilon(z)
            Dual($(funsym)(x),xp*$exp)
        end
    end
    # extend corresponding NaNMath methods
    if funsym in (:sin, :cos, :tan, :asin, :acos, :acosh, :atanh, :log, :log2, :log10,
          :lgamma, :log1p)
        funsym = Expr(:.,:NaNMath,Base.Meta.quot(funsym))
        @eval function $(funsym)(z::Dual)
            x = value(z)
            xp = epsilon(z)
            Dual($(funsym)(x),xp*$(to_nanmath(exp)))
        end
    end
end

# only need to compute exp/cis once
Base.exp(z::Dual) = (expval = exp(value(z)); Dual(expval, epsilon(z)*expval))
Base.cis(z::Dual) = (cisval = cis(value(z)); Dual(cisval, im*epsilon(z)*cisval))

Base.exp10(x::Dual) = (y = exp10(value(x)); Dual(y, y * log(10) * epsilon(x)))

## TODO: should be generated in Calculus
Base.sinpi(z::Dual) = Dual(sinpi(value(z)),epsilon(z)*cospi(value(z))*π)
Base.cospi(z::Dual) = Dual(cospi(value(z)),-epsilon(z)*sinpi(value(z))*π)

Base.checkindex(::Type{Bool}, inds::AbstractUnitRange, i::Dual) = checkindex(Bool, inds, value(i))


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

using .EKI
using Plots
using Plots.PlotMeasures
using LinearAlgebra
using JLD2
using Statistics
#using LaTeXStrings

using .EKI

# utility functions
function get_marginal_sd(param_sample)
    # n_samples, n_params = size(param_sample)
    param_var=var(param_sample,dims=1)
    param_sd=sqrt.(param_var)

    return param_sd

end
function compute_error_new(self::EKIObj,it)
    meang=dropdims(mean(self.g[it],dims=1),dims=1)
    diff = self.g_t - meang
    err = diff' * inv(self.cov)* diff #as the mean of g has a reduced mean due to sample size
    return err
end

"""
Plotting script for the ensemble, It plots the following traced over iterations of EKP:
- The residual of mean difference to the truth data weighted by the covariance (as stored in the EKP)
- The standard deviations of each parameter in the ensemble

I assume 2 parameters for now, but easily extendable to more

requires:
---------
- EKP object
- the min and make of EKP iterations to plot over

"""

######################################################################
function main()

    #object
    @load ekpdir*"ekp.jld2" ekpobj

    #plot from min to max iteration of ekp
    min_ekp_it=1
    max_ekp_it=10

    u=ekpobj.u[min_ekp_it:max_ekp_it]#it x [enssize x param]

    #get sd
    u_sd = get_marginal_sd.(u) #u_sd[i] = [RH_sd,tau_sd]
    u_sd = cat(u_sd...,dims=1) #make into matrix with num_param columns

    #get error
    g_residual = map(x->compute_error_new(ekpobj,x), min_ekp_it:max_ekp_it)

    iterations = 0:size(u_sd)[1]-1 #start at 0 for "ICs"
    #plot - SD of parameters
    gr(size=(500,500))
    Plots.scalefontsizes(1.25)

    plot(iterations,
         u_sd[:,1],
         label="\\theta_{RH}",
         seriestype=:scatter,
         markershape=:circle,
         markercolor=:blue,
         markersize=8,
         msw=0,
         dpi=300,
         legend=:topright,
         framestyle=:box,
         grid=false,
         left_margin=50px,
         bottom_margin=50px)

    plot!(iterations,
          u_sd[:,2],
          seriestype=:scatter,
          markershape=:cross,
          markercolor=:orange,
          markersize=6,
          msw=0,
          label="\\theta_\\tau")

    xlabel!("Iteration")
    ylabel!("Ensemble Standard Deviation")
    savefig(outdir*"sd_ekp.pdf")

    #plot - Residual

    plot(iterations,
         g_residual,
#         yaxis=:log10,
#         ylims=(1,100),
         line = :scatter,
         markershape=:circle,
         markercolor=:grey,
         markersize=8,
         msw=0,
         dpi=300,
         legend=false,
         grid=false,
         framestyle=:box,
         left_margin=75px,
         bottom_margin=50px)
    xlabel!("Iteration")
    ylabel!("Residual")
    savefig(outdir*"residual_ekp.pdf")

end

main()


function generic_tests_ccmatrix(C)
    r = length(C)
    for i = 1:r
        ccm = SymbolicWedderburn.CMMatrix(C, i)
        l = length(C[i])
        @test all([sum(ccm[:,j]) == l for j = 1:r])
    end
    ccm1 = SymbolicWedderburn.CMMatrix(C, 1)
    @test all([isone(ccm1[i,i]) for i = 1:r])
end

@testset "ConjClassMatrix" begin
    @testset "example: Symmetric group (4)" begin
        G = PermutationGroups.SymmetricGroup(4)
        C = conjugacy_classes(G)
        generic_tests_ccmatrix(C)
        @test SymbolicWedderburn.CMMatrix(C, 2)  == [0 1 0 0 0;
                                                     6 0 3 0 2;
                                                     0 4 0 4 0;
                                                     0 0 3 0 4;
                                                     0 1 0 2 0]
        @test SymbolicWedderburn.CMMatrix(C, 3)  == [0 0 1 0 0;
                                                     0 4 0 4 0;
                                                     8 0 4 0 8;
                                                     0 4 0 4 0;
                                                     0 0 3 0 0]
        @test SymbolicWedderburn.CMMatrix(C, 4)  == [0 0 0 1 0;
                                                     0 0 3 0 4;
                                                     0 4 0 4 0;
                                                     6 0 3 0 2;
                                                     0 2 0 1 0]
        @test SymbolicWedderburn.CMMatrix(C, 5)  == [0 0 0 0 1;
                                                     0 1 0 2 0;
                                                     0 0 3 0 0;
                                                     0 2 0 1 0;
                                                     3 0 0 0 2]
    end
    @testset "example: Alternating group (4)" begin
        a = perm"(1,2,3)(4)"
        b = perm"(1,2,4)"
        G = PermGroup([a,b])
        C = [
            Orbit([a,b], one(G)),
            Orbit([a,b], perm"(1,2)(3,4)"),
            Orbit([a,b], perm"(1,2,3)(4)"),
            Orbit([a,b], perm"(1,3,2)(4)"),
            ]

        generic_tests_ccmatrix(C)
        @test SymbolicWedderburn.CMMatrix(C, 1) == [1 0 0 0
                                                    0 1 0 0
                                                    0 0 1 0
                                                    0 0 0 1]
        @test SymbolicWedderburn.CMMatrix(C, 2) == [0 1 0 0
                                                    3 2 0 0
                                                    0 0 3 0
                                                    0 0 0 3]
        @test SymbolicWedderburn.CMMatrix(C, 3) == [0 0 1 0
                                                    0 0 3 0
                                                    0 0 0 4
                                                    4 4 0 0]
        @test SymbolicWedderburn.CMMatrix(C, 4) == [0 0 0 1
                                                    0 0 0 3
                                                    4 4 0 0
                                                    0 0 4 0]
    end

    @testset "example: C₅ ⋊ C₄" begin
        S = [perm"(1,2,4,5,3)", perm"(2,5,3,4)"]
        G = PermGroup(S);

        ccG = [
            Orbit(S, one(G)),
            Orbit(S, perm"(2,3)(4,5)"),
            Orbit(S, perm"(2,4,3,5)"),
            Orbit(S, perm"(2,5,3,4)"),
            Orbit(S, perm"(1,2,4,5,3)"),
            ]
        # the order of cclasses is taken from GAP

        @assert sum(length, ccG) == order(G)

        @test SymbolicWedderburn.CMMatrix(ccG, 1) ==
        [ 1  0  0  0  0
          0  1  0  0  0
          0  0  1  0  0
          0  0  0  1  0
          0  0  0  0  1 ]
        @test SymbolicWedderburn.CMMatrix(ccG, 2) ==
        [ 0  1  0  0  0
          5  0  0  0  5
          0  0  0  5  0
          0  0  5  0  0
          0  4  0  0  0 ]
        @test SymbolicWedderburn.CMMatrix(ccG, 3) ==
        [ 0  0  1  0  0
          0  0  0  5  0
          0  5  0  0  0
          5  0  0  0  5
          0  0  4  0  0 ]
        @test SymbolicWedderburn.CMMatrix(ccG, 4) ==
        [ 0  0  0  1  0
          0  0  5  0  0
          5  0  0  0  5
          0  5  0  0  0
          0  0  0  4  0 ]
        @test SymbolicWedderburn.CMMatrix(ccG, 5) ==
        [ 0  0  0  0  1
          0  4  0  0  0
          0  0  4  0  0
          0  0  0  4  0
          4  0  0  0  3 ]

        generic_tests_ccmatrix(ccG)
        generic_tests_ccmatrix(conjugacy_classes(G)) # might be in different order
    end

    @testset "random subgroups of SymetricGroup(N)" begin
        for i in 2:6
            G = PermutationGroups.SymmetricGroup(i)
            for _ in 1:5
                PG = PermGroup(rand(G, 2))
                generic_tests_ccmatrix(SymbolicWedderburn.conjugacy_classes(PG))
            end
        end
    end
end


#############################################################################
# intersect.jl
# Handles intersections of combinatorial sets.
#############################################################################

import Base: intersect

export intersect, get_elements

mutable struct IntersectAtom <: CombiSet
  head::Symbol
  id_hash::UInt64
  children::Tuple
  elements::ValOrNothing
  baseset::AbstractArray
  cardinality::Int

  function IntersectAtom(elements::ValOrNothing, children::Tuple, baseset::AbstractArray)
    this = new(:intersect, 0, children, elements, baseset, length(baseset))
    this.id_hash = object_id(this)
    return this
  end
end

function intersect(set::AllCombiSet...)
  baseset = Set([])
  sets = collect(set)
  n = length(sets)
  if n < 2
    error("Cannot compute the union of one set.")
  else
    if isa(sets, Array{Number})
      elements = Sets[i]
      for i = 1:n
        elements = intersect(elements, Set(sets[i]))
      end
      elements = collect(Set(elements))
      return elements
    else
      for i = 1:n
        if isa(sets[i], CombiSet)
          baseset = union(baseset, Set(sets[i].baseset))
        end
      end
      baseset = collect(baseset)
      newset = IntersectAtom(nothing, set, baseset)
      return newset
    end
  end
end

function get_elements(x::IntersectAtom)
  n = length(x.children)
  if isa(x.children[1], CombiSet)
    elements = Set(x.children[1].elements)
  else
    elements = x.children[1]
  end
  for i = 2:n
    if isa(x.children[i], CombiSet)
      elements = intersect(elements, Set(get_elements(x.children[i])))
    else
      seti = Set(x.children[i])
      elements = intersect(elements, seti)
    end
  end
  x.elements = collect(elements)
  return x.elements
end


# # Gaussian Process Classification
#
# ## Preliminary steps
#
# ### Loading necessary packages

using Plots
using HTTP, CSV
using DataFrames: DataFrame
using AugmentedGaussianProcesses

# ### Loading the banana dataset from OpenML
data = HTTP.get("https://www.openml.org/data/get_csv/1586217/phpwRjVjk")
data = CSV.read(data.body, DataFrame)
data.Class[data.Class .== 2] .= -1
data = Matrix(data)
X = data[:, 1:2]
Y = data[:, end];

# ### We create a function to visualize the data

function plot_data(X, Y; size=(300, 500))
    return Plots.scatter(
        eachcol(X)...; group=Y, alpha=0.2, markerstrokewidth=0.0, lab="", size=size
    )
end
plot_data(X, Y; size=(500, 500))

# ### Run sparse classification with increasing number of inducing points
Ms = [4, 8, 16, 32, 64]
models = Vector{AbstractGP}(undef, length(Ms) + 1)
kernel = SqExponentialKernel() ∘ ScaleTransform(1.0)
for (i, num_inducing) in enumerate(Ms)
    @info "Training with $(num_inducing) points"
    m = SVGP(
        X,
        Y,
        kernel,
        LogisticLikelihood(),
        AnalyticVI(),
        num_inducing;
        optimiser=false,
        Zoptimiser=false,
    )
    @time train!(m, 20)
    models[i] = m
end
# ### Running the full model
@info "Running full model"
mfull = VGP(X, Y, kernel, LogisticLikelihood(), AnalyticVI(); optimiser=false)
@time train!(mfull, 5)
models[end] = mfull

# ### We create a prediction and plot function on a grid
function compute_grid(model, n_grid=50)
    mins = [-3.25, -2.85]
    maxs = [3.65, 3.4]
    x_lin = range(mins[1], maxs[1]; length=n_grid)
    y_lin = range(mins[2], maxs[2]; length=n_grid)
    x_grid = Iterators.product(x_lin, y_lin)
    y_grid, _ = proba_y(model, vec(collect.(x_grid)))
    return y_grid, x_lin, y_lin
end

function plot_model(model, X, Y, title=nothing; size=(300, 500))
    n_grid = 50
    y_pred, x_lin, y_lin = compute_grid(model, n_grid)
    title = if isnothing(title)
        (model isa SVGP ? "M = $(AGP.dim(model[1]))" : "full")
    else
        title
    end
    p = plot_data(X, Y; size=size)
    Plots.contour!(
        p,
        x_lin,
        y_lin,
        reshape(y_pred, n_grid, n_grid)';
        cbar=false,
        levels=[0.5],
        fill=false,
        color=:black,
        linewidth=2.0,
        title=title,
    )
    if model isa SVGP
        Plots.scatter!(
            p, eachrow(hcat(AGP.Zview(model[1])...))...; msize=2.0, color="black", lab=""
        )
    end
    return p
end;

# ### Now run the prediction for every model and visualize the differences
Plots.plot(
    plot_model.(models, Ref(X), Ref(Y))...; layout=(1, length(models)), size=(1000, 200)
)

# ## Bayesian SVM vs Logistic
# ### We now create a model with the Bayesian SVM likelihood

mbsvm = VGP(X, Y, kernel, BayesianSVM(), AnalyticVI(); optimiser=false)
@time train!(mbsvm, 5)
# ### And compare it with the Logistic likelihood
Plots.plot(
    plot_model.(
        [models[end], mbsvm], Ref(X), Ref(Y), ["Logistic", "BSVM"]; size=(500, 500)
    )...;
    layout=(1, 2),
)


# please instantiate ./test/Manifest.toml for this

using Combinatorics, JLD2, DataStructures
import Statistics

# ===============================
# Grow and Shrink Fuzz Fixtures
# ===============================

# format: [(values, expected_medians)]

grow_shrink_values = [[1.], [1,2], [2,1]]
append!(grow_shrink_values, permutations([1,2,3])|>collect)
append!(grow_shrink_values, permutations([1,2,3,4])|>collect)
for n in [6,7,13,14,50,51], _ in 1:20
    push!(grow_shrink_values, rand(n))
end

function generate_grow_shrink_medians(values)
    N = length(values)
    medians = Array{Float64}(undef, 2N-1)
    # grow phase
    for i in 1:N
        medians[i] = Statistics.median(values[1:i])
    end
    # shrink phase
    for i in N+1:2N-1
        medians[i] = Statistics.median(values[i-N+1:end])
    end
    return medians
end

grow_shrink_medians = map(generate_grow_shrink_medians, grow_shrink_values)

grow_shrink_fixtures = [(grow_shrink_values[i], grow_shrink_medians[i]) for i in eachindex(grow_shrink_values)]

grow_shrink_filename = dirname(@__FILE__) * "/grow_shrink.jld2"
@save grow_shrink_filename grow_shrink_fixtures


#=

N-Queens problem in ConstraintSolver.jl

Performance (timeout 60s):
n:8
  0.004049 seconds (23.04 k allocations: 1.663 MiB)

n:12
  0.018140 seconds (113.31 k allocations: 8.275 MiB)

n:50
  0.102909 seconds (531.49 k allocations: 43.128 MiB)

n:100
  0.297124 seconds (637.23 k allocations: 69.590 MiB, 34.21% gc time)

n:200
  1.427308 seconds (2.57 M allocations: 341.693 MiB, 7.26% gc time)

n:300
  4.453372 seconds (5.99 M allocations: 878.029 MiB, 6.68% gc time)

n:400
 10.641826 seconds (11.54 M allocations: 1.767 GiB, 9.06% gc time)

n:500
status:TIME_LIMIT
 60.136410 seconds (110.37 M allocations: 8.428 GiB, 10.85% gc time)

 77.461974 seconds (132.85 M allocations: 11.574 GiB, 10.32% gc time, 0.21% compilation time)

Model created by Hakan Kjellerstrand, hakank@gmail.com
See also my Julia page: http://www.hakank.org/julia/

=#

using ConstraintSolver, JuMP
using Cbc, GLPK, Ipopt
const CS = ConstraintSolver
include("constraints_utils.jl")

function nqueens(n=8,print_solutions=true,all_solutions=true,timeout=6)

    cbc_optimizer = optimizer_with_attributes(Cbc.Optimizer, "logLevel" => 0)
    glpk_optimizer = optimizer_with_attributes(GLPK.Optimizer)
    ipopt_optimizer = optimizer_with_attributes(Ipopt.Optimizer)

    model = Model(optimizer_with_attributes(CS.Optimizer,   "all_solutions"=> all_solutions,
                                                            # "all_optimal_solutions"=>all_solutions,
                                                            "logging"=>[],

                                                            "traverse_strategy"=>:BFS,
                                                            # "traverse_strategy"=>:DFS,
                                                            # "traverse_strategy"=>:DBFS,

                                                            # "branch_split"=>:Smallest,
                                                            "branch_split"=>:Biggest,
                                                            # "branch_split"=>:InHalf,

                                                            # https://wikunia.github.io/ConstraintSolver.jl/stable/options/#branch_strategy-(:Auto)
                                                            "branch_strategy" => :IMPS, # default
                                                            # "branch_strategy" => :ABS, # Activity Based Search
                                                            # "activity.decay" => 0.999, # default 0.999
                                                            # "activity.max_probes" => 10, # default, 10
                                                            # "activity.max_confidence_deviation" => 20, # default 20

                                                            # "simplify"=>false,
                                                            # "simplify"=>true, # default

                                                            "time_limit"=>timeout,

                                                            # "backtrack" => false, # default true
                                                            # "backtrack_sorting" => false, # default true

                                                            # "lp_optimizer" => cbc_optimizer,
                                                            # "lp_optimizer" => glpk_optimizer,
                                                            # "lp_optimizer" => ipopt_optimizer,
                                        ))

    @variable(model, 1 <= x[1:n] <= n, Int)
    @variable(model, 0 <= t1[1:n] <= n*2, Int)
    @variable(model, -n <= t2[1:n] <= n, Int)

    @constraint(model, x in CS.AllDifferent())
    for i in 1:n
        @constraint(model, x[i] + i == t1[i])
        @constraint(model, x[i] - i == t2[i])
    end
    @constraint(model, t1 in CS.AllDifferent())
    @constraint(model, t2 in CS.AllDifferent())

    # This don't work:
    # """
    # Each variable must be an integer and bounded.
    # Currently the variable index 9 doesn't fulfill this requirements.
    # """
    # @constraint(model, [t1[i] + i for i in 1:n] in CS.AllDifferent())
    # @constraint(model, [x[i] - i for i in 1:n] in CS.AllDifferent())

    #=
    # This is another approach, but slower.
    for i in 1:n, j in i+1:n
        @constraint(model, x[i] != x[j])
        @constraint(model, x[i] + i != x[j] + j)
        @constraint(model, x[i] - i != x[j] - j)
    end
    =#

    # Solve the problem
    optimize!(model)

    status = JuMP.termination_status(model)
    # println("status:$status")
    num_sols = 0
    if status == MOI.OPTIMAL
        num_sols = MOI.get(model, MOI.ResultCount())
        # println("num_sols:$num_sols\n")
        if print_solutions
            for sol in 1:num_sols
                # println("solution #$sol")
                x_val = convert.(Integer,JuMP.value.(x; result=sol))
                # t1_val = convert.(Integer,JuMP.value.(t1; result=sol))
                # t2_val = convert.(Integer,JuMP.value.(t2; result=sol))
                println("x:$x_val")
                # println("t1:$t1_val")
                # println("t2:$t2_val")
                # println()

            end
        end
    else
        println("status:$status")
    end
    println("num_sols: $num_sols")
    return status, num_sols
end

@time nqueens(8)

for n in [8,12,50,100,200,300,400,500]
    println("\nn:$n")
    @time status, num_sols = nqueens(n,false,false,60)
    if status == MOI.TIME_LIMIT
        break
    end
end


using Test

using CompScienceMeshes
using StaticArrays

T = Float64
P = SVector{3,T}

width, height = 1.0, 1.0
mesh = meshrectangle(width, height, 0.5)

# This predicate tests whether a simplex is in the (x==0) plane
function pred(simplex)
    verts = mesh.vertices[simplex]
    for i in 1 : length(verts)
        if abs(verts[i][1]) > eps(Float64)
            return false
        end
    end
    return true
end

# Test if meshes can be intersected
line = meshsegment(width, 1/3, 3)
translate!(line, P(0.0, height, 0.0))

edges = skeleton(mesh, 1)
γ1 = submesh(line, edges)
@test numcells(γ1) == 2
for c in cells(γ1)
    ctr = cartesian(center(chart(edges, c)))
    @test ctr[2] == height
    @test ctr[3] == 0
end
# @test γ1.faces == [
#     index(3,6),
#     index(6,9)]

pred2 = interior_tpredicate(mesh)
γ2 = submesh(pred2, edges)
@test numcells(γ2) == 8

overlaps = overlap_gpredicate(line)
pred3 = x -> overlaps(simplex(vertices(mesh, x)))
pred4 = x -> pred2(x) || pred3(x)
γ3 = submesh(pred4, edges)
@test numcells(γ3) == 10

#test interior_vpredicate
function boundaryvertices(mesh::Mesh)
    vpred = interior_vpredicate(mesh)
    bidx = findall(!vpred,1:numvertices(mesh))
    bnd = mesh.vertices[bidx]
end

tr = meshrectangle(1.0,1.0,1/4)
bv = boundaryvertices(tr)
btr = boundary(tr)
@test length(bv) == numcells(skeleton(btr,0))


using MixedModels
using Test
using StableRNGs

@testset "fixed sigma" begin
    σ = 3
    n = 100
    dat = (; x = ones(n),
            z = collect(1:n),
            y = σ*randn(StableRNG(42), n))

    fmσ1 = fit(MixedModel, @formula(y ~ 0 + (1|z)), dat;
               contrasts=Dict(:z => Grouping()),
               σ=1)
    @test isempty(fixef(fmσ1))
    # verify that we report the exact value requested
    @test fmσ1.σ == 1
    # verify that the constrain actually worked
    @test pwrss(fmσ1) / nobs(fmσ1) ≈ 1.0
    @test only(fmσ1.θ) ≈ σ atol=0.1

    fmσ1 = fit(MixedModel, @formula(y ~ 0 + (1|z)), dat;
               contrasts=Dict(:z => Grouping()),
               σ=3.14)
    @test isempty(fixef(fmσ1))
    # verify that we report the exact value requested
    @test fmσ1.σ == 3.14
    # verify that the constrain actually worked
    @test pwrss(fmσ1) / nobs(fmσ1) ≈ 3.14^2 atol=0.5
    # the shrinkage forces things to zero because 3.14/3 is very close to 0
    @test only(fmσ1.θ) ≈ 0 atol=0.1
end

# specifying sigma was done to allow for doing meta-analytic models
# the example from metafor that doesn't work with lme4 and R-based nlme
# can be done here!
# https://www.metafor-project.org/doku.php/tips:rma_vs_lm_lme_lmer
#
# using RCall
# using MixedModels
# R"""
# library(metafor)
# dat <- escalc(measure="ZCOR", ri=ri, ni=ni, data=dat.molloy2014)
# dat$study <- 1:nrow(dat)
# """
# @rget dat

# fit(MixedModel, @formula(yi ~ 1 + (1 | study)), dat;
#     wts=1 ./ dat.vi,
#     REML=true,
#     contrasts=Dict(:study => Grouping()),
#     σ=1)


export CompanionMatrix, CompanionMatrixTransposed

import Base: *
import LinearAlgebra: transpose, inv

"""
    CompanionMatrix

Represents a matrix of the following structure:

θ1 θ2 θ3 ... θn-1 θn
 1  0  0  ...   0  0
 0  1  0  ...   0  0
 .  .  .  ...   .  .
 .  .  .  ...   .  .
 0  0  0  ...   0  0
 0  0  0  ...   1  0
"""
struct CompanionMatrix{ R <: Real, T <: AbstractVector{R} } <: AbstractMatrix{R}
    θ :: T
end

Base.eltype(::CompanionMatrix{R}) where R = R
Base.size(cmatrix::CompanionMatrix)       = (length(cmatrix.θ), length(cmatrix.θ))
Base.length(cmatrix::CompanionMatrix)     = prod(size(cmatrix))

Base.getindex(cmatrix::CompanionMatrix, i::Int) = getindex(cmatrix, map(r -> r + 1, reverse(divrem(i - 1, first(size(cmatrix)))))...)

function Base.getindex(cmatrix::CompanionMatrix, i::Int, j::Int)
    if i === 1
        return cmatrix.θ[j]
    elseif i === j + 1
        return one(eltype(cmatrix))
    else
        return zero(eltype(cmatrix))
    end
end

struct CompanionMatrixTransposed{ R <: Real, T <: AbstractVector{R} } <: AbstractMatrix{R}
    θ :: T
end

Base.eltype(::CompanionMatrixTransposed{R}) where R = R
Base.size(cmatrix::CompanionMatrixTransposed)       = (length(cmatrix.θ), length(cmatrix.θ))
Base.length(cmatrix::CompanionMatrixTransposed)     = prod(size(cmatrix))

Base.getindex(cmatrix::CompanionMatrixTransposed, i::Int) = getindex(cmatrix, map(r -> r + 1, reverse(divrem(i - 1, first(size(cmatrix)))))...)

function Base.getindex(cmatrix::CompanionMatrixTransposed, i::Int, j::Int)
    if j === 1
        return cmatrix.θ[i]
    elseif j === i + 1
        return one(eltype(cmatrix))
    else
        return zero(eltype(cmatrix))
    end
end

as_companion_matrix(θ::T) where { R, T <: AbstractVector{R} } = CompanionMatrix{R, T}(θ)
as_companion_matrix(θ::T) where { T <: Real }                 = θ

LinearAlgebra.transpose(cmatrix::CompanionMatrix)           = CompanionMatrixTransposed(cmatrix.θ)
LinearAlgebra.transpose(cmatrix::CompanionMatrixTransposed) = CompanionMatrix(cmatrix.θ)

LinearAlgebra.adjoint(cmatrix::CompanionMatrix)           = CompanionMatrixTransposed(cmatrix.θ)
LinearAlgebra.adjoint(cmatrix::CompanionMatrixTransposed) = CompanionMatrix(cmatrix.θ)

LinearAlgebra.inv(t::Union{CompanionMatrix, CompanionMatrixTransposed}) = inv(as_matrix(t))

function as_matrix(cmatrix::CompanionMatrix)
    dim     = first(size(cmatrix))
    S       = zeros(dim, dim)
    S[1, :] = cmatrix.θ
    for i in 2:dim
        S[i, i - 1] = one(eltype(cmatrix))
    end
    S
end

function as_matrix(cmatrix::CompanionMatrixTransposed)
    dim     = first(size(cmatrix))
    S       = zeros(dim, dim)
    S[:, 1] = cmatrix.θ
    for i in 2:dim
        S[i - 1, i] = one(eltype(cmatrix))
    end
    S
end

function Base.:*(tm::CompanionMatrix, v::AbstractVector)
    r = similar(v)
    r[1] = tm.θ' * v
    for i in 1:(length(v) - 1)
        r[i + 1] = v[i]
    end
    return r
end

function Base.:*(tm::CompanionMatrix, m::AbstractMatrix)

    l = length(tm.θ)
    r = similar(m)

    for (index, col) in enumerate(eachcol(m))
        r[1, index] = tm.θ' * col
        for i in 1:(l - 1)
            r[i + 1, index] = col[i]
        end
    end

    return r
end

function Base.:*(m::AbstractMatrix, tm::CompanionMatrixTransposed)

    l = length(tm.θ)
    r = similar(m)

    for (index, row) in enumerate(eachrow(m))
        r[index, 1] = tm.θ' * row
        for i in 1:(l - 1)
            r[index, i + 1] = row[i]
        end
    end

    return r
end

function Base.:*(m::AbstractMatrix, tm::CompanionMatrix)

    l = length(tm.θ)
    r = similar(m)

    return (tm' * m')'
end

function Base.:*(tm::CompanionMatrixTransposed, m::AbstractMatrix)

    l = length(tm.θ)
    r = similar(m)

    for (index, col) in enumerate(eachcol(m))
        for j in 1:l-1
            r[j, index] = tm.θ[j] * col[1] + col[j + 1]
        end
        r[l, index] = tm.θ[l] * col[1]
    end

    return r
end


function hmc_step(current_state::state, log_likelihood::Function, grad_log_likelihood::Function, current_p::Vector{<:Real}, ϵ::Float64, L::Int)

    current_log_likelihood = current_state.log_likelihood[1] + 0.5 * sum(current_p.^2)

    # Evolve p half-step
    current_p += 0.5 * ϵ * current_state.grad_log_likelihood

    # Leap-Frog Integrate
    x1 = current_state.x
    for i = 1:L-1
        x1 += ϵ * current_p
        current_p += ϵ * grad_log_likelihood(x1)
    end
    x1 += ϵ * current_p

    current_p += 0.5 * ϵ * grad_log_likelihood(x1)
    current_p = -1 * current_p

    # Perform comparison
    new_log_likelihood = log_likelihood(x1)
    proposed_log_likelihood = new_log_likelihood + 0.5 * sum(current_p.^2)
    alpha = current_log_likelihood - proposed_log_likelihood
    u = log(rand())
    if u <= alpha
        current_state.x[:] .= x1
        current_state.log_likelihood[1] = new_log_likelihood
        return 1
    else
        return 0
    end

end

function hmc(x0 :: Vector{<:Real}, log_likelihood :: Function, grad_likelihood_function :: Function, num_burnin :: Int, num_samples :: Int)

    # Initialize
    ϵ = 1e-2
    L = 20

    # Initialize
    current_state = state(x0, log_likelihood(x0), grad_likelihood_function(x0))

    # Burnin
    burnin_chain = chain(x0, num_burnin)
    burnin_chain.diagnostics.log_likelihood[1] = log_likelihood(x0)
    for iteration = 1:num_burnin-1
        p = randn(size(x0))
        flag = hmc_step(current_state, log_likelihood, grad_likelihood_function, p, ϵ, L)
        burnin_chain.x[iteration+1] = copy(current_state.x)
        burnin_chain.diagnostics.log_likelihood[iteration+1] = current_state.log_likelihood[1]
        burnin_chain.diagnostics.acceptance_ratio[iteration+1] = ( burnin_chain.diagnostics.acceptance_ratio[iteration] * (iteration - 1) + flag ) / iteration
    end

    if burnin_chain.diagnostics.acceptance_ratio[end] < 0.25
        ϵ *= 0.5
    elseif burnin_chain.diagnostics.acceptance_ratio[end] > 0.9
        ϵ *= 2
    end

    # Sampling
    sampling_chain = chain(current_state.x, num_samples)
    sampling_chain.diagnostics.log_likelihood[1] = log_likelihood(x0)

    for iteration = 1:num_samples-1
        p = randn(size(x0))
        flag = hmc_step(current_state, log_likelihood, grad_likelihood_function, p, ϵ, L)
        sampling_chain.x[iteration+1] = copy(current_state.x)
        sampling_chain.diagnostics.log_likelihood[iteration+1] = current_state.log_likelihood[1]
        sampling_chain.diagnostics.acceptance_ratio[iteration+1] = ( sampling_chain.diagnostics.acceptance_ratio[iteration] * (iteration - 1) + flag ) / iteration
    end

    return burnin_chain, sampling_chain
end


using BenchmarkTools
using Distributions

d = Normal()
X = rand(d, 5000)

@btime pdf.(d, X)
@btime logpdf.(d, X)

using Statistikk

@btime Statistikk.pdf(d, X)
@btime Statistikk.logpdf(d, X)


using LinearAlgebraicRepresentation
using Test

@testset "simplexn" begin
	@testset "extrudeSimplicial" begin
		V = [[0,0] [1,0] [2,0] [0,1] [1,1] [2,1] [0,2] [1,2] [2,2]];
		FV = [[1,2,4],[2,3,5],[3,5,6],[4,5,7],[5,7,8],[6,8,9]];
		model = (V,FV);
		pattern = repeat([1,2,-3],outer=4);
		W,FW = LinearAlgebraicRepresentation.extrudeSimplicial(model, pattern);
		@test typeof(LinearAlgebraicRepresentation.extrudeSimplicial(model, pattern))==
			Tuple{Array{Int64,2},Array{Array{Int64,1},1}}
		@test typeof(W)==Array{Int64,2}
		@test size(W)==(3,117)
		@test typeof(FW)==Array{Array{Int64,1},1}
		@test length(FW)==144
	end

	@testset "simplexGrid" begin
		#@test LinearAlgebraicRepresentation.simplexGrid([0])==([0.0], Array{Int64,1}[])
		@test LinearAlgebraicRepresentation.simplexGrid([1])==([0.0 1.0], Array{Int64,1}[[1, 2]])
		@test LinearAlgebraicRepresentation.simplexGrid([1,1])==([0.0 1.0 0.0 1.0; 0.0 0.0 1.0 1.0],
			Array{Int64,1}[[1, 2, 3], [2, 3, 4]])
		@test typeof(LinearAlgebraicRepresentation.simplexGrid([1,1,1])[2])==Array{Array{Int64,1},1}
		@test length(LinearAlgebraicRepresentation.simplexGrid([1,1,1])[2])==6
		V,CV = LinearAlgebraicRepresentation.simplexGrid([10,10,1])
		@test typeof(V)==Array{Float64,2}
		@test typeof(CV)==Array{Array{Int64,1},1}
		@test size(V)==(3,242)
		@test length(CV)==600
	end

	@testset "simplexFacets" begin
		@test LinearAlgebraicRepresentation.simplexGrid([1,1])[1]==[0 1 0 1; 0 0 1 1]
		@test LinearAlgebraicRepresentation.simplexGrid([1,1])[2]==Array{Int64,1}[[1, 2, 3], [2, 3, 4]]
		#=
		@test typeof(LinearAlgebraicRepresentation.extrudeSimplicial((V,FV), [1])[1])==Array{Float64,2}
		@test typeof(LinearAlgebraicRepresentation.extrudeSimplicial((V,FV), [1])[2])==LinearAlgebraicRepresentation.Cells
		@test size(LinearAlgebraicRepresentation.extrudeSimplicial((V,FV), [1])[1])==(4, 484)
		@test length(LinearAlgebraicRepresentation.extrudeSimplicial((V,FV), [1])[2])==18
		W,CW = LinearAlgebraicRepresentation.extrudeSimplicial((V,FV), [1])
		@test typeof(LinearAlgebraicRepresentation.simplexFacets(CW))==Array{Array{Int64,1},1}
		@test size(LinearAlgebraicRepresentation.simplexFacets(CW))==(56,)
		@test typeof(LinearAlgebraicRepresentation.simplexFacets(CW))==Array{Array{Int64,1},1}
		@test size(LinearAlgebraicRepresentation.simplexFacets(CW))==(56,)
		=#
	end
end


module MeshingBenchmarks

using FileIO
using NRRD
using Meshing
using MeshIO
using GeometryTypes
using BenchmarkTools

function benchmark()
    here = dirname(@__FILE__)

    println(pwd())
    println("CTA-cardio.nrrd loading...")
    ctacardio = @btime load($here*"/../data/CTA-cardio.nrrd")
    q = 100

    samples = 1

    println("CTA-cardio.nrrd loaded")

    println("CTA-cardio.nrrd MarchingCubes Float32 runtime")
    mc = @btime HomogenousMesh{Point{3,Float32}, Face{3,Int}}($ctacardio, MarchingCubes(iso=Float32($q), insidepositive=true))
    for i in 1:samples-1
        @btime HomogenousMesh{Point{3,Float32}, Face{3,Int}}($ctacardio, MarchingCubes(iso=Float32($q), insidepositive=true))
    end

    println("CTA-cardio.nrrd MarchingCubes Float64 runtime")
    @btime HomogenousMesh{Point{3,Float64}, Face{3,Int}}($ctacardio, MarchingCubes(iso=$q, insidepositive=true))
    for i in 1:samples-1
        @btime HomogenousMesh{Point{3,Float64}, Face{3,Int}}($ctacardio, MarchingCubes(iso=$q, insidepositive=true))
    end

    # println("CTA-cardio.nrrd MarchingTetrahedra Float32 runtime")
    # mt = @btime HomogenousMesh{Point{3,Float32},Face{3,Int}}($ctasdf, MarchingTetrahedra($q))
    # for i in 1:samples-1
    #     @btime HomogenousMesh{Point{3,Float32},Face{3,Int}}($ctasdf, MarchingTetrahedra($q))
    # end

    println("Saving files")
    save("ctacardio_mc.ply", mc)
    #save("ctacardio_mt.ply", mt)
end

function brain(samples=1, q=100)

    here = dirname(@__FILE__)

    brain = load(here*"/../Seg3DData/Brain_DataSet/MRI-brain50.nrrd")

    brainsdf = SignedDistanceField(HyperRectangle(Vec(0,0,0), Vec(10,10,10)),brain.data)

    println("Brain.nrrd MarchingCubes Float32 runtime")
    mc_brain = @btime HomogenousMesh{Point{3,Float32},Face{3,Int}}($brainsdf, MarchingCubes(iso=$q, insidepositive=false))
    for i in 1:samples-1
        @btime HomogenousMesh{Point{3,Float32},Face{3,Int}}($brainsdf, MarchingCubes(iso=$q, insidepositive=false))
    end

    save("brian_mc.ply", mc_brain)
end

end # module


#-----------------------------------------------------------------------------------------------
# 1st run - packages installation

# ERROR: LoadError: ArgumentError: Package DataFrames not found in current path:
# - Run `import Pkg; Pkg.add("DataFrames")` to install the DataFrames package.
import Pkg; Pkg.add("DataFrames")
import Pkg;
Pkg.add("CSV")
Pkg.add("StatsBase")

#-----------------------------------------------------------------------------------------------
using DataFrames, CSV, Statistics, StatsBase
data = CSV.read("../externals/Core.Math.Data/data/Pejcic_318.csv", copycols = true)

# @show data.ATV
# println()

for col in eachcol(data);
    if  (
        eltype(col[2]) <: Real)
        println(mean(col[2])
        );
    end;
end;


module CenteredAxes
export Centered

"""
	Centered

Given an array axis of length `N`, return `-N÷2:N÷2-1` for even `N` and `-N÷2:N÷2` for odd `N`.
This may be used to construct or shift an axis of an `OffsetArray`.

# Examples
```julia
julia> A = Float64.(reshape(1:9, 3, 3))
3×3 Array{Float64,2}:
 1.0  4.0  7.0
 2.0  5.0  8.0
 3.0  6.0  9.0

julia> OffsetArray(A, Centered(), Centered())
3×3 OffsetArray(::Array{Float64,2}, -1:1, -1:1) with eltype Float64 with indices -1:1×-1:1:
 1.0  4.0  7.0
 2.0  5.0  8.0
 3.0  6.0  9.0

julia> OffsetArray(A, Centered(), :)
3×3 OffsetArray(::Array{Float64,2}, -1:1, 1:3) with eltype Float64 with indices -1:1×1:3:
 1.0  4.0  7.0
 2.0  5.0  8.0
 3.0  6.0  9.0

julia> A = Float64.(reshape(1:4, 2, 2))
2×2 Array{Float64,2}:
 1.0  3.0
 2.0  4.0

julia> OffsetArray(A, Centered(), Centered())
2×2 OffsetArray(::Array{Float64,2}, -1:0, -1:0) with eltype Float64 with indices -1:0×-1:0:
 1.0  3.0
 2.0  4.0
```
"""
struct Centered end

@inline function Base.to_indices(A, ax::Tuple{Any, Vararg{Any}}, I::Tuple{Centered, Vararg{Any}})
	ax1 = first(ax)
	Lby2 = length(ax1)÷2
	axc = -Lby2:length(ax1) - Lby2 - 1
	(axc, to_indices(A, Base.tail(ax), Base.tail(I))...)
end

end


using Pkg
Pkg.activate(".")

#using LSDescent

using NLPModels, JuMP,  NLPModelsJuMP

using SolverCore, Logging

using LinearOperators

using LinearAlgebra

function test_algo(algo::Function, nlp, B, maxiter )
    reset!(nlp)

    logger = Logging.ConsoleLogger(stderr,Logging.Warn)
    Logging.with_logger(logger) do
        iter, f, normg, B, x = algo(nlp, B₀ = B, maxiter = maxiter)
        @show iter, f, normg
    end

end

function test_algo_stp(algo::Function, nlp, stp, B )
    reset!(nlp)

    logger = Logging.ConsoleLogger(stderr,Logging.Warn)
    Logging.with_logger(logger) do
        stp, B = algo(nlp, stp=stp,  B₀ = B)
        iter, f, normg =  stp.meta.nb_of_stop, stp.current_state.fx, stp.current_state.current_score
        @show iter, f, normg
    end

end

function test_wrapper_stp(algo::Function, nlp, stp; kwargs... )
    reset!(nlp)

    logger = Logging.ConsoleLogger(stderr,Logging.Warn)
    Logging.with_logger(logger) do
        stp, B = algo(nlp, stp=stp; kwargs...)
        iter, f, normg =  stp.meta.nb_of_stop, stp.current_state.fx, stp.current_state.current_score
        @show iter, f, normg
    end

end

# test all solvers with the well known Woods test function
using OptimizationProblems

#include("woods.jl")
#include("genrose.jl")
#nlp = MathOptNLPModel(PureJuMP.dixmaank(40), name="dixmaank")
#nlp = MathOptNLPModel(PureJuMP.dixmaang(100), name="dixmaang")
#nlp = MathOptNLPModel(PureJuMP.srosenbr(10), name="srosenbr")
nlp = MathOptNLPModel(PureJuMP.woods(8), name="woods")
#nlp = MathOptNLPModel(PureJuMP.genrose(100), name="genrose")
n = nlp.meta.nvar

include("bfgsSolver.jl")
include("FormuleN2.jl")

maxiter = 1600
scaling = false
#@info "Version avec matrice"
#bfgs(nlp, maxiter = maxiter)

include("Type.jl")

@info "Version encapsulée opérateur, formule O(n^3)"
include("FormuleNOp.jl")
B = InverseBFGSOperator(Float64, n; scaling = scaling);

test_algo(bfgs, nlp, B, maxiter )
#bfgs(nlp, B₀ = B, maxiter = maxiter)
#Bfull =  Matrix(B)

@info "Version encapsulée opérateur, formule O(n^2)"
include("FormuleN2Op.jl")
B = InverseBFGSOperator(Float64, n; scaling = scaling)

test_algo(bfgs, nlp, B, maxiter )
#bfgs(nlp, B₀ = B, maxiter = maxiter)
#Bfull2 =  Matrix(B)

@info "Version encapsulée opérateur, formule O(n^2), stp"
using Stopping
include("bfgsStop.jl")
stp = NLPStopping(nlp,
                  NLPAtX(nlp.meta.x0)  )
stp.meta.max_iter = maxiter

B = InverseBFGSOperator(Float64, n; scaling = scaling)

test_algo_stp(bfgs_Stop, nlp, stp, B )

@info "Version encapsulée opérateur, formule O(n^2), stp LS"
using OneDmin
include("bfgsStopLS.jl")
stp = NLPStopping(nlp,
                  NLPAtX(nlp.meta.x0)  )
stp.meta.max_iter = maxiter
reinit!(stp)

test_wrapper_stp(M_bfgs_StopLS, nlp, stp; scaling = scaling )

B = InverseBFGSOperator(Float64, n; scaling = scaling)
reinit!(stp)

test_algo_stp(bfgs_StopLS, nlp, stp, B )

@info "Version L-BFGS"
reinit!(stp)

test_wrapper_stp(L_bfgs_Stop, nlp, stp; scaling = scaling )

B = InverseLBFGSOperator(Float64, n, mem=160; scaling = scaling)
reinit!(stp)

test_algo_stp(bfgs_Stop, nlp, stp, B)

;
#bfgs(nlp, B₀ = B, maxiter = maxiter)
#BL = Matrix(B)

#norm(Bfull-BL)


using ADCME
using JLD2

@load "data.jld2" z

n = 64
noise = parse(Float64, ARGS[1])

τ = Variable(1.0)
x = placeholder(Float64, shape = [n,])
y = placeholder(Float64, shape = [n,])
w = placeholder(Float64, shape = [n,])
σ, κ = 0.08, 0.5
α = 0.5
Δt = 0.01
yhat = (
    x + κ * (τ - α*x)*Δt + σ * sqrt(x) * sqrt(Δt) * w  + 1/4*σ^2*Δt*(w^2-1)
)/(1+(1-α)*κ*Δt)

θ = Variable(fc_init([2,20,20,20,1]))
hDξ = sigmoid(squeeze(fc([x yhat], [20,20,20,1], θ)))
Dξ = sigmoid(squeeze(fc([x y], [20,20,20,1], θ)))
# LF = mean(log((1-hDξ)/hDξ))
LF = -mean(log(hDξ))
LD = -mean(log(Dξ) + log(1-hDξ))

function generate_batch_data()
    idx = rand(1:9999, n)
    abs.(z[idx] + noise * randn(n)), abs.(z[idx.+1] + noise * randn(n)), randn(n)
end

optF = RMSPropOptimizer(1e-3).minimize(LF, var_list=[τ])
optD = RMSPropOptimizer(1e-3).minimize(LD, var_list=[θ])

sess = Session(); init(sess)
x0, y0, w0 = generate_batch_data()
    fd = Dict(
        x => x0, y => y0, w => w0
    )
@info run(sess, [LF, LD], feed_dict = fd)

fiter = 1
diter = 5
l1 = []
l2 = []
τs = []

for i = 1:50000
    x0, y0, w0 = generate_batch_data()

    fd = Dict(
        x => x0, y => y0, w => w0
    )

    ld = 0
    for k = 1:diter
        _, ld = run(sess, [optD, LD], feed_dict = fd)
    end

    lf = 0
    for k = 1:fiter
        _, lf = run(sess, [optF, LF], feed_dict = fd)
    end

    τ0 = run(sess, τ)

    # if length(loss)==0
    #     loss = [l1[end]]
    # end
    push!(l1, ld)
    push!(l2, lf)
    push!(τs, τ0)
    if mod(i, 100)==0
        @info i, ld, lf, τ0
    end
end

kappa = join(string.(τs), ',')
db = Database("noise.db")
execute(db, """
CREATE TABLE IF NOT EXISTS tau (
    noise real,
    kappa text
)
""")
execute(db, """
INSERT INTO tau VALUES ($noise, \"$kappa\")
""")
close(db)
# close("all")
# plot(τs)
# savefig("test$noise.png")


import .Utils.Heap: PriorityQueue, dequeue!

"""
    prim_mst_edges(n, edges, weights, [root = 1])

[Prim's algorithm](https://en.wikipedia.org/wiki/Prim%27s_algorithm)
for minimum spanning tree.
Start from node `root` (will become the root node of the spanning
tree).  Return the parent vector and a Boolean vector indicating for
each edge whether the edge is part of the spanning tree.
"""
function prim_mst_edges(
    n::Integer,
    edges::Vector{E},
    weights::Array{Float64},
    root::Int = 1,
) where {E}
    mem = PrimMstMem(edges, n)
    prim_mst_edges(weights, root, mem)
    # swap root element to front
    mem.selected[1], mem.selected[root] = mem.selected[root], mem.selected[1]
    return mem.parent, view(mem.selected, 2:length(mem.selected))
end

struct PrimMstMem
    dist::Vector{Float64}
    parent::Vector{Int}
    neighbors::IncidenceIndex
    selected::Vector{Int}
    pq::PriorityQueue{Int, Float64}

    PrimMstMem(g::EdgeGraph) =
        PrimMstMem(IncidenceIndex(num_nodes(g), g.edges))

    PrimMstMem(g::Graph) =
        PrimMstMem(EdgeGraph(g))

    PrimMstMem(edges::Vector{E}, n::Integer) where {E} =
        PrimMstMem(IncidenceIndex(n, edges))

    PrimMstMem(head::Vector{I}, tail::Vector{I}, n::Integer) where {I} =
        PrimMstMem(IncidenceIndex(n, head, tail))

    function PrimMstMem(neighbors::IncidenceIndex)
        local n = num_nodes(neighbors)
        dist = Vector{Float64}(undef, n)
        parent = Vector{Int}(undef, n)
        selected = Vector{Int}(undef, n)
        pq = PriorityQueue{Int, Float64}(n)
        return new(dist, parent, neighbors, selected, pq)
    end
end

prim_mst_edges(
    weights::Vector{Float64}, root::Int, edges::Vector{E}, n::Integer,
) where {E} =
    prim_mst_edges(weights, root, PrimMstMem(edges, n))

prim_mst_edges(weights::Vector{Float64}, root::Int, mem::PrimMstMem) =
    prim_mst_edges(
        weights,
        mem.dist,
        mem.parent,
        mem.neighbors,
        mem.selected,
        mem.pq,
        root
    )

function prim_mst_edges(
    edge_weight::Vector{Float64},
    dist::Vector{Float64},
    parent::Vector{Int},
    neighbors::IncidenceIndex,
    selected::Vector{Int},
    pq::PriorityQueue{Int, Float64},
    root::Int = 1,
)::Vector{Int}
    @assert root >= 1 "root = $root"
    @assert length(selected) >= length(parent)
    @assert isempty(pq) "length(pq) = $(length(pq))"
    sizehint!(pq, length(parent))

    # see CLRS page 572 (2nd edition)
    dist .= typemax(Float64)
    parent .= 0
    dist[root] = typemin(Float64)
    parent[root] = -root
    selected[root] = -1
    pq[root] = dist[root]
    while !isempty(pq)
        u = dequeue!(pq)
        for (v, eidx) in neighbors[u]
            v == u && continue
            if parent[v] <= 0 && edge_weight[eidx] < dist[v]
                dist[v] = edge_weight[eidx]
                pq[v] = dist[v]    # decrease_key!(v)
                parent[v] = -u
                selected[v] = eidx
            end
        end
        parent[u] *= -1
    end
    return parent
end

function _prim_mst_dbg(
    edge_weight::Vector{Float64},
    dist::Vector{Float64},
    parent::Vector{Int},
    neighbors::IncidenceIndex,
    pq::PriorityQueue{Int, Float64},
    root::Int = 1,
)::Vector{Int}
    @assert root >= 1 "root = $root"
    @assert isempty(pq) "length(pq) = $(length(pq))"
    sizehint!(pq, length(parent))

    dist .= typemax(Float64)
    parent .= 0
    dist[root] = typemin(Float64)
    parent[root] = -root
    pq[root] = dist[root]
    while !isempty(pq)
        u = dequeue!(pq)
        println("fin: ", u - 1, " pi = ", -parent[u] - 1, " was ", dist[u])
        parent[u] *= -1
        for (v, eidx) in neighbors[u]
            v == u && continue
            if parent[v] <= 0 && edge_weight[eidx] < dist[v]
                if isinf(dist[v])
                    println("new: ", v - 1, " --> ", edge_weight[eidx])
                else
                    println("upd: ", v - 1, " --> ", edge_weight[eidx], " was ", dist[v])
                end
                dist[v] = edge_weight[eidx]
                pq[v] = dist[v]
                parent[v] = -u
            end
        end
    end
    return parent
end


using NonuniformResampling1D
using Plots
using Statistics

xin = 1:24
yin = [
    0.10221248850845677,
    -0.3017964617372699,
    0.6313812022052889,
    -0.469567525619184,
    -0.742550821690563,
    0.538340882717059,
    0.7645542960785265,
    -0.42847475695892334,
    -0.2030496164626363,
    -0.6961115394075117,
    0.9553667194266671,
    0.6773297980833843,
    0.6595211409431194,
    0.18762849994721864,
    0.9676546988790669,
    -0.6599170737336646,
    0.17851166014365583,
    0.7417238170162541,
    -0.7312025977875327,
    -0.4015347276663681,
    -0.685997771780074,
    -0.7558107010474338,
    0.12236047934097583,
    -0.8627888068515903]
xout = [4.6, 5.7, 8.2, 13.3, 17.3, 18.2, 19.1, 20, 20.8]
red_ind = 3
blue_ind = 6
red_color = :coral
blue_color = :dodgerblue

w = 1.35
win = hann_window(w)
N = 300
win_y = win.(range(0, w, length=N))
required_points_per_slice = 4

red_left_border = xout[red_ind]*(1-w)+w*xout[red_ind-1]
red_right_border = xout[red_ind]*(1-w)+w*xout[red_ind+1]
blue_left_border = xout[blue_ind]*(1-w)+w*xout[blue_ind-1]
blue_right_border = xout[blue_ind]*(1-w)+w*xout[blue_ind+1]

yout = nuresample(xin, yin, xout, win, required_points_per_slice=required_points_per_slice)

function vertical_lines()
    vline!([xout[red_ind]], c=red_color, linewidth=1)
    vline!([red_left_border, red_right_border], c=red_color, line=:dash)
    vline!([xout[blue_ind]], c=blue_color, linewidth=1)
    vline!([blue_left_border, blue_right_border], c=blue_color, line=:dash)
    # unit lines make the image too crowded
    #vline!([xout[red_ind-1], xout[red_ind+1], xout[blue_ind-1], xout[blue_ind+1]], c=:black, line=:dot)
    hline!([0], c=:gray, z_order=:back)
end

# blue upsampling + points
upsample_step = min(xout[blue_ind]-xout[blue_ind-1], xout[blue_ind+1]-xout[blue_ind])/
    required_points_per_slice*w
upsample_x = range(blue_left_border+upsample_step/2, step=upsample_step,
stop=blue_right_border-upsample_step/2)
in_step = Float64(StepRangeLen(xin).step)
upsample_y = [NonuniformResampling1D.interpolate_point(
    StepRangeLen(xin), yin, up_x, in_step, in_step, lanczos_window()) for up_x in upsample_x]

# assume xout around blue is symmetrical
blue_weights = @. win(abs(upsample_x-xout[blue_ind])/(xout[blue_ind]-xout[blue_ind-1]))
blue_weighted = blue_weights .* upsample_y

plt_in = begin
    plot(xin, yin, m=:circle, c=:black, legend=false, xaxis=false)
    # red window
    red_x_left = range(stop=xout[red_ind], length=N, step=(xout[red_ind]-xout[red_ind-1])/N*w)
    plot!(red_x_left, reverse(win_y), c=red_color)
    red_x_right = range(start=xout[red_ind], length=N, step=(xout[red_ind+1]-xout[red_ind])/N*w)
    plot!(red_x_right, win_y, c=red_color)
    #blue window
    blue_x_left = range(stop=xout[blue_ind], length=N, step=(xout[blue_ind]-xout[blue_ind-1])/N*w)
    plot!(blue_x_left, reverse(win_y), c=blue_color)
    blue_x_right = range(start=xout[blue_ind], length=N, step=(xout[blue_ind+1]-xout[blue_ind])/N*w)
    plot!(blue_x_right, win_y, c=blue_color)
    plot!(upsample_x, upsample_y, c=:deepskyblue, m=:x, markersize=2)

    annotate!([
        (4, -1.1, ("adaptive, asymmetric\nwindow width", 9, red_color)),
        (21.5, .8, ("ad-hoc\nupsampling", 9, blue_color))
    ])

    vertical_lines()
    title!("Input", titlefontsize=13)
end

plt_windowed = begin
    red_start_ind = NonuniformResampling1D.find_first_above_or_equal(red_left_border, StepRangeLen(xin))
    red_middle_ind = NonuniformResampling1D.find_last_below(xout[red_ind], StepRangeLen(xin))
    red_stop_ind = NonuniformResampling1D.find_last_below_or_equal(red_right_border, StepRangeLen(xin))

    red_left_x = xin[red_start_ind:red_middle_ind]
    red_right_x = xin[red_middle_ind+1:red_stop_ind]

    red_left_weights = @. win(abs(red_left_x-xout[red_ind])/(xout[red_ind]-xout[red_ind-1]))
    red_right_weights = @. win(abs(red_right_x-xout[red_ind])/(xout[red_ind+1]-xout[red_ind]))

    red_y = yin[red_start_ind:red_stop_ind]
    red_weighted = [red_left_weights; red_right_weights] .* red_y
    plot(red_start_ind:red_stop_ind, red_weighted,
        m=:circle, c=red_color, legend=false, markerstrokewidth=0)

    plot!(upsample_x, blue_weighted, m=:circle, c=blue_color, markerstrokewidth=0)

    vertical_lines()
    title!("Example Slices Windowed", titlefontsize=13)
end

plt_out = begin
    plot(xout, yout, m=:circle, c=:black, legend=false)
    scatter!([xout[red_ind]], [yout[red_ind]], c=red_color, markerstrokewidth=0)
    scatter!([xout[blue_ind]], [yout[blue_ind]], c=blue_color, markerstrokewidth=0)

    quiver!([xout[red_ind], xout[red_ind]], [-.5, -1.0], quiver=(-[(xout[red_ind]-xout[red_ind-1]), (xout[red_ind]-red_left_border)], [0,0]), c=:black)
    annotate!([
        (mean(xout[red_ind-1:red_ind]), -.7, ("Unit", 9)),
        (mean([xout[red_ind], red_left_border]), -1.2, ("Slice", 9))
    ])

    vertical_lines()
    title!("Output", titlefontsize=13)
end

plot(plt_in, plt_windowed, plt_out, layout=(3,1), link=:x, xticks=false, xaxis=false, yaxis=false)
ylims!(-1.1, 1.1)
savefig("explanation_diagram.svg")


using Test, Random
Random.seed!(1234321)

using HyperDualMatrixTools
using HyperDualNumbers, LinearAlgebra, SparseArrays, SuiteSparse

@testset "Testing HyperDualMatrixTools" begin
    # Chose a size for matrices
    n = 10

    x = randn(n)                 # real-valued      │
    y = randn(n)                 # real-valued      │
    z = randn(n)                 # real-valued      ├─ vector
    w = randn(n)                 # real-valued      │
    d = x + ε₁*y + ε₂*z + ε₁ε₂*w # hyperdual-valued │

    A = randn(n, n)              # real-valued      │
    B = randn(n, n)              # real-valued      │
    C = randn(n, n)              # real-valued      ├─ matrix
    D = randn(n, n)              # real-valued      │
    M = A + ε₁*B + ε₂*C + ε₁ε₂*D # hyperdual-valued │

    @testset "Testing full matrices" begin
        @testset "Testing `\\` without factorization" begin
            @test A \ (A * x) ≈ x
            @test M \ (M * x) ≈ x
            @test A \ (A * d) ≈ d
            @test M \ (M * d) ≈ d
            @test A * (A \ x) ≈ x
            @test M * (M \ x) ≈ x
            @test A * (A \ d) ≈ d
            @test M * (M \ d) ≈ d
        end

        @testset "Testing `\\` with factorization" begin
            Af = factorize(A)
            Mf = factorize(M)
            @test Af \ (A * x) ≈ x
            @test Mf \ (M * x) ≈ x
            @test Af \ (A * d) ≈ d
            @test Mf \ (M * d) ≈ d
            @test A * (Af \ x) ≈ x
            @test M * (Mf \ x) ≈ x
            @test A * (Af \ d) ≈ d
            @test M * (Mf \ d) ≈ d
        end
    end

    @testset "Inplace factorization" begin
        Mf1 =factorize(M)
        Mf2 = factorize(M)
        factorize(Mf2, 2M)
        @test Mf2.Af == Mf1.Af
        @test Mf2.B == 2Mf1.B
        @test Mf2.C == 2Mf1.C
        @test Mf2.D == 2Mf1.D
        factorize(Mf2, 2M, update_factors=true)
        @test Mf2.Af ≠ Mf1.Af
        @test Mf2.B == 2Mf1.B
        @test Mf2.C == 2Mf1.C
        @test Mf2.D == 2Mf1.D
    end

    @testset "Testing sparse matrices" begin
        spA = sparse(A)  # │
        spB = sparse(B)  # │
        spC = sparse(C)  # ├─ sparse matrices
        spD = sparse(D)  # │
        spM = sparse(M)  # │

        @testset "Check that `ε * A` works" begin
            @test spM ≈ spA + ε₁*spB + ε₂*spC + ε₁ε₂*spD
        end

        @testset "Check that `\\` works without factorization" begin
            @test spA \ (spA * x) ≈ x
            @test spM \ (spM * x) ≈ x
            @test spA \ (spA * d) ≈ d
            @test spM \ (spM * d) ≈ d
            @test spA * (spA \ x) ≈ x
            @test spM * (spM \ x) ≈ x
            @test spA * (spA \ d) ≈ d
            @test spM * (spM \ d) ≈ d
        end

        @testset "Check that `\\` works with factorization" begin
            spAf = factorize(spA)
            spMf = factorize(spM)
            @test spAf \ (spA * x) ≈ x
            @test spMf \ (spM * x) ≈ x
            @test spAf \ (spA * d) ≈ d
            @test spMf \ (spM * d) ≈ d
            @test spA * (spAf \ x) ≈ x
            @test spM * (spMf \ x) ≈ x
            @test spA * (spAf \ d) ≈ d
            @test spM * (spMf \ d) ≈ d
        end
    end

    # TODO Add tests specific to `lu`, `qr`, `cholesky`, etc.

    @testset "Testing $f" for f in [:adjoint, :transpose]
        spA = sparse(A)  # │
        spB = sparse(B)  # │
        spC = sparse(C)  # ├─ sparse matrices
        spD = sparse(D)  # │
        spM = sparse(M)  # │

        # Adjoints / transposes
        x₂ = randn(n)                     # real-valued │
        y₂ = randn(n)                     # real-valued │
        z₂ = randn(n)                     # real-valued ├─ vector
        w₂ = randn(n)                     # real-valued │
        d₂ = x₂ + ε₁*y₂ + ε₂*z₂ + ε₁ε₂*w₂ # dual-valued │

        @eval begin
            x₂′ = $f($x₂)
            x′ = $f($x)
            d₂′ = $f($d₂)
            d′ = $f($d)
            A′  = $f($A)
            M′  = $f($M)
            spA′  = $f($spA)
            spM′  = $f($spM)
        end

        # Check that `\` works with adjoints
        @testset "Check that `\\` works with $f" begin
            @test (x₂′ * (A \ x)) ≈ (x′ * (A′ \ x₂))
            @test (d₂′ * (M \ d)) ≈ (d′ * (M′ \ d₂))
            @test (x₂′ * (spA \ x)) ≈ (x′ * (spA′ \ x₂))
            @test (d₂′ * (spM \ d)) ≈ (d′ * (spM′ \ d₂))
        end

        # Check that `\` works with adjoints of factorized versions
        Af = factorize(A)
        Mf = factorize(M)
        spAf = factorize(spA)
        spMf = factorize(spM)
        @eval begin
            Af′  = $f($Af)
            Mf′  = $f($Mf)
            spAf′  = $f($spAf)
            spMf′  = $f($spMf)
        end
        @testset "Check that factorized version with $f" begin
            @test (x₂′ * (Af \ x)) ≈ (x′ * (Af′ \ x₂))
            @test (d₂′ * (Mf \ d)) ≈ (d′ * (Mf′ \ d₂))
            @test (x₂′ * (spAf \ x)) ≈ (x′ * (spAf′ \ x₂))
            @test (d₂′ * (spMf \ d)) ≈ (d′ * (spMf′ \ d₂))
        end
    end

    @testset "Testing isapprox function" begin
        A2 = A .+ 0ε₁
        @test A2 ≈ A
        @test A ≈ A2
        @test (1.0 + 0ε₁) ≈ 1.0
        @test 1.0 ≈ (1.0 + 0ε₁)

        @test ~(M ≈ M .+ ε₁)
    end
end

println("HyperDualMatrixTools tests finished!")


using ApproxFunBase, ApproxFunOrthogonalPolynomials, LinearAlgebra, Test
    import ApproxFunBase: testbandedblockbandedoperator, testblockbandedoperator, testraggedbelowoperator, Block, ldiv_coefficients

@testset "PDE" begin
    @testset "Rectangle Laplace/Poisson" begin
        dx = dy = ChebyshevInterval()
        d = dx × dy
        g = Fun((x,y)->exp(x)*cos(y),∂(d))

        B = Dirichlet(d)

        testblockbandedoperator(B)
        testbandedblockbandedoperator(Laplacian(d))
        testbandedblockbandedoperator(Laplacian(d)+0.0I)

        A=[Dirichlet(d);Laplacian(d)]

        @time u=A\[g,0.]
        @test u(.1,.2) ≈ real(exp(0.1+0.2im))

        A=[Dirichlet(d);Laplacian(d)+0.0I]
        @time u=A\[g,0.]

        @test u(.1,.2) ≈ real(exp(0.1+0.2im))

        ## Poisson
        f=Fun((x,y)->exp(-10(x+.2)^2-20(y-.1)^2),ChebyshevInterval()^2,500)  #default is [-1,1]^2
        d=domain(f)
        A=[Dirichlet(d);Laplacian(d)]
        @time  u=\(A,[zeros(∂(d));f];tolerance=1E-7)
        @test ≈(u(.1,.2),-0.04251891975068446;atol=1E-5)
    end

    @testset "Bilaplacian" begin
        dx = dy = ChebyshevInterval()
        d = dx × dy
        Dx = Derivative(dx); Dy = Derivative(dy)
        L = Dx^4⊗I + 2*Dx^2⊗Dy^2 + I⊗Dy^4

        testbandedblockbandedoperator(L)

        B = Dirichlet(dx) ⊗ Operator(I,dy)
        testraggedbelowoperator(B)

        A=[Dirichlet(dx) ⊗ Operator(I,dy);
                Operator(I,dx)  ⊗ Dirichlet(dy);
                Neumann(dx) ⊗ Operator(I,dy);
                Operator(I,dx) ⊗ Neumann(dy);
                 L]

        testraggedbelowoperator(A)

        @time u=\(A,[[1,1],[1,1],[0,0],[0,0],0];tolerance=1E-5)
        @test u(0.1,0.2) ≈ 1.0
    end

    @testset "Schrodinger" begin
        dx=0..1; dt=0.0..0.001
        C=Conversion(Chebyshev(dx)*Ultraspherical(1,dt),Ultraspherical(2,dx)*Ultraspherical(1,dt))
        testbandedblockbandedoperator(C)
        testbandedblockbandedoperator(Operator{ComplexF64}(C))

        d = dx × dt

        x,y = Fun(d)
        @test x(0.1,0.0001) ≈ 0.1
        @test y(0.1,0.0001) ≈ 0.0001

        V = x^2

        Dt=Derivative(d,[0,1]);Dx=Derivative(d,[1,0])

        ϵ = 1.
        u0 = Fun(x->exp(-100*(x-.5)^2)*exp(-1/(5*ϵ)*log(2cosh(5*(x-.5)))),dx)
        L = ϵ*Dt+(.5im*ϵ^2*Dx^2)
        testbandedblockbandedoperator(L)

        @time u = \([timedirichlet(d);L],[u0,[0.,0.],0.];tolerance=1E-5)
        @test u(0.5,0.001) ≈ 0.857215539785593+0.08694948835021317im  # empircal from ≈ schurfact
    end

    @testset "Neumann" begin
        d=ChebyshevInterval()^2

        @time u=\([Neumann(d); Laplacian(d)-100.0I],[[[1,1],[1,1]],0.];tolerance=1E-12)
        @test u(.1,.9) ≈ 0.03679861429138079
    end

    @testset "Transport" begin
        dx=ChebyshevInterval(); dt=Interval(0,2.)
        d=dx × dt
        Dx=Derivative(d,[1,0]);Dt=Derivative(d,[0,1])
        x,y=Fun(identity,d)
        @time u=\([I⊗ldirichlet(dt);Dt+x*Dx],[Fun(x->exp(-20x^2),dx);0.];tolerance=1E-12)

        @test u(0.1,0.2) ≈ 0.8745340845783758  # empirical
    end

    @testset "Bilaplacian" begin
        dx=dy=ChebyshevInterval()
        d=dx × dy
        Dx=Derivative(dx);Dy=Derivative(dy)
        L=Dx^4⊗I+2*Dx^2⊗Dy^2+I⊗Dy^4

        testbandedblockbandedoperator(L)

        A=[ldirichlet(dx)⊗Operator(I,dy);
                rdirichlet(dx)⊗Operator(I,dy);
                Operator(I,dx)⊗ldirichlet(dy);
                Operator(I,dx)⊗rdirichlet(dy);
                lneumann(dx)⊗Operator(I,dy);
                rneumann(dx)⊗Operator(I,dy);
                Operator(I,dx)⊗lneumann(dy);
                Operator(I,dx)⊗rneumann(dy);
                L]

        # Checks bug in constructor
        f=Fun((x,y)->real(exp(x+1.0im*y)),component(rangespace(A)[1],1),22)
        @test f(-1.,0.1) ≈ real(exp(-1+0.1im))
        f=Fun((x,y)->real(exp(x+1.0im*y)),component(rangespace(A)[1],1))
        @test f(-1.,0.1) ≈ real(exp(-1+0.1im))

        F=[Fun((x,y)->real(exp(x+1.0im*y)),rangespace(A)[1]);
            Fun((x,y)->real(exp(x+1.0im*y)),rangespace(A)[2]);
            Fun((x,y)->real(exp(x+1.0im*y)),rangespace(A)[3]);
            Fun((x,y)->real(exp(x+1.0im*y)),rangespace(A)[4]);
            Fun((x,y)->real(exp(x+1.0im*y)),rangespace(A)[5]);
            Fun((x,y)->real(exp(x+1.0im*y)),rangespace(A)[6]);
            Fun((x,y)->-imag(exp(x+1.0im*y)),rangespace(A)[7]);
            Fun((x,y)->-imag(exp(x+1.0im*y)),rangespace(A)[8]);
            0]

        @time u=\(A,F;tolerance=1E-10)

        @test u(0.1,0.2)  ≈ exp(0.1)*cos(0.2)

        dx=dy=ChebyshevInterval()
        d=dx×dy
        Dx=Derivative(dx);Dy=Derivative(dy)
        L=Dx^4⊗I+2*Dx^2⊗Dy^2+I⊗Dy^4

        A=[(ldirichlet(dx)+lneumann(dx))⊗Operator(I,dy);
                (rdirichlet(dx)+rneumann(dx))⊗Operator(I,dy);
                Operator(I,dx)⊗(ldirichlet(dy)+lneumann(dy));
                Operator(I,dx)⊗(rdirichlet(dy)+rneumann(dy));
                (ldirichlet(dx)-lneumann(dx))⊗Operator(I,dy);
                (rdirichlet(dx)-rneumann(dx))⊗Operator(I,dy);
                Operator(I,dx)⊗(ldirichlet(dy)-lneumann(dy));
                Operator(I,dx)⊗(rdirichlet(dy)-rneumann(dy));
                L]

        u=\(A,[fill(1.0,8);0];tolerance=1E-5)
        @test u(0.1,0.2) ≈ 1.0

        F=[2Fun((x,y)->real(exp(x+1.0im*y)),rangespace(A)[1]);
            2Fun((x,y)->real(exp(x+1.0im*y)),rangespace(A)[2]);
            Fun((x,y)->real(exp(x+1.0im*y))-imag(exp(x+1.0im*y)),rangespace(A)[3]);
            Fun((x,y)->real(exp(x+1.0im*y))-imag(exp(x+1.0im*y)),rangespace(A)[4]);
            0;
            0;
            Fun((x,y)->real(exp(x+1.0im*y))+imag(exp(x+1.0im*y)),rangespace(A)[7]);
            Fun((x,y)->real(exp(x+1.0im*y))+imag(exp(x+1.0im*y)),rangespace(A)[8]);
            0]

        u=\(A,F;tolerance=1E-10)

        @test u(0.1,0.2)  ≈ exp(0.1)*cos(0.2)
    end

    @testset "Check resizing" begin
        d=ChebyshevInterval()^2
        A=[Dirichlet(d);Laplacian()+100I]
        QRR = qr(A)
        @time ApproxFunBase.resizedata!(QRR.R_cache,:,2000)
        @test norm(QRR.R_cache.data[1:200,1:200] - A[1:200,1:200]) ==0

        @time ApproxFunBase.resizedata!(QRR,:,200)
        j=56
        v=QRR.R_cache.op[1:100,j]
        @test norm(ldiv_coefficients(QRR.Q,v;maxlength=300)[j+1:end]) < 100eps()

        j=195
        v=QRR.R_cache.op[1:ApproxFunBase.colstop(QRR.R_cache.op,j),j]
        @test norm(ldiv_coefficients(QRR.Q,v;maxlength=1000)[j+1:end]) < 100eps()

        j=300
        v=QRR.R_cache.op[1:ApproxFunBase.colstop(QRR.R_cache.op,j),j]
        @test norm(ldiv_coefficients(QRR.Q,v;maxlength=1000)[j+1:end]) < j*20eps()

        @test ApproxFunBase.colstop(QRR.R_cache.op,195)-194 == ApproxFunBase.colstop(QRR.H,195)

        QR1 = qr(A)
        @time ApproxFunBase.resizedata!(QR1.R_cache,:,1000)
        QR2 = qr([Dirichlet(d);Laplacian()+100I])
        @time ApproxFunBase.resizedata!(QR2.R_cache,:,500)
        n=450;QR1.R_cache.data[1:n,1:n]-QR2.R_cache.data[1:n,1:n]|>norm
        @time ApproxFunBase.resizedata!(QR2.R_cache,:,1000)
        N=450;QR1.R_cache.data[1:N,1:N]-QR2.R_cache.data[1:N,1:N]|>norm
        N=1000;QR1.R_cache.data[1:N,1:N]-QR2.R_cache.data[1:N,1:N]|>norm

        QR1 = qr(A)
        @time ApproxFunBase.resizedata!(QR1,:,1000)
        QR2 = qr([Dirichlet(d);Laplacian()+100I])
        @time ApproxFunBase.resizedata!(QR2,:,500)
        @time ApproxFunBase.resizedata!(QR2,:,1000)

        @test norm(QR1.H[1:225,1:1000]-QR2.H[1:225,1:1000]) ≤ 10eps()

        QR1 = qr(A)
        @time ApproxFunBase.resizedata!(QR1,:,5000)
        @time u=\(QR1,[ones(∂(d));0.];tolerance=1E-7)

        @test norm((Dirichlet(d)*u-ones(∂(d))).coefficients) < 1E-7
        @test norm((A*u-Fun([ones(∂(d));0.])).coefficients) < 1E-7
        @test norm(((A*u)[2]-(Laplacian(space(u))+100I)*u).coefficients) < 1E-10
        @test eltype(ApproxFunBase.promotedomainspace(Laplacian(),space(u))) == Float64
        @test eltype(ApproxFunBase.promotedomainspace(Laplacian()+100I,space(u))) == Float64
        @test norm(((A*u)[2]-(Laplacian()+100I)*u).coefficients) < 1E-10
        @test norm((Laplacian()*u+100*u - (A*u)[2]).coefficients) < 10E-10
        @time v=\(A,[ones(∂(d));0.];tolerance=1E-7)
        @test norm((u-v).coefficients) < 100eps()

        @test u(0.1,1.) ≈ 1.0
        @test u(0.1,-1.) ≈ 1.0
        @test u(1.,0.1) ≈ 1.0
        @test u(-1.,0.1) ≈ 1.0
    end

    @testset "Chebyshev Dirichlet Kronecker" begin
        S=ChebyshevDirichlet()^2
        ff=(x,y)->exp(x)*cos(y)
        u=Fun(ff,S)

        for KO in [Operator(I,factor(S,1))⊗rdirichlet(factor(S,1)),
                    rdirichlet(factor(S,1))⊗Operator(I,factor(S,2))]
            testblockbandedoperator(KO)
            @test norm((KO*u-Fun(ff,rangespace(KO))).coefficients) ≤ 1E-10
        end
    end

    @testset "Small diffusion" begin
        dx=ChebyshevInterval();dt=Interval(0,0.2)
        d=dx×dt
        Dx=Derivative(d,[1,0]);Dt=Derivative(d,[0,1])
        x,t=Fun(dx×dt)

        B=0.0
        C=0.0
        V=B+C*x
        ε=0.1
        f=Fun(x->exp(-30x^2),dx)
        u=\([timedirichlet(d);Dt-ε*Dx^2-V*Dx],[f;zeros(3)];tolerance=1E-6)

        @test u(.1,.2) ≈ 0.496524222625512
        B=0.1
        C=0.2
        V=B+C*x
        u=\([timedirichlet(d);Dt-ε*Dx^2-V*Dx],[f;zeros(3)];tolerance=1E-7)
        @test u(.1,.2) ≈ 0.46810331039791464
    end

    @testset "concatenate InterlaceOperator" begin
        a=Fun(x -> 0 ≤ x ≤ 0.5 ? 0.5 : 1, Domain(-1..1) \ [0,0.5])
        @test a(0.1) == 0.5
        @test a(0.7) == 1.0
        s=space(a)
        # Bx=[ldirichlet(s);continuity(s,0)]
        # TODO: this should concat
        dt=Interval(0,2.)
        Dx=Derivative(s);Dt=Derivative(dt)
        Bx=[ldirichlet(s);continuity(s,0)]

        @test ApproxFunBase.rangetype(rangespace(continuity(s,0))) == Vector{Float64}
        @test ApproxFunBase.rangetype(rangespace(Bx)) == Vector{Any}
        @test ApproxFunBase.rangetype(rangespace(Bx⊗Operator(I,Chebyshev()))) == Vector{Any}

        rhs = Fun([0,[0,[0,0]],0],rangespace([I⊗ldirichlet(dt);Bx⊗I;I⊗Dt+(a*Dx)⊗I]))
        @test rhs(-0.5,0.0) == [0,[0,[0,0]],0]

        u=\([I⊗ldirichlet(dt);Bx⊗I;I⊗Dt+(a*Dx)⊗I],
            [Fun(x->exp(-20(x+0.5)^2),s),[[0],[0,0]],0.0];tolerance=1E-2)

        @test u(-0.4,0.1) ≈ u(-0.5,0.0) atol = 0.0001
    end
end


#=
renormalize!

getTensorM
truncMatrixU
matU2tenU
getNewTensorT
retateTensorT90
normalize
=#

function renormalize!{T}(lattice::Classical2dSquareLattice{T}, dimM::Int)
	tensorT, impurityTensorTtilde = getTensorT(lattice)

	tenmatM = getTensorM(tensorT)
	matrixMMd = tenmatM.matrix * transpose(tenmatM.matrix)
	Ul, lambdaVector, Uld = svd(matrixMMd)
	trunUl = truncMatrixU(Ul,dimM)
	tensorUl = matU2tenU(trunUl)

	# contract tensors to get the bigger tensorT
	newTensorT = getNewTensorT(tensorT,tensorUl)

	# rotate tensorT 90 degrees to the clockwise
	newRotTensorT = rotateTensorT90(newTensorT)

	# normalize tensorT
	newRotNormTensorT, normalizationFactor = normalize(newRotTensorT)

	### similar thing to do to tensorTtilde
	newTensorTtilde = getNewTensorT(tensorT,tensorUl,impurityTensorTtilde)
	newRotTensorTtilde = rotateTensorT90(newTensorTtilde)
	newRotNormTensorTtilde = newRotTensorTtilde ./normalizationFactor

	setTensorT!(lattice, newRotNormTensorT)
	setTensorT!(lattice, newRotNormTensorTtilde; tilde = true)
	return normalizationFactor
end

function getTensorM{T}(tensorT::Array{T,4} ; returntenmat = true, mirror = false)
	if mirror == false
		@tensor begin
			tensorM[x1,x2,x1p,x2p,y,yp] := tensorT[x1,x1p,y,j] * tensorT[x2,x2p,j,yp]
		end
	else
		@tensor begin
			tensorM[x1,x2p,x1p,x2,y,yp] := tensorT[x1,x1p,y,j] * tensorT[x2,x2p,yp,j]
		end
	end
	returntenmat ? answer = tensor2tenmat(tensorM,[1,2],[3,4,5,6]) : answer = tensorM
	return answer
end

function truncMatrixU(matrixU::Matrix, dimM::Int)
	# truncate the column dimension of matrix U to the dimM
	if size(matrixU,2) > dimM
		matrixU = Array(view(matrixU,:,1:dimM))
	end
	return matrixU
end

function matU2tenU(matrixU::Matrix)
	size1 = Int(sqrt(size(matrixU,1)))
	tenmatU = Tenmat(matrixU,[1,2],[3],(size1,size1,size(matrixU,2)))
	tensorU = tenmat2tensor(tenmatU)
	return tensorU
end

function getNewTensorT{T}(tensorT::Array{T,4},tensorU::Array{T,3}, tensorTtilde::Array{T,4}= tensorT; mirror = false)
	# contract U,T,T,U tensors to make new tensorT
	# according to the Fig.1.(b)
	# the third input tensorTtilde is for impurity tensor
	# 	without the third input, it uses tensorT twice
	if mirror == false
		@tensor begin
			UT[x2,x,x1p,y,i] := tensorU[x1,x2,x] * tensorTtilde[x1,x1p,y,i]
			UTT[x,x1p,y,x2p,yp] := UT[x2,x,x1p,y,i] * tensorT[x2,x2p,i,yp]
			newTensorT[x,xp,y,yp] := UTT[x,x1p,y,x2p,yp] * tensorU[x1p,x2p,xp]
		end
	else
		@tensor begin
			UT[x2,x,x1p,y,i] := tensorU[x1,x2,x] * tensorTtilde[x1,x1p,y,i]
			UTT[x,x1p,y,x2p,yp] := UT[x2,x,x1p,y,i] * tensorT[x2,x2p,yp,i]
			newTensorT[x,xp,y,yp] := UTT[x,x1p,y,x2p,yp] * tensorU[x1p,x2p,xp]
		end
	end
	return newTensorT
end


using Gomalish
using CxxWrap
using Test
using Random

for T in [Float64, Float32]
    @testset "dot $T" begin
        stdv1 = StdVector(T[1,2,3])
        stdv2 = StdVector(T[-1,-2,-3])
        mv1 = vector{T}(stdv1)
        mv2 = vector{T}(stdv2)
        print_all(mv1)
        print_all(mv2)
        @test Gomalish.dot(mv1, mv2) ≈ -T(14)
        @test mv1[1] ≈ stdv1[1]
        @test mv1[2] ≈ stdv1[2]
        @test mv1[3] ≈ stdv1[3]
    end
end

for T in [Float64, Float32]
    Random.seed!(12345)
    @testset "cg $T" begin
        coo = COO{T}("matrixfiles/sample.mtx")
        print_all(coo)
        A = CRS{T}(coo)
        jl_x = one(T) .+ rand(T, get_row(A))
        jl_b = one(T) .+ rand(T, get_row(A))
        x = vector{T}(StdVector(jl_x))
        b = vector{T}(StdVector(jl_b))

        solver = CG{CRS{T},T}()
        precond = Jacobi{CRS{T},T}()
        set_create_precond(solver, precond)
        set_apply_precond(solver, precond)
        set_tol(solver, 1.0e-12)
        set_maxiter(solver, get_row(A))
        print_all(A)
        println("---")
        print_all(x)
        println("---")
        print_all(b)

        if solver_check(solve(solver, A, x, b))
            return
        end
        jl_A = T[
            2 -1 0
            -1 2 -1
            0 -1 2
        ]
        println("Result by Gomalish/monolish")
        print_all(x)
        println("Result by Julia")
        jl_x = jl_A \ jl_b
        println(jl_x)
        @test x[1] ≈ jl_x[1]
        @test x[2] ≈ jl_x[2]
        @test x[3] ≈ jl_x[3]
    end
end

for T in [Float64, Float32]
    @testset "cg-impl $T" begin
        DIM = 100
    coo = tridiagonal_toeplitz_matrix(DIM, T(11.0), T(-1.0));
        A = CRS{T}(coo);

        println("===== Matrix informatrion =====")
        println("# of rows : ", get_row(A))
        println("# of cols : ", get_col(A))
        println("# of nnz  : ", get_nnz(A))

        # initial x is rand(0~1)
        x = vector{T}(get_row(A), T(0.0), T(1.0))
        # initial b is {1, 1, 1, ...,1}
        b = vector{T}(get_row(A), T(1.0))
        tol::T = 1.0e-12

        r = vector{T}(get_row(A), T(0.0))
        p = vector{T}(get_row(A), T(0.0))
        q = vector{T}(get_row(A), T(0.0))

        matvec(A, x, q)
        sub(b, q, r)
        Gomalish.copy(r, p)

        for iter in 1:get_row(A)
            matvec(A, p, q)
            tmp = Gomalish.dot(r, r)
            α = tmp / Gomalish.dot(p, q);
            axpy(α, p, x)
            axpy(-α, q, r)
            β = Gomalish.dot(r, r) / tmp
            xpay(β, r, p)
            resid = nrm2(r);
            println(iter + 1)
            @show resid
            (resid < tol) && return
            isnan(resid) && return
        end
    end
end

for T in [Float64, Float32]
    @testset "template solver $T" begin
        Random.seed!(12345)
        coo = COO{T}("matrixfiles/sample.mtx")
        A = CRS{T}(coo)
        jl_x = one(T) .+ rand(T, get_row(A))
        jl_b = one(T) .+ rand(T, get_row(A))
        x = vector{T}(StdVector(jl_x))
        b = vector{T}(StdVector(jl_b))
        solver = CG{CRS{T},T}()
        precond = Jacobi{CRS{T},T}()
        set_create_precond(solver, precond)
        set_apply_precond(solver, precond)
        set_tol(solver, 1.0e-12)
        set_maxiter(solver, get_row(A))
        if solver_check(solve(solver, A, x, b))
        return
        end
        print_all(x)
        jl_A = T[
            2 -1 0
            -1 2 -1
            0 -1 2
        ]
        println("Result by Gomalish/monolish")
        print_all(x)
        println("Result by Julia")
        jl_x = jl_A \ jl_b
        println(jl_x)
        @test x[1] ≈ jl_x[1]
        @test x[2] ≈ jl_x[2]
        @test x[3] ≈ jl_x[3]
    end
end

for T in [Float64, Float32]
    @testset "template solver part2 $T" begin
        Random.seed!(12345)
        coo = COO{T}("matrixfiles/sample.mtx")
        A = Dense{T}(coo)
        jl_x = one(T) .+ rand(T, get_row(A))
        jl_b = one(T) .+ rand(T, get_row(A))
        x = vector{T}(StdVector(jl_x))
        b = vector{T}(StdVector(jl_b))
        solver = BiCGSTAB{Dense{T},T}()
        precond = Gomalish.none{Dense{T},T}()
        set_create_precond(solver, precond)
        set_apply_precond(solver, precond)
        set_tol(solver, 1.0e-12)
        set_maxiter(solver, get_row(A))
        if solver_check(solve(solver, A, x, b))
        return
        end
        print_all(x)
        jl_A = T[
            2 -1 0
            -1 2 -1
            0 -1 2
        ]
        println("Result by Gomalish/monolish")
            print_all(x)
        println("Result by Julia")
        jl_x = jl_A \ jl_b
        println(jl_x)
        @test x[1] ≈ jl_x[1]
        @test x[2] ≈ jl_x[2]
        @test x[3] ≈ jl_x[3]
    end
end

for T in [Float64, Float32]
    @testset "template solver part3 $T" begin
        Random.seed!(12345)
        coo = COO{T}("matrixfiles/sample.mtx")
        A = Dense{T}(coo)
        jl_x = one(T) .+ rand(T, get_row(A))
        jl_b = one(T) .+ rand(T, get_row(A))
        x = vector{T}(StdVector(jl_x))
        b = vector{T}(StdVector(jl_b))
        solver = Gomalish.LU{Dense{T},T}()
        @test Gomalish.name(solver) == "monolish::equation::LU"
        precond = Gomalish.none{Dense{T},T}()
        set_create_precond(solver, precond)
        set_apply_precond(solver, precond)
        set_tol(solver, 1.0e-12)
        set_maxiter(solver, get_row(A))
        if solver_check(solve(solver, A, x, b))
            return
        end
        print_all(x)
        jl_A = T[
            2 -1 0
           -1 2 -1
            0 -1 2
        ]
        println("Result by Gomalish/monolish")
        print_all(x)
        println("Result by Julia")
        jl_x = jl_A \ jl_b
        println(jl_x)
        @test x[1] ≈ jl_x[1]
        @test x[2] ≈ jl_x[2]
        @test x[3] ≈ jl_x[3]
    end
end