"""ratjets.py -- exact (eps, x, u)-jets for the realizability chain of a metric on a 3-manifold.

A quantity is carried as a 3-index array Q[i][j][k] = coefficient of eps^i x^j u^k, with i < NE, j <= JX, k <= JU,
where x = r - r0 and u = theta - theta0 are local coordinates about a point and eps is a deformation parameter.
Coefficients are exact (rationals, or exact algebraic numbers where a square root of a non-square appears). No
numerical identification (sympy's nsimplify) is used anywhere: it can replace an algebraic number by a different one
that agrees to about 15 digits.
Derivatives are index shifts. Expanding about a point on the circle with rational (cos theta0, sin theta0) keeps every
trigonometric coefficient rational.

The chain: metric g -> Christoffel symbols -> Ricci tensor and scalar -> M = -G (as an endomorphism)
-> K~ = sqrt(det M) M^{-1}, lowered with g -> Codazzi tensor C_ajk = nabla_a K~_jk - nabla_j K~_ak.
Codazzi needs three derivatives of g, so with JX = JU = 4 the values at the expansion point are exact for every
eps order carried.

Metrics here are invariant under phi -> -phi and independent of phi, so the third coordinate derivative is zero.
"""
import sympy as sp
from sympy import Rational as R

JX, JU = 4, 4
NE = 3
C0, S0 = R(3, 5), R(4, 5)


def set_point(c, s):
    """Set (cos theta0, sin theta0); must lie on the unit circle."""
    global C0, S0
    c, s = sp.sympify(c), sp.sympify(s)
    assert sp.simplify(c**2 + s**2 - 1) == 0, "set_point: not on the unit circle"
    C0, S0 = c, s


def zero():
    return [[[R(0)] * (JU + 1) for _ in range(JX + 1)] for _ in range(NE)]


def const(c):
    a = zero()
    a[0][0][0] = sp.sympify(c)
    return a


def add(a, b):
    return [[[a[i][j][k] + b[i][j][k] for k in range(JU + 1)] for j in range(JX + 1)] for i in range(NE)]


def sub(a, b):
    return [[[a[i][j][k] - b[i][j][k] for k in range(JU + 1)] for j in range(JX + 1)] for i in range(NE)]


def smul(c, a):
    c = sp.sympify(c)
    return [[[c * a[i][j][k] for k in range(JU + 1)] for j in range(JX + 1)] for i in range(NE)]


def mul(a, b):
    out = zero()
    for i1 in range(NE):
        for j1 in range(JX + 1):
            for k1 in range(JU + 1):
                c1 = a[i1][j1][k1]
                if c1 == 0:
                    continue
                for i2 in range(NE - i1):
                    for j2 in range(JX + 1 - j1):
                        for k2 in range(JU + 1 - k1):
                            c2 = b[i2][j2][k2]
                            if c2 == 0:
                                continue
                            out[i1 + i2][j1 + j2][k1 + k2] += c1 * c2
    return out


def _series(a, coef):
    """sum_n coef(n) w^n with a = a0 (1 + w); returns (a0, sum)."""
    a0 = a[0][0][0]
    w = sub(smul(1 / a0, a), const(1))
    acc, term = const(1), const(1)
    for n in range(1, NE + JX + JU + 1):
        term = mul(term, w)
        if all(term[i][j][k] == 0 for i in range(NE) for j in range(JX + 1) for k in range(JU + 1)):
            break
        acc = add(acc, smul(coef(n), term))
    return a0, acc


def inv(a):
    assert a[0][0][0] != 0, "inv: constant term vanishes"
    a0, acc = _series(a, lambda n: (-1)**n)
    return smul(1 / a0, acc)


def sqrtj(a):
    assert a[0][0][0] > 0, "sqrtj: constant term not positive"
    coefs = {0: R(1)}

    def coef(n):
        if n not in coefs:
            coefs[n] = coef(n - 1) * (R(1, 2) - (n - 1)) / n
        return coefs[n]
    a0, acc = _series(a, coef)
    return smul(sp.sqrt(a0), acc)


def dx(a):
    out = zero()
    for i in range(NE):
        for j in range(JX):
            for k in range(JU + 1):
                out[i][j][k] = (j + 1) * a[i][j + 1][k]
    return out


def du(a):
    out = zero()
    for i in range(NE):
        for j in range(JX + 1):
            for k in range(JU):
                out[i][j][k] = (k + 1) * a[i][j][k + 1]
    return out


def val(a, i):
    """value of the eps^i part at the expansion point (exact; no numerical identification)"""
    return sp.radsimp(sp.sympify(a[i][0][0]))


def at_eps(a, n):
    """move the eps^0 content of a jet into the eps^n slot"""
    out = zero()
    for j in range(JX + 1):
        for k in range(JU + 1):
            out[n][j][k] = a[0][j][k]
    return out


def cos_u():
    a = zero()
    cu = [R(1), R(0), R(-1, 2), R(0), R(1, 24)][:JU + 1]
    su = [R(0), R(1), R(0), R(-1, 6), R(0)][:JU + 1]
    for k in range(JU + 1):
        a[0][0][k] = C0 * cu[k] - S0 * su[k]
    return a


def sin_u():
    a = zero()
    cu = [R(1), R(0), R(-1, 2), R(0), R(1, 24)][:JU + 1]
    su = [R(0), R(1), R(0), R(-1, 6), R(0)][:JU + 1]
    for k in range(JU + 1):
        a[0][0][k] = S0 * cu[k] + C0 * su[k]
    return a


