"""Exact coefficients of the metallic expansion of the finite-wire scalar.

Supplement to Section 13 of the paper, "The metallic expansion" (s:metallic,
s:met-scalar).  The script implements the definitions of that section
literally and evaluates the coefficients I_0, ..., I_3 of

    I_X(z) = sum_m X^m I_m(z) + O(X^(M+1))

in exact arithmetic.  Notation (as in Section 13):

  * D = [1, oo) x [-1, 1] with coordinates (B, F), measure dB dF,
    d = B - F and s = B + F;
  * A = d_B (B^2 - 1) d_B + d_F (1 - F^2) d_F and C = d A d^{-1};
  * conjugated operators  A_t P = e^{z t d} A(e^{-z t d} P)  and
    C_t P = e^{z t d} C(e^{-z t d} P), acting on polynomials P;
  * the recursion of the singlet state (e:met-singlet-recursion)
        P_0 = 1,  P_k(t) = int_0^t C_r P_{k-1}(r) dr,  f_k(t) = e^{-z t d} P_k(t);
  * the flat terms
        g_0^(k)(b; c) = e^{-z b d} f_k(c),
        g_j^(k)(b; c) = int_0^b e^{-z (b - r) d} A g_{j-1}^(k)(r; c) dr;
    writing g_j^(k)(b; c) = e^{-z (b + c) d} G_j^(k)(b; c), this is
        G_0^(k) = P_k(c),  G_j^(k)(b; c) = int_0^b A_{r + c} G_{j-1}^(k)(r; c) dr;
  * the time-simplex integral T_m(d, s; z) of
        e^{z d} sum_{i + k + j = m} f_i(a) g_j^(k)(b; c),   c = 1 - a - b,
    over a, b >= 0, a + b <= 1 (the exponentials cancel, since a + b + c = 1);
  * I_m(z) = int_D e^{-z d} T_m dB dF, evaluated through the radial moments
        int_D h(d) s^j dB dF = int_0^oo h(d) m_j(d) dd,
        m_j(d) = ((d + 2)^{j + 1} - |d - 2|^{j + 1}) / (2 (j + 1)).

The value of I_m is computed twice: once with the radial moments m_j (split at
d = 2 for even j) and once by direct integration in (B, F).  The script checks
the (d, s) forms of A_t and C_t, which it derives from their definitions in
Section 13 (the paper does not display them), the polynomials T_0, T_1, T_2,
and

    I_0 = (1 - e^{-2z}) / (2 z^2),  I_1 = 0,  I_2 = (1 - e^{-2z}) / 45,
    I_3 = -8 z / 945,

and the spectral-trace forms 1/45 = 2 Tr G_N^2, 8/945 = 8 Tr G_N^3, where
G_N has eigenvalues (pi n)^{-2}, n >= 1.  The integrals are evaluated for real
z > 0; both sides of each identity are holomorphic in Re z > 0.

Requires sympy (tested with sympy 1.14).  Run with:  python -I metallic_coefficients.py
Expected output ends with "ALL CHECKS PASSED" (about 15 seconds).
"""

import sympy as sp

B, F, z = sp.symbols("B F z", positive=True)
a, b, c, r, t, dd = sp.symbols("a b c r t dd", nonnegative=True)
d = B - F
M_MAX = 3


def conj_A(R, tau):
    """A_tau R = e^{z tau d} A(e^{-z tau d} R) for a rational function R of B, F."""
    DB = lambda Q: sp.diff(Q, B) - z * tau * Q   # e^{z tau d} d_B e^{-z tau d}
    DF = lambda Q: sp.diff(Q, F) + z * tau * Q   # e^{z tau d} d_F e^{-z tau d}
    return sp.expand(DB((B**2 - 1) * DB(R)) + DF((1 - F**2) * DF(R)))


def conj_C(P, tau):
    """C_tau P = e^{z tau d} d A(d^{-1} e^{-z tau d} P), a rational function of B, F."""
    return sp.cancel(d * conj_A(P / d, tau))


def check_explicit_forms():
    """Check the (d, s) forms of A_t and C_t against their definitions in Section 13."""
    u, v = sp.symbols("u v")            # u = d, v = s
    P = u**3 * v**2 + 2 * u * v**4 + u**2 + 5 * v
    sub = {u: B - F, v: B + F}
    Dt = lambda Q: sp.diff(Q, u) - z * t * Q
    A_ds = (u * v * (Dt(Dt(P)) + sp.diff(P, v, 2))
            + (u**2 + v**2 - 4) * Dt(sp.diff(P, v)) + 2 * v * Dt(P) + 2 * u * sp.diff(P, v))
    C_ds = A_ds - 2 * v * Dt(P) - (u**2 + v**2 - 4) / u * sp.diff(P, v)
    PB = P.subs(sub)
    assert sp.expand(A_ds.subs(sub) - conj_A(PB, t)) == 0
    assert sp.simplify(C_ds.subs(sub) - conj_C(PB, t)) == 0
    assert sp.expand(conj_C(sp.Integer(1), t) - d * (B + F) * z**2 * t**2) == 0
    assert sp.expand(conj_C(d * PB, t) - d * conj_A(PB, t)) == 0
    print("explicit forms of A_t and C_t, C_t 1 = d s z^2 t^2 and C_t(dP) = d A_t P: checked")


