#!/usr/bin/env python
# coding: utf-8

# # Orthogonal and QR Iteration
# 
# Follow a real $8\times8$ matrix through orthogonal iteration, Hessenberg reduction, and **complex Wilkinson-shifted QR** to an upper-triangular Schur form. Cells marked `#clear` contain short algorithms to type in class; surrounding instructions are supplied.

# In[1]:


import numpy as np
import numpy.linalg as la
import matplotlib.pyplot as pt

np.set_printoptions(precision=3, suppress=True, linewidth=100)
rng = np.random.default_rng(3)
n = 8
T0 = np.diag([8., 6., 4., 3., 3., 2., -1., .5])
T0 += .35 * np.triu(rng.normal(size=(n, n)), 1)
T0[3, 4], T0[4, 3] = -2., 2.             # One block with eigenvalues 3 +/- 2i.
U, _ = la.qr(rng.normal(size=(n, n)))
A = U @ T0 @ U.T
print("Eigenvalues:", np.sort_complex(la.eigvals(A)))


# The matrix is dense and nonsymmetric, with six real eigenvalues and one conjugate pair.
# The supplied plotter shows the numerical fill pattern (white means $|M_{ij}| \le 10^{-12}\|A\|_F$).

# In[2]:


def show_fill(frames):
    fig, axes = pt.subplots(1, len(frames), figsize=(2.5 * len(frames), 2.7), squeeze=False)
    for ax, (M, title) in zip(axes[0], frames):
        ax.imshow(np.abs(M) > 1e-12 * la.norm(A), cmap="Greys", vmin=0, vmax=1)
        ax.set(title=title, xticks=range(n), yticks=range(n))
    pt.tight_layout()
    pt.show()


# ## 1. Orthogonal iteration and QR iteration
# 
# Orthogonal iteration: $AQ_k=Q_{k+1}S_{k+1}$, starting with $Q_0=I$.
# QR iteration: $B_k=V_kR_k$, $B_{k+1}=R_kV_k$, starting with $B_0=A$.
# Track $W_{k+1}=W_kV_k$ so $B_k=W_k^TAW_k$.
# 
# **Sign convention:** NumPy's `la.qr` does *not* require positive diagonal entries in $R$.
# The two orthonormal bases agree up to column signs: $W_k=Q_kD_k$, where
# $D_k$ is diagonal with entries $\pm1$. Thus
# $B_k=D_k(Q_k^TAQ_k)D_k$; we adjust signs **only for the comparison**, not for either algorithm.

# In[3]:


Q_orth, B, W, k = np.eye(n), A.copy(), np.eye(n), 0
# Orthogonal iteration: QR-factor A @ Q_orth; keep its orthogonal factor.
# QR iteration: QR-factor B, form R @ V; also retain V to track the similarity.


# In[4]:


def orthogonal_iteration(A, Q):
    Q, _ = la.qr(A @ Q)
    return Q

def qr_iteration(B):
    V, R = la.qr(B)
    return R @ V, V


# Repeat the next cell with **Ctrl-Enter** to advance both algorithms.
# Compare the raw matrices and the comparison adjusted for column signs. Why can the first difference be large?

# In[5]:


Q_orth = orthogonal_iteration(A, Q_orth)
B, V = qr_iteration(B)
W = W @ V
k += 1
signs = np.where(np.sum(Q_orth * W, axis=0) < 0, -1., 1.)
B_orth = Q_orth.T @ A @ Q_orth
B_aligned = signs[:, None] * B_orth * signs[None, :]
print(f"Step {k}: raw difference {la.norm(B-B_orth)/la.norm(A):.2e}, "
      f"after sign alignment {la.norm(B-B_aligned)/la.norm(A):.2e}")
show_fill([(B, f"Unshifted QR: step {k}")])


# ## 2. Reduce to upper Hessenberg form
# 
# For column $j$, use a Householder reflector on rows $j+1,\ldots,n-1$.
# Apply it on both sides to preserve $H=Z^TAZ$. The supplied `House(x)` returns a unit
# vector $v$ defining $I-2vv^T$ (or $v=0$ for a zero input).
# Why do we preserve the first subdiagonal rather than eliminate it too?

# In[6]:


def House(x):
    """Unit vector v for the Householder reflector I - 2vv^T; zero if x is zero."""
    v = np.array(x, dtype=float, copy=True)
    normx = la.norm(v)
    if normx != 0:
        v[0] += np.copysign(normx, v[0])   # Choose sign to avoid cancellation.
        v /= la.norm(v)
    return v

H, Z, j = A.copy(), np.eye(n), 0
hess_frames = [(H.copy(), "Dense A")]
# Use House(H[j+1:, j]) and apply I - 2vv.T from both sides; also update Z.


# In[7]:


def hessenberg_step(H, Z, j):
    v = House(H[j+1:, j])
    H[j+1:, :] -= 2 * np.outer(v, v @ H[j+1:, :])
    H[:, j+1:] -= 2 * np.outer(H[:, j+1:] @ v, v)
    Z[:, j+1:] -= 2 * np.outer(Z[:, j+1:] @ v, v)
    H[j+2:, j] = 0.                     # Clear computed roundoff below the subdiagonal.
    return H, Z


# Run the next cell in place to eliminate one column at a time (six reflectors).
# 

# In[14]:


if j < n - 2:
    H, Z = hessenberg_step(H, Z, j)
    j += 1
    hess_frames.append((H.copy(), f"After reflector {j}"))
show_fill([hess_frames[-1]])


# In[15]:


