#!/usr/bin/env python3
"""Check the covariance, chiral tensors and fourth-order determinant identity."""
from __future__ import annotations

from itertools import combinations
import sympy as sp


def equal(actual: sp.Expr, expected: sp.Expr = sp.S.Zero) -> None:
    difference = sp.cancel(sp.expand(actual - expected))
    if difference != 0:
        raise AssertionError(f"Nonzero difference: {difference}")


def matrix_equal(actual: sp.MatrixBase, expected: sp.MatrixBase) -> None:
    if actual.shape != expected.shape:
        raise AssertionError("Matrix dimensions differ.")
    for value in actual - expected:
        equal(value)


def curl(momentum: sp.MatrixBase) -> sp.Matrix:
    result = sp.zeros(6, 4)
    for row, (mu, nu) in enumerate(combinations(range(4), 2)):
        result[row, nu] = momentum[mu]
        result[row, mu] = -momentum[nu]
    return result


def check_covariances() -> None:
    eta, k, om, kap = sp.symbols("eta k omega kappa", nonzero=True, real=True)
    metric = sp.Matrix([[0, -1, 0, 0], [-1, eta**2, 0, 0],
                        [0, 0, 1, 0], [0, 0, 0, 1]])
    inverse = metric.inv()
    p = sp.Matrix([-om, -k, kap, 0])
    p2 = (p.T * inverse * p)[0]
    equal(p2, kap**2 - 2*k*om - eta**2*om**2)
    sharp = inverse * p
    hessian = sharp * sharp.T - p2 * inverse
    covariant = -metric/p2
    embeddings = [
        sp.Matrix([[1, 0, 0], [0, 0, 0], [0, 1, 0], [0, 0, 1]]),
        sp.Matrix([[1, 0, 0], [0, 1, 0], [0, 0, 0], [0, 0, 1]])]
    normals = [sp.Matrix([0, 1, 0, 0]), sp.Matrix([0, 0, 1, 0])]
    expected = [
        sp.Matrix([[k**2+eta**2*kap**2, kap*(k+eta**2*om), 0],
                   [kap*(k+eta**2*om), 2*k*om+eta**2*om**2, 0],
                   [0, 0, -p2]]),
        sp.Matrix([[eta**2*kap**2, kap**2, 0], [kap**2, om**2, 0],
                   [0, 0, eta**2*om**2-kap**2]])]
    for i, (embedding, normal) in enumerate(zip(embeddings, normals)):
        substitution = {} if i == 0 else {k: 0}
        p_here = p.subs(substitution)
        h_here = hessian.subs(substitution)
        cov_here = covariant.subs(substitution)
        reduced = embedding.T * h_here * embedding
        matrix_equal(reduced, expected[i])
        projection = sp.eye(4) - p_here*normal.T/(normal.T*p_here)[0]
        embedded = embedding * reduced.inv() * embedding.T
        matrix_equal(embedded, projection*cov_here*projection.T)
        c = curl(p_here)
        matrix_equal(c*embedded*c.T, c*cov_here*c.T)
    # Strict-null transverse/flux completeness, including the reversed flux basis.
    pp, pm, p1, p2t = sp.symbols("pplus pminus p1 p2", nonzero=True)
    mom = sp.Matrix([pp, pm, p1, p2t])
    g0 = sp.Matrix([[0, -1, 0, 0], [-1, 0, 0, 0],
                    [0, 0, 1, 0], [0, 0, 0, 1]])
    norm = (mom.T*g0*mom)[0]
    n = sp.Matrix([0, 1, 0, 0])
    basis = sp.Matrix.hstack(sp.Matrix([0, p1/pp, 1, 0]),
                            sp.Matrix([0, p2t/pp, 0, 1]), n/pp)
    reversed_basis = sp.Matrix.hstack(basis[:, 0], basis[:, 1], -basis[:, 2])
    actual = basis*sp.diag(1/norm, 1/norm, 1)*reversed_basis.T
    expected_comp = g0/norm - (mom*n.T+n*mom.T)/(norm*(mom.T*g0*n)[0])
    matrix_equal(actual, expected_comp)


