"""Numerical check of the real parts along the deformation (Appendix B and p:rep-contour).

Run: python3 rep_homotopy_check.py   (numpy and scipy)

With R_t(P, X) = P - i Lambda_B ((1 - t)(f_0(X) - c) + X) as in e:contour-homotopy (k = 2,
s = s_E for a random E in (-2, 2)), it checks at random points:
  * the identity e:contour-vertex and the bound e:contour-vertex-bound of l:contour-real(a);
  * the identity e:contour-edge and its bound -(w/2)||P - P'||^2 (l:contour-real(b));
  * the bound of l:contour-real(c) for the link to a point R_* of the cycle;
  * the per-site real part displayed after p:rep-contour (m = 1, c = eta, eta' = 0), and the
    per-site bound in e:rep-homotopy with the constant of the proof of p:rep-contour(b),
    C_{k,eta}(1 + |zeta|^4) = k(s_E^2 + eta^2) + (s_E + eta)|zeta|^2 + |zeta|^4 + k eta s_E/2 + k eta^2/6.
It prints the largest value of (left side - bound) for each bound, which must be <= 0.
"""
import numpy as np
from scipy.linalg import sqrtm

rng = np.random.default_rng(5)
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)


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


worst = {"e:contour-vertex identity (abs. error)": 0.0, "e:contour-vertex-bound": -np.inf,
         "e:contour-edge identity (abs. error)": 0.0, "e:contour-edge bound": -np.inf,
         "l:contour-real(c) bound": -np.inf, "display after p:rep-contour (abs. error)": 0.0,
         "e:rep-homotopy, per site": -np.inf}
TRIALS = 4000
for trial in range(TRIALS):
    E = rng.uniform(-1.95, 1.95)
    s = np.sqrt(1 - E**2 / 4)
    f0 = lambda Z: sqrtm(s**2 * I + Z @ Z)
    t = rng.uniform(0, 1)
    X, Xp, P, Pp = offd(rng.exponential(3)), offd(rng.exponential(3)), bd(), bd()
    F, Fp = f0(X), f0(Xp)
    zeta = (rng.standard_normal(n) + 1j * rng.standard_normal(n)) * rng.exponential(1)
    zz = np.outer(zeta, zeta.conj())

    # l:contour-real(a)
    c, etap, m = rng.uniform(0, 1), rng.uniform(0, 1), rng.uniform(0.1, 2)
    Rt = lambda P_, F_, X_, c_: P_ - 1j * Lam @ ((1 - t) * (F_ - c_ * I) + X_)
    R = Rt(P, F, X, c)
    val = (-m / 2 * np.trace(R @ R) - 1j * np.trace(R @ (zz + etap * I) @ Lam)).real
    Zm = (1 - t) * (F - c * I) + X
    pred = (-m / 2 * (hs(P) ** 2 + hs(X) ** 2 - (1 - t) ** 2 * hs(F - c * I) ** 2)
            - (zeta.conj() @ Zm @ zeta).real - (1 - t) * etap * (np.trace(F).real - 2 * k * c))
    worst["e:contour-vertex identity (abs. error)"] = max(worst["e:contour-vertex identity (abs. error)"], abs(val - pred) / (1 + abs(val)))
    etab = m * c + etap
    nz = np.linalg.norm(zeta)
    bound = (-m / 2 * hs(P) ** 2 - etab / 4 * np.trace(F).real + m * k * (s**2 + c**2) + 2 * k * etap * c
             + (s + c) * nz**2 + nz**4 / m + k * etab * s / 2 + k * etab**2 / (6 * m))
    worst["e:contour-vertex-bound"] = max(worst["e:contour-vertex-bound"], val - bound)

    # l:contour-real(b)
    w = rng.uniform(0, 3)
    D = Rt(P, F, X, c) - Rt(Pp, Fp, Xp, c)
    val = (-w / 2 * np.trace(D @ D)).real
    pred = -w / 2 * (hs(P - Pp) ** 2 + hs(X - Xp) ** 2 - (1 - t) ** 2 * hs(F - Fp) ** 2)
    worst["e:contour-edge identity (abs. error)"] = max(worst["e:contour-edge identity (abs. error)"], abs(val - pred) / (1 + abs(val)))
    worst["e:contour-edge bound"] = max(worst["e:contour-edge bound"], val + w / 2 * hs(P - Pp) ** 2)

    # l:contour-real(c), c = 0
    Xs, Ps = offd(rng.exponential(1)), bd()
    Fs = f0(Xs)
    Rs = Ps - 1j * Lam @ (Fs + Xs)
    D = Rt(P, F, X, 0) - Rs
    val = (-w / 2 * np.trace(D @ D)).real
    gs2 = (Fs + Xs) / s
    bound = (-w / 2 * hs(P - Ps) ** 2 - w * s / (2 * op(gs2)) * np.trace(F).real + 2 * w * k * s**2
             + w / 2 * hs(Fs) ** 2)
    worst["l:contour-real(c) bound"] = max(worst["l:contour-real(c) bound"], val - bound)

    # the display after p:rep-contour and e:rep-homotopy (per site), with m = 1, c = eta, eta' = 0
    eta = rng.uniform(0.001, 1)
    R = Rt(P, F, X, eta)
    val = (-0.5 * np.trace(R @ R) - 1j * (zeta.conj() @ Lam @ R @ zeta)).real
    pred = (-0.5 * hs(P) ** 2 - t * (2 - t) / 2 * hs(X) ** 2 - (1 - t) ** 2 * eta * np.trace(F).real
            + (1 - t) ** 2 * k * (s**2 + eta**2) - (zeta.conj() @ ((1 - t) * (F - eta * I) + X) @ zeta).real)
    worst["display after p:rep-contour (abs. error)"] = max(worst["display after p:rep-contour (abs. error)"], abs(val - pred) / (1 + abs(val)))
    C = k * (s**2 + eta**2) + (s + eta) * nz**2 + nz**4 + k * eta * s / 2 + k * eta**2 / 6
    bound = -0.5 * hs(P) ** 2 - eta / 4 * np.trace(F).real + C
    worst["e:rep-homotopy, per site"] = max(worst["e:rep-homotopy, per site"], val - bound)

print(f"{TRIALS} random points")
ok = True
for key, v in worst.items():
    print(f"  {key:<45s} {v: .2e}")
    ok &= (v < 1e-8) if "error" in key else (v <= 1e-9)
print("all identities and bounds hold" if ok else "FAILED")