# Finish any remaining columns; this also makes Run All work.
while j < n - 2:
    H, Z = hessenberg_step(H, Z, j)
    j += 1
    hess_frames.append((H.copy(), f"After reflector {j}"))
H0, Z0 = H.copy(), Z.copy()
show_fill([hess_frames[i] for i in (0, 1, 3, 6)])
print("Hessenberg similarity error:", la.norm(H - Z.T @ A @ Z) / la.norm(A))


# ## 3. Complex Wilkinson-shifted QR
# 
# For the active leading $m\times m$ Hessenberg matrix $X$, choose **Wilkinson's shift** $\mu$: the eigenvalue of the trailing $2\times2$ block closest to $x_{mm}$.
# With complex arithmetic, one QR step suffices even when $\mu$ is nonreal; no real double shift is needed.

# In[16]:


def wilkinson_shift(X):
    """Trailing 2 x 2 eigenvalue closest to the bottom-right entry."""
    C = X[-2:, -2:]
    roots = la.eigvals(C)
    return roots[np.argmin(np.abs(roots - C[-1, -1]))]


# Factor $X-\mu I=QR$ and update
# \[
# X^+=RQ+\mu I=Q^*XQ.
# \]
# To maintain $A=ZHZ^*$ after a deflation, also transform the upper-right coupling and accumulate $Z$. The explicit dense QR below is for teaching ($O(m^3)$ per step); an implicit Hessenberg QR step costs $O(m^2)$.

# In[17]:


def shifted_qr_step(H, Z, m):
    mu = wilkinson_shift(H[:m, :m])
    Q, R = la.qr(H[:m, :m] - mu * np.eye(m))
    H[:m, :m] = np.triu(R @ Q + mu * np.eye(m), -1)  # Clear roundoff below the subdiagonal.
    H[:m, m:] = Q.conj().T @ H[:m, m:]
    Z[:, :m] = Z[:, :m] @ Q
    return mu


# ### Stop when the trailing subdiagonal is small
# 
# For active size $m$, stop iterating on its bottom eigenvalue when
# \[
# |h_{m,m-1}|\le\max\left\{\tau\bigl(|h_{m-1,m-1}|+|h_{mm}|\bigr),\;\epsilon_{\rm mach}\|H\|_\infty\right\},\qquad\tau=10^{-12}.
# \]
# Then set that subdiagonal entry to zero, **deflate one eigenvalue** ($m\leftarrow m-1$), and continue. Stop when $m=1$; there are no remaining subdiagonal entries. A step cap guards against failure to converge.

# In[18]:


tol = 1e-12

def deflation_threshold(H, m):
    scale = abs(H[m-2, m-2]) + abs(H[m-1, m-1])
    return max(tol * scale, np.finfo(float).eps * la.norm(H, np.inf))

def advance(H, Z, m):
    if m <= 1:
        return m, "complete"
    if abs(H[m-1, m-2]) <= deflation_threshold(H, m):
        H[m-1, m-2] = 0.
        return m-1, "deflate 1 x 1"
    mu = shifted_qr_step(H, Z, m)
    return m, f"QR with shift {mu:.4g}"

H, Z, m = H0.astype(complex), Z0.astype(complex), n
qr_steps = 0
qr_frames = [(H.copy(), "Hessenberg")]


# Repeat the next cell with **Ctrl-Enter**. It either takes one Wilkinson-shifted QR step or deflates the bottom eigenvalue when its subdiagonal is small enough. Watch the magnitude and cutoff.

# In[19]:


if m > 1:
    old_m = m
    before = abs(H[m-1, m-2])
    cutoff = deflation_threshold(H, m)
    m, action = advance(H, Z, m)
    qr_steps += (m == old_m)
    qr_frames.append((H.copy(), f"{action}\nActive size {m}"))
    print(f"{action}; trailing |h|: {before:.2e} -> "
          f"{abs(H[old_m-1, old_m-2]):.2e} (cutoff {cutoff:.2e})")
show_fill([qr_frames[-1]])


# In[20]:


# Finish the QR iteration if the repeatable cell has not yet converged.
for _ in range(100):
    if m == 1:
        break
    old_m = m
    m, action = advance(H, Z, m)
    qr_steps += (m == old_m)
    qr_frames.append((H.copy(), f"{action}\nActive size {m}"))
else:
    raise RuntimeError("QR iteration did not converge in 100 further advances.")

show_fill([qr_frames[i] for i in np.linspace(0, len(qr_frames)-1, 6, dtype=int)])
print(f"Converged after {qr_steps} Wilkinson-shifted QR steps (plus deflations).")


# ## 4. Read the complex Schur form
# 
# The result $A=ZHZ^*$ has a **triangular** $H$; its diagonal entries are the eigenvalues, including the nonreal conjugate pair. Complex Schur vectors are generally not eigenvectors, but their leading groups span invariant subspaces.

# In[21]:


print("Complex Schur form:\n", H)
print("Eigenvalues from Schur form: ", np.sort_complex(np.diag(H)))
print("Eigenvalues of A:          ", np.sort_complex(la.eigvals(A)))
print("Largest subdiagonal entry:", np.max(abs(np.diag(H, -1))))
print("Relative reconstruction error:", la.norm(A - Z @ H @ Z.conj().T) / la.norm(A))
print("Unitarity error:", la.norm(Z.conj().T @ Z - np.eye(n)))


# Observe the two nonreal eigenvalues on the diagonal; no $2\times2$ real Schur block is required. The matrix may remain dense **above** the diagonal, but all entries below it converge to zero (to the chosen tolerance).