# 教材位置：第9章，第9.2节；图 9-1。
# 用同一个二维矩阵依次展示 V^T、Σ、U 的几何作用。
# A=[[2,1],[0,1]] 不是对称矩阵；不可把特征分解直接当作 SVD。
# A、U、Vh、Sigma 均为 (2,2)，s 为 (2,)，circle 为 (2,361)。
# 坐标、矩阵元素均无量纲；Σ 是奇异值对角矩阵，∑ 才是求和符号。
# 正交矩阵保持长度和夹角，也可能包含反射，不一定是旋转。
# 固定矩阵无随机抽样；输出奇异值、正交因子行列式及重建误差。
# 执行 python e39_svd2d.py 生成四阶段中文图 svd2d.png。

import numpy as np
import matplotlib.pyplot as plt
from plot_style import configure_chinese_plot

A = np.array([[2.0, 1.0], [0.0, 1.0]])
U, s, Vh = np.linalg.svd(A, full_matrices=False)
Sigma = np.diag(s)
# 手算 A^T A=[[4,2],[2,2]]，特征方程 λ²−6λ+4=0。
# 特征值是 3±sqrt(5)，取非负平方根得到降序奇异值。
manual_s = np.sqrt([3 + np.sqrt(5), 3 - np.sqrt(5)])
assert np.allclose(s, manual_s)
assert np.allclose(A.T @ A, [[4, 2], [2, 2]])
assert np.allclose(U.T @ U, np.eye(2))
assert np.allclose(Vh @ Vh.T, np.eye(2))
assert np.allclose(U @ Sigma @ Vh, A)
# NumPy 的 Vh 存放 V 的共轭转置；本例全为实数，所以 Vh=V^T。
# 奇异向量正负可同时翻转；只核对重建和对应方向关系。
assert np.allclose(A @ Vh.T, U * s)
print("手算、数值奇异值：", manual_s, s)
print("正交因子 U 与 V^T 的行列式：", np.linalg.det(U), np.linalg.det(Vh))
print("矩阵重建误差：", np.linalg.norm(U @ Sigma @ Vh - A))

# 所有圆周点作为列向量存放，左乘矩阵同时映射全部点。
theta = np.linspace(0, 2 * np.pi, 361)
circle = np.vstack([np.cos(theta), np.sin(theta)])
after_v = Vh @ circle
after_sigma = Sigma @ after_v
after_u = U @ after_sigma
assert np.allclose(after_u, A @ circle)
assert np.allclose(np.linalg.norm(after_v, axis=0), 1)
# 追踪同一批输入基向量的像，而非每幅图重新画标准基。
stage_matrices = [np.eye(2), Vh, Sigma @ Vh, A]
stage_points = [circle, after_v, after_sigma, after_u]
titles = ["输入单位圆", "施加 V^T：正交变换", "施加 Σ：伸缩", "施加 U：正交变换"]
# 四个面板使用同样的坐标范围和等比例单位，圆与椭圆形状可直接比较。
# 坐标基不同于奇异向量；基向量箭头只是追踪映射的辅助标记。

configure_chinese_plot()
fig, axes = plt.subplots(2, 2, figsize=(6.3, 5.6), constrained_layout=True)
for ax, points, matrix, title in zip(
    axes.flat, stage_points, stage_matrices, titles
):
    ax.plot(*points, color="tab:blue", lw=1.5)
    # 橙点和短箭头追踪输入圆上的同一个起点及前进方向。
    ax.scatter(*points[:, 0], color="tab:orange", zorder=4, s=20)
    ax.annotate("", xy=points[:, 30], xytext=points[:, 15],
                arrowprops={"arrowstyle": "->", "color": "tab:orange"})
    for column, color in zip(matrix.T, ["tab:green", "tab:red"]):
        ax.annotate("", xy=column, xytext=(0, 0),
                    arrowprops={"arrowstyle": "->", "color": color})
    ax.axhline(0, color="gray", lw=0.5)
    ax.axvline(0, color="gray", lw=0.5)
    ax.set(xlim=(-2.7, 2.7), ylim=(-2.7, 2.7),
           xlabel="第一坐标", ylabel="第二坐标", title=title)
    ax.set_aspect("equal")
    ax.set_xticks([-2, 0, 2])
    ax.set_yticks([-2, 0, 2])
    ax.grid(alpha=0.2)
# 颜色约定在源码中明确：绿、红箭头分别为第一、第二输入基向量的像。
# 末图应与直接计算 A@circle 相同，不能另换矩阵来追求视觉效果。
fig.savefig("svd2d.png", dpi=180, bbox_inches="tight")
plt.close(fig)