def endpoint_polynomials():
    P = [sp.Integer(1)]
    for k in range(1, M_MAX + 1):
        integrand = conj_C(P[k - 1].subs(t, r), r)
        assert integrand.is_polynomial(B, F), "singular coefficient in the recursion"
        integrand = sp.expand(integrand)
        Pk = sp.expand(sp.integrate(integrand, (r, 0, t)))
        assert sp.cancel(Pk / d).is_polynomial(B, F), "P_k not divisible by d"
        P.append(Pk)
    return P


def flat_polynomials(P):
    """G[k][j] = G_j^(k)(b; c) for k + j <= M_MAX."""
    G = {}
    for k in range(M_MAX + 1):
        G[k] = [sp.expand(P[k].subs(t, c))]
        for j in range(1, M_MAX - k + 1):
            prev = G[k][j - 1].subs(b, r)
            G[k].append(sp.expand(sp.integrate(conj_A(prev, r + c), (r, 0, b))))
    return G


def simplex_polynomial(m, P, G):
    """T_m = int over a, b >= 0, a + b <= 1 of sum P_i(a) G_j^(k)(b; 1 - a - b)."""
    total = sp.Integer(0)
    for i in range(m + 1):
        for k in range(m - i + 1):
            j = m - i - k
            total += P[i].subs(t, a) * G[k][j].subs(c, 1 - a - b)
    total = sp.expand(total)
    inner = sp.integrate(total, (a, 0, 1 - b))
    return sp.expand(sp.integrate(sp.expand(inner), (b, 0, 1)))


def radial_moment(j):
    """m_j(d) on (0, 2) and on (2, oo)."""
    low = sp.expand(((dd + 2)**(j + 1) - (2 - dd)**(j + 1)) / (2 * (j + 1)))
    high = sp.expand(((dd + 2)**(j + 1) - (dd - 2)**(j + 1)) / (2 * (j + 1)))
    return low, high


def integrate_radial(T):
    """int_D e^{-z d} T dB dF through the radial moments m_j(d)."""
    u, v = sp.symbols("u v")
    T_ds = sp.expand(T.subs({B: (v + u) / 2, F: (v - u) / 2}))
    poly = sp.Poly(T_ds, v)
    result = sp.Integer(0)
    for (j,), coeff in poly.terms():
        low, high = radial_moment(j)
        h = coeff.subs(u, dd)
        result += sp.integrate(sp.exp(-z * dd) * h * low, (dd, 0, 2))
        result += sp.integrate(sp.exp(-z * dd) * h * high, (dd, 2, sp.oo))
    return sp.simplify(result)


def integrate_direct(T):
    """int_1^oo int_{-1}^1 e^{-z (B - F)} T dF dB, monomial by monomial."""
    poly = sp.Poly(sp.expand(T), B, F)
    result = sp.Integer(0)
    for (p, q), coeff in poly.terms():
        IB = sp.integrate(B**p * sp.exp(-z * B), (B, 1, sp.oo))
        IF = sp.integrate(F**q * sp.exp(z * F), (F, -1, 1))
        result += coeff * IB * IF
    return sp.simplify(result)


def main():
    check_explicit_forms()
    P = endpoint_polynomials()
    print("P_1(t) =", sp.factor(P[1]))
    G = flat_polynomials(P)
    expected = {
        0: (1 - sp.exp(-2 * z)) / (2 * z**2),
        1: sp.Integer(0),
        2: (1 - sp.exp(-2 * z)) / 45,
        3: -8 * z / 945,
    }
    # T_0, T_1, T_2 in the variables d = B - F, s = B + F (Section 13 displays the
    # polynomial of the second-order term in (e:met-second-polynomial))
    s_ = B + F
    displayed_T = {
        0: sp.Rational(1, 2),
        1: s_ * (d * z**2 / 12 - z / 6),
        2: (-d * z / 15 + (d**2 - 2) * z**2 / 15 + (4 * d - d**3) * z**3 / 90
            + s_**2 * (z**2 / 10 - d * z**3 / 15 + d**2 * z**4 / 120)),
    }
    assert sp.expand(conj_A(sp.Integer(1), t) - (z**2 * t**2 * d * s_ - 2 * z * t * s_)) == 0
    n = sp.symbols("n", integer=True, positive=True)
    trace = lambda p: sp.summation((sp.pi * n)**(-2 * p), (n, 1, sp.oo))
    ok = True
    for m in range(M_MAX + 1):
        T = simplex_polynomial(m, P, G)
        if m in displayed_T:
            agrees = sp.expand(T - displayed_T[m]) == 0
            ok = ok and agrees
            print(f"T_{m} agrees with the displayed polynomial: {agrees}")
        I_rad = integrate_radial(T)
        I_dir = integrate_direct(T)
        same = sp.simplify(I_rad - I_dir) == 0
        match = sp.simplify(I_rad - expected[m]) == 0
        ok = ok and same and match
        print(f"I_{m}(z) = {sp.expand(I_rad)}"
              f"   [radial = direct: {same}; as stated: {match}]")
    tr2, tr3 = trace(2), trace(3)
    print("Tr G_N^2 =", tr2, "; 2 Tr G_N^2 = 1/45:", sp.simplify(2 * tr2 - sp.Rational(1, 45)) == 0)
    print("Tr G_N^3 =", tr3, "; 8 Tr G_N^3 = 8/945:", sp.simplify(8 * tr3 - sp.Rational(8, 945)) == 0)
    print("ALL CHECKS PASSED" if ok else "SOME CHECK FAILED")


if __name__ == "__main__":
    main()