"""Numerical check of l:rep-dictionary (Section 4), in the convention G(z) = (H - z)^{-1}.

Run: python3 rep_dictionary_check.py   (numpy only; about a minute)

(1) Matrix side.  For a fixed random Hermitian H of size N = 5 and k = 2, the Taylor
    coefficients of Sdet(E + i eta0 Gamma_0 + Y - H)^{-1} (e:rep-TY before the average) in
    the local sources Y are computed by Cauchy integrals and compared with the integrands of
    the left sides of e:rep-dict-perm, e:rep-dict-det (m = 3, so that the sign (-1)^m is
    visible), e:rep-dict-imperm, e:rep-dict-ldos (two flavors at distinct sites) and
    e:rep-dict-moment (q = 2).  These identities hold for each H, before the expectation.
(2) Chain side.  For a supermatrix E - S with random complex BB and FF blocks and vanishing
    odd blocks, the Taylor coefficients of Sdet(Id + (E - S)^{-1} Y)^{-1} (e:rep-insertion)
    are compared with the single-site insertions of e:rep-dict-perm, e:rep-dict-det,
    e:rep-dict-imperm (gamma_1 of e:rep-gamma) and e:rep-dict-moment.  Two flavors at one
    site give gamma_1 gamma_2 + (e_1^-)^T y e_2^+ (e_2^-)^T y e_1^+, not gamma_1 gamma_2,
    which is why e:rep-dict-ldos requires distinct sites (remark after l:rep-dictionary).
(3) Saddle values (e:rep-saddle-inverse and the paragraph after it).  At S = kappa_E - i s_E Lambda,
    -(E - S)^{-1} has retarded entries m_sc(E + i0) and advanced entries m_sc(E - i0), and
    gamma_a(S) = s_E.
"""
import itertools
import math

import numpy as np

rng = np.random.default_rng(7)
TOL = 1e-9
failures = []
max_err = [0.0]


def report(label, got, want):
    err = abs(got - want)
    ok = err < TOL * (1 + abs(want))
    max_err[0] = max(max_err[0], err / (1 + abs(want)))
    print(f"{'ok  ' if ok else 'FAIL'} {label:<58s} {complex(got):.10f}   {complex(want):.10f}")
    if not ok:
        failures.append(label)


def coeff(fun, nvar, mon, r=0.05, M=24):
    """Taylor coefficient of t^mon of fun(t), holomorphic near 0, by Cauchy integrals."""
    th = 2 * np.pi * np.arange(M) / M
    acc = 0
    for idx in itertools.product(range(M), repeat=nvar):
        t = r * np.exp(1j * th[list(idx)])
        acc += fun(t) * np.prod(np.exp(-1j * np.array(mon) * th[list(idx)]))
    return acc / (M**nvar * r ** sum(mon))


def perm(A):
    n = A.shape[0]
    return sum(np.prod([A[i, p[i]] for i in range(n)]) for p in itertools.permutations(range(n)))


N, k = 5, 2
A = rng.standard_normal((N, N)) + 1j * rng.standard_normal((N, N))
H = (A + A.conj().T) / 2
E, eta = 0.4, 0.3
z0 = E + 1j * eta
G = lambda w: np.linalg.inv(H - w * np.eye(N))  # G(z) = (H - z)^{-1}
zs = np.array([z0] * k + [np.conj(z0)] * k)  # diagonal of E + i eta0 Gamma_0 (BB and FF)


def sdet_inv(YB, YF):
    """Sdet(D + Y)^{-1} = det(FF block)/det(BB block), D = E + i eta0 Gamma_0 - H, Y = (Y_i) on sites."""
    BB = np.kron(np.eye(N), np.diag(zs)) - np.kron(H, np.eye(2 * k))
    FF = BB.copy()
    for i, Yi in YB.items():
        BB[2 * k * i:2 * k * (i + 1), 2 * k * i:2 * k * (i + 1)] += Yi
    for i, Yi in YF.items():
        FF[2 * k * i:2 * k * (i + 1), 2 * k * i:2 * k * (i + 1)] += Yi
    return np.linalg.det(FF) / np.linalg.det(BB)


def Eab(a, b):
    M = np.zeros((2 * k, 2 * k), complex)
    M[a, b] = 1
    return M


def e_plus(a):
    v = np.zeros(2 * k, complex)
    v[a] = 1
    v[k + a] = 1
    return v


def e_minus(a):
    v = np.zeros(2 * k, complex)
    v[a] = 1
    v[k + a] = -1
    return v / 2j


V = lambda a: np.outer(e_plus(a), e_minus(a))  # e_a^+ (e_a^-)^T

print("(1) matrix side, fixed H, k = 2, N = 5")
sites = [0, 2, 4]
Gz = G(z0)
ImG = (Gz - Gz.conj().T) / (2j)
f = lambda t: sdet_inv({s: t[a] * Eab(0, 0) for a, s in enumerate(sites)}, {})
report("e:rep-dict-perm, m = 3: per G", coeff(f, 3, [1, 1, 1]), perm(Gz[np.ix_(sites, sites)]))
f = lambda t: sdet_inv({}, {s: t[a] * Eab(0, 0) for a, s in enumerate(sites)})
report("e:rep-dict-det, m = 3: coefficient vs (-1)^3 det G", coeff(f, 3, [1, 1, 1]),
       -np.linalg.det(Gz[np.ix_(sites, sites)]))
