module TGST
# routine tests for the Symmetric Tridiagonal eigensolvers in GenericSchur

using LinearAlgebra
using Test
using GenericSchur
using GenericSchur: geigen!, geigvals!, QRIteration, DivideAndConquer

# some interesting examples from the LAPACK papers

function wilkinson(n, ::Type{T}) where {T}
    if mod(n, 2) == 1
        nh = (n - 1) >>> 1
        D = T.(vcat(nh:-1:1, [0], 1:nh))
    else
        D = abs.(T.(collect(-(n - 1):2:(n - 1))) / 2)
    end
    E = fill(T(1), n - 1)
    return D, E
end

function clement(n, ::Type{T}) where {T}
    D = zeros(T, n)
    E = [sqrt(T(i * (n + 1 - i))) for i in 1:(n - 1)]
    return D, E
end

function legendre(n, ::Type{T}) where {T}
    D = zeros(T, n)
    E = [T(i) / sqrt(T((2i - 1) * (2i) + 1)) for i in 2:n]
    return D, E
end

function laguerre(n, ::Type{T}) where {T}
    D = T.(collect(3:2:(2 * n + 1)))
    E = T.(collect(2:n))
    return D, E
end

function hermite(n, ::Type{T}) where {T}
    D = zeros(T, n)
    E = [sqrt(T(i)) for i in 1:(n - 1)]
    return D, E
end

# LAPACK-style classes
function st_testmat1(n, itype, ::Type{T} = Float64; dtype = :normal) where {T}
    dtype ∈ [:normal, :unif, :uniform] || throw(ArgumentError("incomprehensible dtype"))
    λ = zeros(T, n)
    ulp = eps(real(T))
    kpar = 1 / ulp
    if itype == 1
        λ[1] = 1
        for j in 2:n
            λ[j] = 1 / T(j)
        end
    elseif itype == 2
        λ[1:(n - 1)] .= 1
        λ[n] = T(1) / n
    elseif itype == 3
        for j in 1:n
            λ[j] = 1 / (kpar^((j - 1) / (n - 1)))
        end
    elseif itype == 4
        for j in 1:n
            λ[j] = T(1) - (T(j - 1) / T(n - 1)) * (1 - 1 / kpar)
        end
    elseif itype == 5
        λ .= exp.(-log(n) * rand(n))
    elseif itype == 6
        if dtype == :normal
            λ .= randn(n)
        else
            λ .= rand(dtype, n)
        end
    elseif itype == 7
        λ[1:(n - 1)] .= (1:(n - 1)) * ulp
        λ[n] = 1
    elseif itype == 8
        λ[1] = ulp
        for j in 2:(n - 1)
            λ[j] = 1 + sqrt(ulp) * j
        end
        λ[n] = 2
    elseif itype == 9
        λ[1] = 1
        for j in 2:n
            λ[j] = λ[j - 1] + 100 * ulp
        end
    end
    Q, r = qr(randn(n, n))
    Q = T.(Matrix(Q)) # until LinearAlgebra is fixed
    # FIXME: sign adjustment
    A = Q' * diagm(λ) * Q

    F = GenericSchur._hessenberg!(A)
    S = F.H
    return S, λ
end
const itypes1 = 1:9
# downselect so subsetting is reasonable
const itypes1a = [1, 3, 4, 5, 6]

