"""Ordinary (non-super) check of V_F(0) = 0 for the tubular half-density.

V = V_+ (+) V_-, dim V_+ = p, dim V_- = q, n = p + q.  R_* = diag(a I_p, b I_q).
Tubular map (X_m, U) -> exp(X) (R_* + U) exp(-X), X_m off-diagonal, U block diagonal.
F(U)^2 := det of the Jacobian of this map at X = 0 (computed by finite differences
of the actual matrix exponential map, not from the differential formula).
Metric tr(XY) on gl(p) (+) gl(q); orthonormal real basis of Hermitian matrices;
for holomorphic f, Delta f = sum_alpha d^2/dt^2 f(U + t E_alpha).
Predictions (gamma = 1/(a-b), u = p, v = q):
  Delta log F(0) = -gamma^2 u v (u+v),  |grad log F(0)|^2 = +gamma^2 u v (u+v).
"""
import mpmath as mp

mp.mp.dps = 60


def expm(X):
    return mp.expm(X)


def blocks_basis(p, q):
    """Orthonormal basis (w.r.t. tr(XY)) of Hermitian block-diagonal matrices."""
    n = p + q
    out = []
    for off, m in ((0, p), (p, q)):
        for i in range(m):
            E = mp.zeros(n, n)
            E[off + i, off + i] = 1
            out.append(E)
        for i in range(m):
            for j in range(i + 1, m):
                E = mp.zeros(n, n)
                E[off + i, off + j] = 1 / mp.sqrt(2)
                E[off + j, off + i] = 1 / mp.sqrt(2)
                out.append(E)
                E = mp.zeros(n, n)
                E[off + i, off + j] = 1j / mp.sqrt(2)
                E[off + j, off + i] = -1j / mp.sqrt(2)
                out.append(E)
    return out


def coords_basis(p, q):
    """Complex coordinate basis: off-diagonal (m) entries and block-diagonal (l) entries."""
    n = p + q
    m_basis, l_basis = [], []
    for i in range(n):
        for j in range(n):
            E = mp.zeros(n, n)
            E[i, j] = 1
            if (i < p) == (j < p):
                l_basis.append(E)
            else:
                m_basis.append(E)
    return m_basis, l_basis


def jac_det(S, p, q, step=mp.mpf('1e-25')):
    m_basis, l_basis = coords_basis(p, q)
    n = p + q
    cols = []
    for E in m_basis:  # derivative in X direction at X = 0
        Rp = expm(step * E) * S * expm(-step * E)
        Rm = expm(-step * E) * S * expm(step * E)
        cols.append((Rp - Rm) / (2 * step))
    for E in l_basis:  # derivative in U direction
        cols.append(E)
    # flatten to n^2 x n^2 matrix
    J = mp.zeros(n * n, n * n)
    for c, M in enumerate(cols):
        for i in range(n):
            for j in range(n):
                J[i * n + j, c] = M[i, j]
    return mp.det(J)


def logF_jac(U, a, b, p, q):
    n = p + q
    S = mp.zeros(n, n)
    for i in range(n):
        S[i, i] = a if i < p else b
    S = S + U
    return jac_det(S, p, q)


def logF_def(U, a, b, p, q):
    """F = det of L_S : A -> S_+ A - A S_- on Hom(C^q, C^p)."""
    n = p + q
    S = mp.zeros(n, n)
    for i in range(n):
        S[i, i] = a if i < p else b
    S = S + U
    Sp = S[0:p, 0:p]
    Sm = S[p:n, p:n]
    L = mp.zeros(p * q, p * q)
    for i in range(p):
        for al in range(q):
            # image of E_{i al}
            A = mp.zeros(p, q)
            A[i, al] = 1
            B = Sp * A - A * Sm
            for k in range(p):
                for be in range(q):
                    L[k * q + be, i * q + al] = B[k, be]
    return mp.det(L) ** 2


def check(p, q, a, b, U0=None, use_jac=True, hstep=mp.mpf('1e-12')):
    n = p + q
    basis = blocks_basis(p, q)
    if U0 is None:
        U0 = mp.zeros(n, n)
    g = (lambda U: logF_jac(U, a, b, p, q)) if use_jac else (lambda U: logF_def(U, a, b, p, q))
    g0 = g(U0)
    f = lambda U: mp.log(g(U) / g0) / 2   # log F relative to base point: no branch jumps
    f0 = mp.mpf(0)
    lap = 0
    grad2 = 0
    for E in basis:
        fp = f(U0 + hstep * E)
        fm = f(U0 - hstep * E)
        lap += (fp - 2 * f0 + fm) / hstep**2
        grad2 += ((fp - fm) / (2 * hstep))**2
    return f0, lap, grad2


if __name__ == '__main__':
    s = mp.mpf('0.7')
    a, b = -1j * s, 1j * s
    gamma = 1 / (a - b)
    print('root: a = -is, b = is, s = 0.7, gamma^2 =', mp.nstr(gamma**2, 12))
    for (p, q) in [(1, 1), (2, 1), (1, 2), (2, 2), (3, 1), (2, 3)]:
        f0, lap, g2 = check(p, q, a, b)
        f0 = mp.log(logF_jac(mp.zeros(p + q, p + q), a, b, p, q)) / 2
        pred = gamma**2 * p * q * (p + q)
        # F^2 from Jacobian vs det(L_S)^2 (log-ratio up to i*pi*k for the sign)
        r = logF_jac(mp.zeros(p + q, p + q), a, b, p, q) / logF_def(mp.zeros(p + q, p + q), a, b, p, q)
        print(f'(p,q)=({p},{q}): det(Jacobian)/det(L_S)^2 at U=0 = {mp.nstr(r, 8)} (should be +-1)')
        print('   Delta log F(0)    =', mp.nstr(lap, 12), '  predicted', mp.nstr(-pred, 12))
        print('   |grad log F(0)|^2 =', mp.nstr(g2, 12), '  predicted', mp.nstr(pred, 12))
        print('   V_F(0) = sum      =', mp.nstr(lap + g2, 6))
    # generic complex root values and a random nonzero U (exact harmonicity in the ordinary case)
    mp.mp.dps = 60
    import random
    random.seed(1)
    p, q = 2, 2
    a, b = mp.mpc('0.3', '-1.1'), mp.mpc('-0.2', '0.9')
    U0 = mp.zeros(4, 4)
    for i in range(4):
        for j in range(4):
            if (i < 2) == (j < 2):
                U0[i, j] = mp.mpc(random.uniform(-.15, .15), random.uniform(-.15, .15))
    f0, lap, g2 = check(p, q, a, b, U0=U0, use_jac=True)
    print('generic a,b, random U != 0, (2,2): V_F(U) =', mp.nstr(lap + g2, 6),
          ' (Delta log F =', mp.nstr(lap, 8), ')')