f = lambda t: sdet_inv({s: t[a] * V(0) for a, s in enumerate(sites)}, {})
report("e:rep-dict-imperm, m = 3: per Im G", coeff(f, 3, [1, 1, 1]), perm(ImG[np.ix_(sites, sites)]))
f = lambda t: sdet_inv({0: t[0] * V(0), 3: t[1] * V(1)}, {})
report("e:rep-dict-ldos, flavors 1, 2 at sites 0, 3", coeff(f, 2, [1, 1]), ImG[0, 0] * ImG[3, 3])
f = lambda t: sdet_inv({1: t[0] * V(0) + t[1] * V(1)}, {})
report("two flavors at one site, matrix side: (Im G_11)^2", coeff(f, 2, [1, 1]), ImG[1, 1] ** 2)
f = lambda t: sdet_inv({1: t[0] * Eab(0, 0) + t[1] * Eab(k, k)}, {})
report("e:rep-dict-moment, q = 2: |G_11|^4", coeff(f, 2, [2, 2]), abs(Gz[1, 1]) ** 4)

print("(2) chain side, random numeric blocks y = [(E-S)^{-1}]^BB, yF = [(E-S)^{-1}]^FF")
y = rng.standard_normal((2 * k, 2 * k)) + 1j * rng.standard_normal((2 * k, 2 * k))
yF = rng.standard_normal((2 * k, 2 * k)) + 1j * rng.standard_normal((2 * k, 2 * k))
Z = np.zeros((2 * k, 2 * k))
ins = lambda YB, YF: np.linalg.det(np.eye(2 * k) + yF @ YF) / np.linalg.det(np.eye(2 * k) + y @ YB)
gam = lambda a: -(e_minus(a) @ y @ e_plus(a))  # e:rep-gamma
f = lambda t: ins(t[0] * Eab(0, 0), Z)
report("e:rep-dict-perm insertion: -y_11", coeff(f, 1, [1]), -y[0, 0])
f = lambda t: ins(Z, t[0] * Eab(0, 0))
report("e:rep-dict-det insertion: coefficient yF_11, so G <-> -yF", coeff(f, 1, [1]), yF[0, 0])
f = lambda t: ins(t[0] * V(0), Z)
report("e:rep-dict-imperm insertion: gamma_1", coeff(f, 1, [1]), gam(0))
q = 2
f = lambda t: ins(t[0] * Eab(0, 0) + t[1] * Eab(k, k), Z)
pred = sum(math.comb(q, r) ** 2 * (y[0, 0] * y[k, k]) ** (q - r) * (y[0, k] * y[k, 0]) ** r for r in range(q + 1))
report("e:rep-dict-moment insertion, q = 2", coeff(f, 2, [q, q]), pred)
f = lambda t: ins(t[0] * V(0) + t[1] * V(1), Z)
coincident = gam(0) * gam(1) + (e_minus(0) @ y @ e_plus(1)) * (e_minus(1) @ y @ e_plus(0))
report("two flavors at one site: gamma_1 gamma_2 + cross term", coeff(f, 2, [1, 1]), coincident)
print(f"     (gamma_1 gamma_2 alone = {complex(gam(0) * gam(1)):.10f}: e:rep-dict-ldos needs distinct sites)")

print("(3) saddle values at S = kappa_E - i s_E Lambda")
for Ev in [0.0, 0.7, -1.3]:
    kap, s = Ev / 2, np.sqrt(1 - Ev**2 / 4)
    Lam = np.diag([1.0] * k + [-1.0] * k)
    S = kap * np.eye(2 * k) - 1j * s * Lam
    Rinv = np.linalg.inv(Ev * np.eye(2 * k) - S)
    report(f"E = {Ev}: (E - S)^(-1) = S (e:rep-saddle-inverse)", np.abs(Rinv - S).max(), 0)
    msc = lambda w: (-w + np.sqrt(w - 2 + 0j) * np.sqrt(w + 2 + 0j)) / 2  # Stieltjes transform, Im > 0 for Im w > 0
    report(f"E = {Ev}: retarded entry of -(E - S)^(-1) vs m_sc(E + i0)", -Rinv[0, 0], msc(Ev + 1e-14j))
    report(f"E = {Ev}: advanced entry of -(E - S)^(-1) vs m_sc(E - i0)", -Rinv[k, k], msc(Ev - 1e-14j))
    y0 = Rinv  # BB block of (E - S)^{-1} on the saddle
    report(f"E = {Ev}: gamma_1(S) = s_E", -(e_minus(0) @ y0 @ e_plus(0)), s)

print(f"largest relative error {max_err[0]:.1e}")
print("all checks passed" if not failures else f"FAILED: {failures}")