function st_testmat2(n, itype, ::Type{T} = Float64; dtype = :normal) where {T <: Real}
    dtype ∈ [:normal, :unif, :uniform] || throw(ArgumentError("incomprehensible dtype"))
    λ = zeros(T, n)
    ulp = eps(real(T))
    if itype == 0
        S = SymTridiagonal(zeros(T, n), zeros(T, n - 1))
        λ = zeros(T, n)
    elseif itype == 1
        S = SymTridiagonal(ones(T, n), zeros(T, n - 1))
        λ = ones(T, n)
    elseif itype == 2
        S = SymTridiagonal(fill(T(2), n), ones(T, n - 1))
        λ = nothing
    elseif itype == 3
        S = SymTridiagonal(wilkinson(n, T)...)
        λ = nothing
    elseif itype == 4
        S = SymTridiagonal(clement(n, T)...)
        λ = T.(collect(-n:2:n))
    elseif itype == 5
        S = SymTridiagonal(legendre(n, T)...)
        λ = nothing
    elseif itype == 6
        S = SymTridiagonal(laguerre(n, T)...)
        λ = nothing
    elseif itype == 7
        S = SymTridiagonal(hermite(n, T)...)
        λ = nothing
    elseif itype == 8
        # exercise the scaling logic, all at once
        small = floatmin(T) / ulp
        large = floatmax(T) * ulp
        S1 = SymTridiagonal(wilkinson(n, T)...)
        S = gluemats(small * S1, S1, large * S1)
    else
        throw(ArgumentError("unimplemented itype"))
    end
    return S, λ
end
const itypes2 = 1:8
const broken_cases2 = [(BigFloat, QRIteration(), 8)]

function gluemats(mat1::SymTridiagonal{T}, mats...; gval = zero(T)) where {T}
    D = diag(mat1)
    E = diag(mat1, 1)
    for mat in mats
        append!(D, diag(mat))
        push!(E, gval)
        append!(E, diag(mat, 1))
    end
    return SymTridiagonal(D, E)
end

function batch(n, alg, ::Type{T} = Float64; thresh = 50, quiet = true) where {T}
    dmax = 0.0; vmax = 0.0
    for itype in itypes1
        @testset "class 1 type $itype" begin
            A, _ = st_testmat1(n, itype, T)
            de, ve = runtest(diag(A), diag(A, 1), alg)
            dmax = max(dmax, de)
            vmax = max(vmax, ve)
            @test de < thresh
            @test ve < thresh
        end
    end
    quiet || println("peak errors: $dmax $vmax")
    dmax = 0.0; vmax = 0.0
    for itype in itypes2
        @testset "class 2 type $itype" begin
            A, _ = st_testmat2(n, itype, T)
            de, ve = runtest(diag(A), diag(A, 1), alg)
            dmax = max(dmax, de)
            vmax = max(vmax, ve)
            if (T, alg, itype) in broken_cases2
                @test_broken de < thresh
            else
                @test de < thresh
            end
            @test ve < thresh
        end
    end
    quiet || println("peak errors: $dmax $vmax")
    return nothing
end


maxnorm(A) = maximum(abs.(vec(A)))

function runtest(D::AbstractVector{T}, E, alg; λ = nothing, V = nothing) where {T <: Real}
    if λ === nothing
        λ, V = geigen!(SymTridiagonal(copy(D), copy(E)), alg)
    end
    n = length(D)
    m = size(V, 2)
    ntol = min(m, n)
    myeps = eps(T)
    if m == n
        resnorm = maxnorm(V * Diagonal(λ) * V' - SymTridiagonal(D, E))
    else
        resnorm = maxnorm(Diagonal(λ) - V' * SymTridiagonal(D, E) * V)
    end
    Tnorm = max(opnorm(SymTridiagonal(D, E), 1), floatmin(T))
    if resnorm < Tnorm
        d_err = (resnorm / Tnorm) / (myeps * ntol)
    elseif Tnorm < 1
        d_err = (min(resnorm, Tnorm * ntol) / Tnorm) / (myeps * ntol)
    else
        d_err = min(resnorm / Tnorm, ntol) / (myeps * ntol)
    end
    if m == n
        v_err = maxnorm(V * V' - I) / (myeps * ntol)
    else
        v_err = maxnorm(V' * V - I) / (myeps * ntol)
    end
    return d_err, v_err
end

for T in [Float32, Float64, BigFloat], alg in [DivideAndConquer(), QRIteration()]
    @testset "SymTridiagonal $T $alg" begin
        batch(32, alg, T)
    end
end

end # module
