//  [python]
N = 50
H = torch.zeros(N, N)
for i in range(N):
    H[max(i - 1, 0), i] = 0.5
    H[i, max(i - 1, 0)] = 0.5
    H[i, i] = 1
# H = torch.randn(N, N)
# H = H@H.t()/N
I = torch.eye(N)
for count, eps in enumerate([0, 1e-2]):
    if count == 0:
        ax = ax1
        num_iterations = 150000
    else:
        ax = ax2
        num_iterations = 150000

    if eps == 0:
        H1 = H.clone()
    else:
        L, U = torch.linalg.eigh((H @ H.t() + eps**2 * I).to(torch.float64))
        H1 = (U @ torch.diag(torch.sqrt(L)) @ U.t()).to(torch.float32)

    # fitting on Lie group
    Q, invQ = torch.eye(N), torch.eye(N)
    if eps == 0:
        step = 1.0
    else:
        step = 0.1
    Loss = []
    for i in range(num_iterations):
        loss = torch.linalg.matrix_norm(Q.t() @ Q @ H1 - I)
        Loss.append(loss.item())
        
        v = torch.randn(N, 1)
        h = H @ v + eps * torch.randn(N, 1)

        psgd.update_precond_newton_math_(Q, invQ, v, h, step, "2nd", 0.0)

    ax.semilogy(range(1, 1 + num_iterations), Loss, "k")
    
    # fitting on Lie group
    Q, invQ = torch.eye(N), None
    if eps == 0:
        step = 1.0
    else:
        step = 0.1
    Loss = []
    for i in range(num_iterations):
        loss = torch.linalg.matrix_norm(Q.t() @ Q @ H1 - I)
        Loss.append(loss.item())
        
        v = torch.randn(N, 1)
        h = H @ v + eps * torch.randn(N, 1)

        psgd.update_precond_newton_math_(Q, invQ, v, h, step, "2nd", 0.0)

    ax.semilogy(range(1, 1 + num_iterations), Loss, "b")

    # closed-form solution
    Loss = []
    hh = torch.eye(N)
    for i in range(num_iterations):
        L, U = torch.linalg.eigh(hh)
        # L[L < 1e-6] = 1e-6
        P = U @ torch.diag(torch.rsqrt(L)) @ U.t()
        loss = torch.linalg.matrix_norm(P @ H1 - I)
        Loss.append(loss.item())
        
        v = torch.randn(N, 1)
        h = H @ v + eps * torch.randn(N, 1)

        if (i + 1) / (i + 2) < 0.999:
            hh = (i + 1) / (i + 2) * hh + 1 / (i + 2) * (h @ h.t())
        else:
            hh = 0.999 * hh + 0.001 * (h @ h.t())

    ax.semilogy(range(1, 1 + num_iterations), Loss, "r")

    # bfgs
    Loss = []
    P = torch.eye(N)
    for i in range(num_iterations):
        loss = torch.linalg.matrix_norm(P @ H1 - I)
        Loss.append(loss.item())
        
        v = torch.randn(N, 1)
        h = H @ v + eps * torch.randn(N, 1)
        if v.t() @ h < 0:
            h = -h  # to avoid P<0

        P = (
            P
            + (v.t() @ h + h.t() @ P @ h) * (v @ v.t()) / (h.t() @ v) ** 2
            - (P @ h @ v.t() + v @ h.t() @ P) / (v.t() @ h)
        )

    ax.semilogy(range(1, 1 + num_iterations), Loss, "m")

    ax.legend(
        [
            r"PSGD, GL$(n,\mathbb{R})$",
            r"PSGD, Triangular",
            r"$P=(E[hh^T])^{-0.5}$",
            "BFGS",
        ],
        fontsize=7,
    )
    ax.set_xlabel("Iterations", fontsize=8)
    ax.tick_params(labelsize=6)
    if count == 0:
        ax.set_ylabel(r"$||PH' - I||_F$", fontsize=8)
        ax.set_title(r"(a) Clean $Hv$", fontsize=8)
    else:
        # ax.set_ylabel(r"$||PH' - I||_F$", fontsize=8)
        ax.set_title(r"(b) Noisy $Hv$", fontsize=8)