"""Numerical check of the identities for the Schafer--Wegner cycle (Section 4 and Appendix B).

Run: python3 rep_cycle_check.py   (numpy and scipy)

At random points (k = 2, random E in (-2, 2), s_E = (1 - E^2/4)^{1/2}) it checks:
  * the inequality |det A| >= det Im A for Im A > 0 (proof of l:rep-sdet);
  * l:rep-cycle(a): f_0(X) is block diagonal and commutes with X, g_X^2 > 0, det g_X = 1,
    g_X Lambda_B g_X = Lambda_B, and Tr f_0(X) >= ||X||;
  * l:rep-cycle(b): Im(Lambda_B (E - S^BB)) = s_E g_X^2, |det(E - S^BB)|^{-1} <= s_E^{-2k}, and
    ||(E - S^BB)^{-1}||_op <= s_E^{-2}(s_E + 2||X||_op);
  * l:rep-cycle(c): the identity e:rep-translation;
  * l:rep-cycle(d): Q_X^2 = Id and -i Lambda_B (f_0(X) + X) = -i s_E Q_X;
  * e:contour-edge-identity, and ||f_0(X) - f_0(X')|| <= ||X - X'||;
  * e:rep-saddle-inverse: (E - S)^{-1} = S for S = kappa_E - i s_E Q_X.
"""
import numpy as np
from scipy.linalg import sqrtm

rng = np.random.default_rng(3)
k = 2
n = 2 * k
Lam = np.diag([1.0] * k + [-1.0] * k).astype(complex)
I = np.eye(n)
hs = lambda A: np.sqrt(np.trace(A.conj().T @ A).real)
op = lambda A: np.linalg.norm(A, 2)
failures = []


def check(label, ok):
    if not ok:
        failures.append(label)


def offd(scale):
    B = scale * (rng.standard_normal((k, k)) + 1j * rng.standard_normal((k, k)))
    X = np.zeros((n, n), complex)
    X[:k, k:] = B
    X[k:, :k] = B.conj().T
    return X


def bd():
    A = rng.standard_normal((n, n)) + 1j * rng.standard_normal((n, n))
    P = (A + A.conj().T) / 2
    P[:k, k:] = 0
    P[k:, :k] = 0
    return P


TRIALS = 500
for trial in range(TRIALS):
    E = rng.uniform(-1.95, 1.95)
    s, kap = np.sqrt(1 - E**2 / 4), E / 2
    X, Xp, P = offd(rng.exponential(2)), offd(rng.exponential(2)), bd()
    f0 = lambda Z: sqrtm(s**2 * I + Z @ Z)
    F, Fp = f0(X), f0(Xp)
    g2, g2p = (F + X) / s, (Fp + Xp) / s
    g = sqrtm(g2)

    # |det A| >= det Im A
    A = rng.standard_normal((n, n)) + 1j * rng.standard_normal((n, n))
    B = rng.standard_normal((n, n)) + 1j * rng.standard_normal((n, n))
    A = (A + A.conj().T) / 2 + 1j * (B @ B.conj().T + 0.1 * I)
    ImA = (A - A.conj().T) / (2j)
    check("|det A| >= det Im A", abs(np.linalg.det(A)) >= np.linalg.det(ImA).real * (1 - 1e-10))

    # l:rep-cycle(a)
    check("(a) f0 block diagonal", np.allclose(F[:k, k:], 0) and np.allclose(F[k:, :k], 0))
    check("(a) f0 commutes with X", np.allclose(F @ X, X @ F))
    check("(a) g^2 > 0", np.linalg.eigvalsh((g2 + g2.conj().T) / 2).min() > 0)
    check("(a) det g = 1", np.isclose(np.linalg.det(g), 1))
    check("(a) g Lam g = Lam", np.allclose(g @ Lam @ g, Lam))
    check("(a) Tr f0 >= ||X||", np.trace(F).real >= hs(X) - 1e-10)

    # l:rep-cycle(b)
    SBB = P - 1j * Lam @ (F + X)
    M = Lam @ (E * I - SBB)
    check("(b) Im(Lam(E - S)) = s g^2", np.allclose((M - M.conj().T) / (2j), s * g2))
    check("(b) |det(E - S)|^-1 <= s^-2k", 1 / abs(np.linalg.det(E * I - SBB)) <= s ** (-2 * k) * (1 + 1e-10))
    check("(b) operator norm bound", op(np.linalg.inv(E * I - SBB)) <= s**-2 * (s + 2 * op(X)) * (1 + 1e-10))

    # l:rep-cycle(c), e:rep-translation
    eta, theta = rng.uniform(0, 1), rng.uniform(0, 1)
    R = SBB + 1j * theta * eta * Lam
    lhs = (-0.5 * np.trace(R @ R)).real
    rhs = -0.5 * hs(P) ** 2 - theta * eta * np.trace(F).real + k * (s**2 + theta**2 * eta**2)
    check("(c) e:rep-translation", np.isclose(lhs, rhs))

    # l:rep-cycle(d)
    Q = np.linalg.inv(g) @ Lam @ g
    check("(d) Q^2 = Id", np.allclose(Q @ Q, I))
    check("(d) -i Lam (f0 + X) = -i s Q", np.allclose(-1j * Lam @ (F + X), -1j * s * Q))

    # e:contour-edge-identity and the contraction
    lhs = hs(X - Xp) ** 2 - hs(F - Fp) ** 2
    rhs = s**2 * np.trace(g2 @ np.linalg.inv(g2p) + np.linalg.inv(g2) @ g2p - 2 * I).real
    check("e:contour-edge-identity", np.isclose(lhs, rhs, rtol=1e-8, atol=1e-8) and rhs >= -1e-9)
    check("f0 is a contraction", hs(F - Fp) <= hs(X - Xp) * (1 + 1e-10))

    # e:rep-saddle-inverse
    S = kap * I - 1j * s * Q
    check("e:rep-saddle-inverse", np.allclose(np.linalg.inv(E * I - S), S))

print(f"{TRIALS} random points")
print("all identities and bounds hold" if not failures else f"FAILED: {sorted(set(failures))}")