# Exact check of the odd Wick identity for one pair (alpha, beta) of odd normal coordinates.
# Generators: a, b (output), a2, b2 (input). Grassmann algebra over sympy expressions.
import itertools, sympy as sp

GENS = ['a', 'b', 'a2', 'b2']
def mono_mul(m1, m2):
    # m1, m2: tuples of generator indices in increasing order; returns (sign, monomial) or (0, None)
    if set(m1) & set(m2):
        return 0, None
    seq = list(m1) + list(m2)
    sign = 1
    # bubble sort counting transpositions
    for i in range(len(seq)):
        for j in range(len(seq) - 1 - i):
            if seq[j] > seq[j + 1]:
                seq[j], seq[j + 1] = seq[j + 1], seq[j]
                sign = -sign
    return sign, tuple(seq)

class G:
    def __init__(self, d=None):
        self.d = {k: v for k, v in (d or {}).items() if sp.simplify(v) != 0}
    @staticmethod
    def gen(name):
        return G({(GENS.index(name),): sp.Integer(1)})
    @staticmethod
    def sc(c):
        return G({(): sp.sympify(c)})
    def __add__(self, o):
        o = o if isinstance(o, G) else G.sc(o)
        d = dict(self.d)
        for k, v in o.d.items():
            d[k] = d.get(k, 0) + v
        return G(d)
    __radd__ = __add__
    def __neg__(self):
        return G({k: -v for k, v in self.d.items()})
    def __sub__(self, o):
        return self + (-(o if isinstance(o, G) else G.sc(o)))
    def __mul__(self, o):
        o = o if isinstance(o, G) else G.sc(o)
        d = {}
        for k1, v1 in self.d.items():
            for k2, v2 in o.d.items():
                s, m = mono_mul(k1, k2)
                if s:
                    d[m] = d.get(m, 0) + s * v1 * v2
        return G(d)
    __rmul__ = lambda self, o: G.sc(o) * self
    def exp(self):
        # nilpotent part: exp via series (degree <= 4)
        c0 = self.d.get((), 0)
        n = self - c0
        res, term = G.sc(1), G.sc(1)
        for k in range(1, 6):
            term = term * n * sp.Rational(1, k)
            res = res + term
        return res * sp.exp(c0)
    def coeff(self, mono):
        return sp.simplify(self.d.get(tuple(sorted(GENS.index(g) for g in mono)), 0) *
                           (1 if list(mono) == sorted(mono, key=GENS.index) else -1))

a, b, a2, b2 = (G.gen(n) for n in GENS)
W, m = sp.symbols('W m', positive=True)
A = W**2 + m / 2
c = sp.sqrt(W**2 * m + m**2 / 4)
mu = W**2 / (A + c)
N = -1 / W**2

def berezin_in(F):
    # integral over (a2, b2) with the convention  int a2 b2 da2 db2 = 1:
    # return the element of the (a, b) algebra given by the coefficient of a2*b2 (moved to the left)
    out = {}
    i2, j2 = GENS.index('a2'), GENS.index('b2')
    for k, v in F.d.items():
        if i2 in k and j2 in k:
            rest = tuple(x for x in k if x not in (i2, j2))
            # write monomial as rest * a2 b2 ; k is sorted with a, b before a2, b2 -> already rest + (i2, j2)
            out[rest] = out.get(rest, 0) + v
    return G(out)

def K(f):
    # odd kernel with the quadratic half-weights, conjugated by gamma = exp(-c a b)
    ker = (G.sc(-W**2) * (a - a2) * (b - b2) - G.sc(m / 2) * (a * b + a2 * b2)).exp() * N
    gam_in = (G.sc(-c) * a2 * b2).exp()
    gam_out_inv = (G.sc(c) * a * b).exp()
    return gam_out_inv * berezin_in(ker * gam_in * f)

def subst_in(f):
    # rename output generators of f (a, b) to input generators (a2, b2)
    d = {}
    for k, v in f.d.items():
        newk = tuple(sorted(GENS.index({'a': 'a2', 'b': 'b2'}[GENS[x]]) for x in k))
        d[newk] = v
    return G(d)

def full_diff(img, expect):
    # all Grassmann components of img - expect, so that an extra component of the image
    # (for example an alpha*beta term in the image of 1) is detected
    keys = set(img.d) | set(expect.d)
    return {k: sp.simplify(sp.radsimp(img.d.get(k, 0) - expect.d.get(k, 0))) for k in keys}


if __name__ == '__main__':
    # heat kernel alone maps 1 to 1
    heat = (G.sc(-W**2) * (a - a2) * (b - b2)).exp() * N
    print('heat 1 =', {k: sp.simplify(v) for k, v in berezin_in(heat).d.items()})
    print('mu =', sp.simplify(mu), '  mu^-1 =', sp.simplify(1 / mu))
    # e:tube-odd-wick: 1 -> mu^{-1}, alpha -> alpha, beta -> beta, alpha beta + 1/(2c) -> mu (alpha beta + 1/(2c)),
    # compared as full images (every Grassmann component), not as ratios on the monomials of the input
    targets = [('1', G.sc(1), G.sc(1) * (1 / mu)),
               ('a', a, a),
               ('b', b, b),
               ('ab+1/(2c)', a * b + G.sc(1 / (2 * c)), (a * b + G.sc(1 / (2 * c))) * mu)]
    ok = True
    for name, f, expect in targets:
        img = K(subst_in(f))
        diff = full_diff(img, expect)
        num = max([abs(complex(sp.N(v.subs({W: 3.7, m: 0.83})))) for v in diff.values()] + [0.0])
        print(name, '-> image minus expected image, all components:', diff, ' numeric max |diff| =', num)
        ok = ok and (all(v == 0 for v in diff.values()) or num < 1e-12)
    print('ALL IMAGES AGREE' if ok else 'MISMATCH')