# # Combining paradigms: the algebraic Riccati equation
#
# The continuous algebraic Riccati equation (CARE)
# $$A^T X + X A - X B B^T X + Q = 0$$
# is solved every day in floating point (`care` in Matlab, Kleinman's iteration). We
# show three ways of using exact and certified tools around such a solver:
#
# 1. *certify* the floating-point solution: a Kantorovich test proves that an exact
#    solution lies in a tiny ball around it;
# 2. *refine* it to any precision by interval Newton;
# 3. *solve exactly*: compute all the real solutions, and prove that the one found
#    is the only stabilizing one.
#
# The example: the system $\dot x = Ax + Bu$ with one input, where
# $$A = \begin{pmatrix} 0 & 1\\ 2 & -1\end{pmatrix}, \qquad B = \begin{pmatrix} 0\\ 1\end{pmatrix}.$$
# Without control it is unstable: the eigenvalues of $A$ are $1$ and $-2$. We look for
# the control $u$ that minimizes $\int_0^\infty (x^Tx + u^2)\,dt$ (weights $Q = I$ on the
# state, $R = 1$ on the input). It is $u = -B^TXx$, where $X$ is the solution of the
# CARE that makes the closed loop $A - BB^TX$ stable.

#nb # Notebook only: Nemo and AbstractAlgebra objects are shown as plain text
#nb # (their HTML and LaTeX output throws on purpose, which IJulia 1.34 no longer catches).
#nb filter!(m -> !(m == MIME("text/latex") || (m isa AbstractVector && MIME("text/html") in m)), IJulia.ijulia_mime_types)

using LinearAlgebra, AbstractAlgebra, PACE, PaceControl

A = [0 1; 2 -1]
B = [0; 1;;]
Q = [1 0; 0 1]

# ## 1. Floating point: Kleinman's iteration
#
# Kleinman's method improves a feedback $u = -Kx$ step by step. A *gain* is the row
# vector $K$ of this feedback; it is *stabilizing* when the closed loop $\dot x =
# (A-BK)x$ is stable, that is when the eigenvalues of $A - BK$ have negative real
# parts. The method needs one to start. Here $K_0 = (3\ \ 1)$ gives
# $A - BK_0 = \begin{pmatrix} 0 & 1\\ -1 & -2\end{pmatrix}$, of characteristic polynomial
# $(s+1)^2$: both eigenvalues are $-1$.
#
# At each step, with $A_k = A - BK_k$ (the closed loop), the feedback $u = -K_kx$ gives
# $\dot x = A_kx$ and the cost $\int_0^\infty x^T(Q + K_k^TK_k)x\,dt = x_0^TX_{k+1}x_0$ with
# $X_{k+1} = \int_0^\infty e^{A_k^Tt}(Q+K_k^TK_k)e^{A_kt}\,dt$, which converges because $A_k$
# is stable. Differentiating under the integral, $X_{k+1}$ is the solution of the
# Lyapunov equation
# $$A_k^T X_{k+1} + X_{k+1}A_k + Q + K_k^T K_k = 0,$$
# which is *linear* in $X_{k+1}$ (no $XBB^TX$ term): $3$ equations in $3$ unknowns here,
# solved below in the form $(I\otimes A_k^T + A_k^T\otimes I)\,\mathrm{vec}(X) =
# -\mathrm{vec}(Q + K_k^TK_k)$. The next gain is $K_{k+1} = B^T X_{k+1}$. It is again
# stabilizing: by a theorem of Kleinman (1968), when $K_0$ is stabilizing and $Q$ is
# positive definite, every $K_k$ is stabilizing, and the iterates converge
# quadratically to the stabilizing
# solution of the CARE. On the example, the first step gives, exactly,
# $X_1 = \begin{pmatrix}10 & 5\\ 5 & 3\end{pmatrix}$ and $K_1 = (5\ \ 3)$, and $A - BK_1$ has
# characteristic polynomial $(s+1)(s+3)$.