def chiral_matrix(r: tuple[sp.Expr, ...]) -> sp.Matrix:
    r0, r1, r2, r3 = r
    spatial = sp.Matrix([r1, r2, r3])
    cross = sp.Matrix([[0, -r3, r2], [r3, 0, -r1], [-r2, r1, 0]])
    return (r0**2-r1**2-r2**2-r3**2)*sp.eye(3)+2*spatial*spatial.T+2*r0*cross


def check_chiral_tensors() -> None:
    r = sp.symbols("r0:4")
    t = sp.symbols("t0:4")
    nr, nt = chiral_matrix(r), chiral_matrix(t)
    rr, tt, rt = sum(x*x for x in r), sum(x*x for x in t), sum(x*y for x, y in zip(r, t))
    matrix_equal(nr*nr.T, rr**2*sp.eye(3))
    m = nr*nt.T
    equal(sp.trace(m), 4*rt**2-rr*tt)
    equal(sp.trace(m*m), 16*rt**4-16*rr*tt*rt**2+3*rr**2*tt**2)
    # Hodge projection of x wedge contraction.
    x = sp.Matrix(sp.symbols("x0:4"))
    c = curl(x)
    exterior = c*c.T
    pairs = list(combinations(range(4), 2))
    star = sp.zeros(6)
    for i, pair in enumerate(pairs):
        for j, other in enumerate(pairs):
            indices = (*pair, *other)
            if len(set(indices)) == 4:
                inv = sum(indices[a] > indices[b] for a in range(4) for b in range(a+1, 4))
                star[i, j] = (-1)**inv
    matrix_equal(star*star, sp.eye(6))
    for sign in (-1, 1):
        projector = (sp.eye(6)+sign*star)/2
        matrix_equal(projector*exterior*projector, sum(t*t for t in x)*projector/2)


def check_source_determinant() -> None:
    t = sp.symbols("t")
    t0 = sp.diag(2, 3)
    k0 = sp.diag(5, 7)
    matrices = [
        (sp.Matrix([[1, 2], [2, -1]]), sp.Matrix([[1, -2], [3, 1]]), sp.Matrix([[2, 1], [1, 3]])),
        (sp.Matrix([[2, -1], [-1, 4]]), sp.Matrix([[2, 1], [-1, 2]]), sp.Matrix([[1, -3], [-3, 2]]))]
    for L, C, V in matrices:
        h = (t0+t*L).row_join(t*C.T).col_join((t*C).row_join(k0+t*V))
        ratio = sp.expand(h.det()/(t0.det()*k0.det()))
        coefficient = sp.series(-sp.log(ratio)/2, t, 0, 5).removeO().expand().coeff(t, 4)
        G = k0.inv()
        U = [L, -C.T*G*C, C.T*G*V*G*C, -C.T*G*V*G*V*G*C]
        A1, A2, A3, A4 = [t0.inv()*u for u in U]
        trace_value = sp.trace(A1**4/8-A1**2*A2/2+A2**2/4+A1*A3/2-A4/2+(G*V)**4/8)
        equal(coefficient, trace_value)


def check_propagator() -> None:
    tau, eta, k, kap = sp.symbols("tau eta k kappa", positive=True)
    root = sp.sqrt(k*k+eta*eta*kap*kap)
    plus = sp.exp((k-root)*tau/eta**2)/(2*root)
    minus = sp.exp((k+root)*tau/eta**2)/(2*root)
    for part in (plus, minus):
        equal(-eta**2*sp.diff(part, tau, 2)+2*k*sp.diff(part, tau)+kap**2*part)
    equal(-eta**2*(sp.diff(plus, tau)-sp.diff(minus, tau)).subs(tau, 0), sp.S.One)


if __name__ == "__main__":
    check_covariances()
    check_chiral_tensors()
    check_source_determinant()
    check_propagator()
    print("Gauge covariance, chiral tensors, source determinant and propagator identities passed.")