def poly_x(coeffs, eps_order=0):
    """a function of r alone, from its Taylor coefficients about r0 (already divided by k!)"""
    a = zero()
    for j, c in enumerate(coeffs):
        if j <= JX:
            a[eps_order][j][0] = sp.sympify(c)
    return a


def taylor(expr, r, r0):
    """exact Taylor coefficients of expr(r) about r0, up to x^JX"""
    return [sp.radsimp(sp.simplify(sp.diff(expr, r, k).subs(r, r0) / sp.factorial(k))) for k in range(JX + 1)]


def det3(X):
    t = zero()
    for (p, q, s, sg) in [(0, 1, 2, 1), (1, 2, 0, 1), (2, 0, 1, 1), (0, 2, 1, -1), (2, 1, 0, -1), (1, 0, 2, -1)]:
        t = add(t, smul(sg, mul(X[0][p], mul(X[1][q], X[2][s]))))
    return t


def adj3(X):
    A = [[zero() for _ in range(3)] for _ in range(3)]
    for i in range(3):
        for j in range(3):
            r1, r2 = [a for a in range(3) if a != i]
            c1, c2 = [a for a in range(3) if a != j]
            cof = sub(mul(X[r1][c1], X[r2][c2]), mul(X[r1][c2], X[r2][c1]))
            A[j][i] = smul((-1)**(i + j), cof)
    return A


def chain(g):
    """Run the realizability chain on a symmetric 3x3 array of jets g (coordinates r, theta, phi).
    Returns a dict with M (endomorphism), K (K~ lowered) and a function cod(a, j, k)."""
    D = [dx, du, lambda a: zero()]
    dg = det3(g)
    dgi = inv(dg)
    ag = adj3(g)
    ginv = [[mul(ag[i][j], dgi) for j in range(3)] for i in range(3)]
    for i in range(3):
        for j in range(3):
            acc = zero()
            for m in range(3):
                acc = add(acc, mul(g[i][m], ginv[m][j]))
            assert acc[0][0][0] == (1 if i == j else 0) and all(acc[e][0][0] == 0 for e in range(1, NE)), "inverse"
    d = [[[D[k](g[i][j]) for k in range(3)] for j in range(3)] for i in range(3)]
    Gam = [[[None] * 3 for _ in range(3)] for _ in range(3)]
    for l in range(3):
        for i in range(3):
            for j in range(3):
                acc = zero()
                for m in range(3):
                    acc = add(acc, mul(ginv[l][m], sub(add(d[m][i][j], d[m][j][i]), d[i][j][m])))
                Gam[l][i][j] = smul(R(1, 2), acc)
    Ric = [[zero() for _ in range(3)] for _ in range(3)]
    for i in range(3):
        for j in range(3):
            acc = zero()
            for l in range(3):
                acc = add(acc, D[l](Gam[l][i][j]))
                acc = sub(acc, D[j](Gam[l][i][l]))
                for m in range(3):
                    acc = add(acc, mul(Gam[l][l][m], Gam[m][i][j]))
                    acc = sub(acc, mul(Gam[l][j][m], Gam[m][i][l]))
            Ric[i][j] = acc
    Rs = zero()
    for i in range(3):
        for j in range(3):
            Rs = add(Rs, mul(ginv[i][j], Ric[i][j]))
    Mlow = [[smul(-1, sub(Ric[i][j], smul(R(1, 2), mul(Rs, g[i][j])))) for j in range(3)] for i in range(3)]
    M = [[zero() for _ in range(3)] for _ in range(3)]
    for i in range(3):
        for j in range(3):
            acc = zero()
            for m in range(3):
                acc = add(acc, mul(ginv[i][m], Mlow[m][j]))
            M[i][j] = acc
    dM = det3(M)
    sq = sqrtj(dM)
    aM = adj3(M)
    idM = inv(dM)
    Kmix = [[mul(sq, mul(aM[i][j], idM)) for j in range(3)] for i in range(3)]
    K = [[zero() for _ in range(3)] for _ in range(3)]
    for i in range(3):
        for j in range(3):
            acc = zero()
            for m in range(3):
                acc = add(acc, mul(g[i][m], Kmix[m][j]))
            K[i][j] = acc

    def cod(a, j, k):
        t = sub(D[a](K[j][k]), D[j](K[a][k]))
        for m in range(3):
            t = sub(t, mul(Gam[m][a][j], K[m][k]))
            t = sub(t, mul(Gam[m][a][k], K[j][m]))
            t = add(t, mul(Gam[m][j][a], K[m][k]))
            t = add(t, mul(Gam[m][j][k], K[a][m]))
        return t
    return {"M": M, "Kmix": Kmix, "K": K, "cod": cod, "Gam": Gam}


CODIDX = [(a, b, c) for (a, b) in ((0, 1), (0, 2), (1, 2)) for c in range(3)]


def _selftest():
    got = inv(poly_x([1, 1]))
    assert [got[0][j][0] for j in range(5)] == [1, -1, 1, -1, 1]
    got = sqrtj(poly_x([1, 1]))
    assert [got[0][j][0] for j in range(4)] == [R(1), R(1, 2), R(-1, 8), R(1, 16)]
    assert du(cos_u())[0][0][0] == -sin_u()[0][0][0]
    p = add(mul(sin_u(), sin_u()), mul(cos_u(), cos_u()))
    assert p[0][0][0] == 1 and all(p[0][0][k] == 0 for k in range(1, JU + 1))


_selftest()