function kleinman(A, B, Q, K; iters = 30)
    n = size(A, 1)
    X = zeros(n, n)
    for _ in 1:iters
        Ak = A - B * K
        L = kron(I(n), Ak') + kron(Ak', I(n))
        X = reshape(-(L \ vec(Q + K' * K)), n, n)
        K = B' * X
    end
    X
end

K0 = [3.0 1.0]                       ## closed loop (s+1)^2
Xf = kleinman(Float64.(A), Float64.(B), Float64.(Q), K0)

# The residual is at the level of rounding errors:

norm(A' * Xf + Xf * A - Xf * B * B' * Xf + Q)

# A small residual is not a proof: it says that $X_f$ solves a nearby equation, not
# that an exact solution is nearby.

# ## 2. The CARE as a polynomial system
#
# With $X = \begin{pmatrix} x_1 & x_2\\ x_2 & x_3\end{pmatrix}$, the CARE is a system
# of three quadratic equations in three unknowns, with rational coefficients.

R, (x1, x2, x3) = polynomial_ring(QQ, ["x1", "x2", "x3"])
X = [x1 x2; x2 x3]
E = A' * X + X * A - X * B * B' * X + Q
sys = [E[1, 1], E[1, 2], E[2, 2]]

# ## 3. Certify: the Kantorovich test
#
# At the point $X_f$, LACE computes in interval arithmetic the constants of
# Kantorovich's theorem. If $C \le 1$, there is exactly one solution of the system in
# the ball of radius $2B_0$ around the point.

pt = BigFloat.([Xf[1, 1], Xf[1, 2], Xf[2, 2]])
r, C = kantorovich_test(sys, pt, 63, 63, 30)

# $C$ is far below $1$: the floating-point solution is within about $4\cdot 10^{-15}$ of an
# exact solution, and this is now a theorem. The test costs one evaluation of the
# system, of its Jacobian and of a bound on its second derivatives at one point.

# ## 4. Refine: interval Newton
#
# Interval Newton contracts a *box* known to contain the solution. The Kantorovich
# ball gives one: the box of half-width $2B_0$ around the point. (Started from the
# point itself, the iteration only computes a Newton step, which is not an enclosure.)
# Here 200 bits, with a working precision of 256 bits.

box = [BigInterval(x - right(r), x + right(r)) for x in pt]
status, Xbox, = interval_newton(sys, box, 256, 256, 200, 200, 10; refine = Int32(1), verbose = Int32(0))
status, [Float64(log2(diam(b))) for b in Xbox]

# The status `1` means that the box is proved to contain a zero (some Newton step
# mapped the box into its interior) and that the target precision is reached. The
# exact solution is
# known here, $x_2 = 2+\sqrt5$, $x_3 = \sqrt5$, $x_1 = x_2x_3 + x_2 - 2x_3 = 7+\sqrt5$, and
# it is in the boxes:

exact = setprecision(BigFloat, 512) do
    r5 = sqrt(BigFloat(5))
    [7 + r5, 2 + r5, r5]
end
all(is_inside.(exact, Xbox))

# ## 5. Solve exactly: all the solutions
#
# The Kantorovich test is local: it says nothing about the other solutions. Exact
# solving (a rational univariate representation, then certified isolation, in LACE)
# gives all the real solutions in boxes:

sols = isolate(sys, symbols(R); output_precision = 30)

# The closed loop $A - BB^TX$ has characteristic polynomial $s^2 + (1+x_3)s + (x_2-2)$.
# Its coefficients are known as intervals at every solution. On interval coefficients,
# `hurwitz_stable` applies Kharitonov's theorem (worksheet 1), which decides which
# solutions are stabilizing:

closed_loop(x) = [one(x[3]), 1 + x[3], x[2] - 2]          ## highest degree first
stabilizing = [hurwitz_stable(closed_loop(x)) for x in sols]

# Exactly one solution is stabilizing, and it is the one certified in step 3: its box
# lies in the Kantorovich ball of radius $2B_0$ around the floating-point point.

xs = sols[findfirst(stabilizing)]
all(abs(xs[i] - pt[i]) < r for i in eachindex(pt))

# ## 6. Scaling
#
# How do the certificate (step 3) and exact solving (step 5) behave when the size
# grows? We build larger systems in the same way as the example, $\dot x = Ax + Bu$
# with one input, of dimension $n$:
#
# - $A$ has ones above the diagonal, last row $(-a_1, \dots, -a_n)$, and zeros
#   elsewhere, so that its characteristic polynomial is
#   $s^n + a_n s^{n-1} + \dots + a_2 s + a_1$; we take $a_k = (-1)^k(1 + k \bmod 3)$,
#   arbitrary values that make $A$ unstable ($a_1 = -2 < 0$);
# - $B = (0, \dots, 0, 1)^T$, $Q = I$, $R = 1$;
# - the initial gain $K_0$ is chosen so that all the eigenvalues of $A - BK_0$ are $-1$.
#
# The unknown $X$ is a symmetric $n\times n$ matrix: $n(n+1)/2$ unknowns, and as many
# quadratic equations. For each $n$ the table gives the Kantorovich constant `C` at
# the Kleinman solution (certified when $C \le 1$), the time of Kleinman plus the
# test (`t_numeric_certified`, in seconds), the number of real solutions of the CARE
# and the time to compute them all (`t_exact`).

function care(a)
    n = length(a)
    A = zeros(Rational{BigInt}, n, n)
    for i in 1:n-1
        A[i, i+1] = 1
    end
    A[n, :] = -a
    B = zeros(Rational{BigInt}, n, 1)
    B[n, 1] = 1
    R, xs = polynomial_ring(QQ, ["x$(i)_$(j)" for i in 1:n for j in i:n])
    X = Matrix{eltype(xs)}(undef, n, n)
    k = 0
    for i in 1:n, j in i:n
        k += 1
        X[i, j] = X[j, i] = xs[k]
    end
    E = A' * X + X * A - X * B * B' * X + I(n)
    [E[i, j] for i in 1:n for j in i:n], R, Float64.(A), Float64.(B)
end

function compare(n)
    a = Rational{BigInt}[(-1)^k * (1 + k % 3) for k in 1:n]
    sys, R, Af, Bf = care(a)
    K0 = Float64.([binomial(n, k) for k in n:-1:1])' .- Float64.(a')
    t_num = @elapsed begin
        Xf = kleinman(Af, Bf, Matrix(1.0I, n, n), K0)
        pt = BigFloat.([Xf[i, j] for i in 1:n for j in i:n])
        r, C = kantorovich_test(sys, pt, 63, 63, 10)
    end
    t_exact = @elapsed sols = isolate(sys, symbols(R); output_precision = 20)
    (n = n, unknowns = length(sys), C = Float64(right(C)), t_numeric_certified = t_num,
     real_solutions = length(sols), t_exact = t_exact)
end

compare(2);   ## compilation
[compare(n) for n in 2:9]

# On this structured problem, exact solving of the whole system is competitive with
# the certification of a single solution up to $n \approx 8$: the Kantorovich test
# bounds second derivatives over $n(n+1)/2$ unknowns, and its cost grows quickly too.
# Beyond, exact solving becomes the more expensive one: $5$ s for $n = 10$ on a laptop,
# eight times more than for $n = 9$, while the number of real solutions doubles every two
# orders. The constant $C$ also grows with $n$ and would eventually exceed $1$ at 63
# bits, which calls for a more precise point (interval Newton) or the preconditioned
# test (`flags = KANTO_PRECOND`).

# ## 7. A preconditioned Kantorovich test
#
# Beyond $n = 11$, the test fails at the Kleinman point: $C > 1$. The point is not worse
# than before, but the constant of the test depends on how the equations are scaled.
# With `flags = KANTO_PRECOND`, LACE applies the test to the system multiplied by an
# approximate inverse of its Jacobian matrix at the point, which does not change its
# solutions; the second derivatives of the equations are then combined before taking
# absolute values, so that their cancellations are kept. For each $n$: $C$ and the time
# of the test, without and with
# preconditioning (only the numerical part, no exact solving).

function kanto_compare(n)
    a = Rational{BigInt}[(-1)^k * (1 + k % 3) for k in 1:n]
    sys, R, Af, Bf = care(a)
    K0 = Float64.([binomial(n, k) for k in n:-1:1])' .- Float64.(a')
    Xf = kleinman(Af, Bf, Matrix(1.0I, n, n), K0)
    pt = BigFloat.([Xf[i, j] for i in 1:n for j in i:n])
    t0 = @elapsed (_, C0) = kantorovich_test(sys, pt, 63, 63, 10)
    t1 = @elapsed (_, C1) = kantorovich_test(sys, pt, 63, 63, 10;
                                         flags = PACE.KANTO_PRECOND)
    (n = n, C = Float64(right(C0)), t = t0,
     C_precond = Float64(right(C1)), t_precond = t1)
end

[kanto_compare(n) for n in 10:14]

# Without preconditioning the test fails from $n = 12$ ($C = 16$); with it, the same
# points are certified up to $n = 13$ ($C \approx 0.13$), for about three times the cost
# of the test. At $n = 14$ both fail: the point itself must then be refined.

# ## Summary
#
# | Step | Tool | What it gives |
# |---|---|---|
# | Kleinman in `Float64` | LinearAlgebra | a fast approximation, no guarantee |
# | Kantorovich test | LACE | a proof that an exact solution is within $2B_0$ |
# | Interval Newton | LACE | that solution to any precision |
# | Exact solving + Kharitonov | LACE, PACE | all solutions, and uniqueness of the stabilizing one |
#
# The certificate is local and its cost is polynomial in the size; exact solving is
# global and its cost grows with the number of solutions. Use the first to trust a
# numerical result, the second to know that nothing else exists.