using Krylov, LinearOperators
using LinearAlgebra, Printf, SparseArrays

# Identity matrix.
eye(n::Int) = sparse(1.0 * I, n, n)

# Symmetric quasi-definite systems and variants
n = m = 5
A = [2^(i/j)*j + (-1)^(i-j) * n*(i-1) for i = 1:n, j = 1:n]
b = ones(n)
M = diagm(0 => [3.0 * i for i = 1:n])
N = diagm(0 => [5.0 * i for i = 1:n])
c = -b

# [I   A] [x] = [b]
# [Aᴴ -I] [y]   [c]
(x, y, stats) = tricg(A, b, c)
K = [eye(m) A; A' -eye(n)]
B = [b; c]
r = B - K * [x; y]
resid = norm(r)
@printf("TriCG: Relative residual: %8.1e\n", resid)

# [-I   A] [x] = [b]
# [ Aᴴ  I] [y]   [c]
(x, y, stats) = tricg(A, b, c, flip=true)
K = [-eye(m) A; A' eye(n)]
B = [b; c]
r = B - K * [x; y]
resid = norm(r)
@printf("TriCG: Relative residual: %8.1e\n", resid)

# [I   A] [x] = [b]
# [Aᴴ  I] [y]   [c]
(x, y, stats) = tricg(A, b, c, spd=true)
K = [eye(m) A; A' eye(n)]
B = [b; c]
r = B - K * [x; y]
resid = norm(r)
@printf("TriCG: Relative residual: %8.1e\n", resid)

# [-I    A] [x] = [b]
# [ Aᴴ  -I] [y]   [c]
(x, y, stats) = tricg(A, b, c, snd=true)
K = [-eye(m) A; A' -eye(n)]
B = [b; c]
r = B - K * [x; y]
resid = norm(r)
@printf("TriCG: Relative residual: %8.1e\n", resid)

# [τI    A] [x] = [b]
# [ Aᴴ  νI] [y]   [c]
(τ, ν) = (1e-4, 1e2)
(x, y, stats) = tricg(A, b, c, τ=τ, ν=ν)
K = [τ*eye(m) A; A' ν*eye(n)]
B = [b; c]
r = B - K * [x; y]
resid = norm(r)
@printf("TriCG: Relative residual: %8.1e\n", resid)

# [M⁻¹  A  ] [x] = [b]
# [Aᴴ  -N⁻¹] [y]   [c]
(x, y, stats) = tricg(A, b, c, M=M, N=N, verbose=1)
K = [inv(M) A; A' -inv(N)]
H = BlockDiagonalOperator(M, N)
B = [b; c]
r = B - K * [x; y]
resid = sqrt(dot(r, H * r))
@printf("TriCG: Relative residual: %8.1e\n", resid)

TriCG: Relative residual:  2.7e-11
TriCG: Relative residual:  7.8e-12
TriCG: Relative residual:  1.2e-11
TriCG: Relative residual:  4.1e-11
TriCG: Relative residual:  7.0e-10
TriCG: system of 10 equations in 10 variables
    k     ‖rₖ‖     βₖ₊₁     γₖ₊₁  timer
    0  1.1e+01  6.7e+00  8.7e+00  0.31s
    1  5.8e+00  2.6e+02  3.0e+02  0.46s
    2  2.5e+00  2.9e+02  2.2e+02  0.46s
    3  1.2e+00  2.4e+01  5.5e+01  0.46s
    4  1.9e-01  3.4e-01  8.1e-01  0.46s
    5  1.2e-08  6.7e-08  7.7e-08  0.46s

TriCG: Relative residual:  1.2e-08