"""Super check of V_F(0) = 0, and of the Taylor coefficients, with a finite Grassmann algebra and floating-point complex coefficients.

Elements: dict {(eexp, omono): complex}, eexp = exponents of even jet variables
(total degree <= 2, jets around a base point x0), omono = sorted tuple of odd generators.

V_+ has parities par_p (0 = even, 1 = odd), V_- has parities par_m.
U_sigma entries: even jet variable if parities agree, odd generator otherwise.
S_+ = a I + U_+, S_- = b I + U_-.  L_S A = S_+ A - A S_- on H = Hom(V_-, V_+),
written as an even supermatrix on vec(A) (entries A_{i alpha}, parity |i| + |alpha|).
F = Ber(L_S).  Laplacian of the supertrace metric Str(X Y) on each block,
odd coefficient kappa fixed by |grad Q|^2 = 4 Q on the odd part of Q = Str X^2.
"""
import itertools, cmath, random

DEG = 2


class SA:
    nE = 0

    def __init__(self, t=None):
        self.t = {} if t is None else {k: v for k, v in t.items() if abs(v) > 1e-300}

    @staticmethod
    def const(c):
        return SA({((0,) * SA.nE, ()): complex(c)})

    @staticmethod
    def evar(k, x0=0.0):
        e = [0] * SA.nE
        e[k] = 1
        return SA({((0,) * SA.nE, ()): complex(x0), (tuple(e), ()): 1.0})

    @staticmethod
    def ogen(k):
        return SA({((0,) * SA.nE, (k,)): 1.0})

    def __add__(self, o):
        o = o if isinstance(o, SA) else SA.const(o)
        t = dict(self.t)
        for k, v in o.t.items():
            t[k] = t.get(k, 0) + v
        return SA(t)

    __radd__ = __add__

    def __neg__(self):
        return SA({k: -v for k, v in self.t.items()})

    def __sub__(self, o):
        return self + (-(o if isinstance(o, SA) else SA.const(o)))

    def __rsub__(self, o):
        return SA.const(o) - self

    def __mul__(self, o):
        if not isinstance(o, SA):
            return SA({k: v * o for k, v in self.t.items()})
        t = {}
        for (e1, m1), v1 in self.t.items():
            for (e2, m2), v2 in o.t.items():
                if sum(e1) + sum(e2) > DEG:
                    continue
                if set(m1) & set(m2):
                    continue
                seq = list(m1) + list(m2)
                # sign of sorting permutation
                sgn = 1
                for i in range(len(seq)):
                    for j in range(i + 1, len(seq)):
                        if seq[i] > seq[j]:
                            sgn = -sgn
                key = (tuple(a + b for a, b in zip(e1, e2)), tuple(sorted(seq)))
                t[key] = t.get(key, 0) + sgn * v1 * v2
        return SA(t)

    __rmul__ = __mul__

    def body(self):
        return self.t.get(((0,) * SA.nE, ()), 0)

    def inv(self):
        c = self.body()
        n = self * (1 / c) - 1
        res = SA.const(1)
        p = SA.const(1)
        for _ in range(DEG + 40):
            p = p * (-n)
            if not p.t:
                break
            res = res + p
        return res * (1 / c)

    def d_even(self, k):
        t = {}
        for (e, m), v in self.t.items():
            if e[k] > 0:
                e2 = list(e)
                e2[k] -= 1
                key = (tuple(e2), m)
                t[key] = t.get(key, 0) + v * e[k]
        return SA(t)

    def d_odd(self, k):  # left derivative
        t = {}
        for (e, m), v in self.t.items():
            if k in m:
                pos = m.index(k)
                key = (e, m[:pos] + m[pos + 1:])
                t[key] = t.get(key, 0) + (-1) ** pos * v
        return SA(t)

    def at_base(self):
        """Grassmann polynomial (even jet variables set to 0)."""
        z = (0,) * SA.nE
        return SA({(e, m): v for (e, m), v in self.t.items() if e == z})

    def maxabs(self):
        return max([abs(v) for v in self.t.values()] + [0])


def det_even(M):
    """Leibniz determinant for a square matrix of even (commuting) elements."""
    n = len(M)
    if n == 0:
        return SA.const(1)
    res = SA.const(0)
    for perm in itertools.permutations(range(n)):
        sgn = 1
        for i in range(n):
            for j in range(i + 1, n):
                if perm[i] > perm[j]:
                    sgn = -sgn
        term = SA.const(sgn)
        for i in range(n):
            term = term * M[i][perm[i]]
        res = res + term
    return res


def inv_even(M):
    n = len(M)
    d = det_even(M)
    di = d.inv()
    out = [[None] * n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            minor = [[M[r][c] for c in range(n) if c != i] for r in range(n) if r != j]
            out[i][j] = det_even(minor) * ((-1) ** (i + j)) * di
    return out


def matmul(A, B):
    n, m, k = len(A), len(B), len(B[0]) if B else 0
    return [[sum((A[i][l] * B[l][j] for l in range(m)), SA.const(0)) for j in range(k)] for i in range(n)]


def ber(M, par):
    ev = [i for i in range(len(par)) if par[i] == 0]
    od = [i for i in range(len(par)) if par[i] == 1]
    A = [[M[i][j] for j in ev] for i in ev]
    B = [[M[i][j] for j in od] for i in ev]
    C = [[M[i][j] for j in ev] for i in od]
    D = [[M[i][j] for j in od] for i in od]
    if not od:
        return det_even(A)
    Di = inv_even(D)
    if ev:
        BDC = matmul(matmul(B, Di), C)
        Sch = [[A[i][j] - BDC[i][j] for j in range(len(ev))] for i in range(len(ev))]
        num = det_even(Sch)
    else:
        num = SA.const(1)
    return num * det_even(D).inv()


def setup(par_p, par_m, x0p=None, x0m=None):
    """Return variable maps and S_+, S_- entries."""
    evars, ovars = [], []  # entries (block, i, j)
    for blk, par in (('+', par_p), ('-', par_m)):
        for i in range(len(par)):
            for j in range(len(par)):
                (evars if par[i] == par[j] else ovars).append((blk, i, j))
    SA.nE = len(evars)
    return evars, ovars


def build(par_p, par_m, a, b, x0):
    evars, ovars = setup(par_p, par_m)
    ent = {}
    for k, key in enumerate(evars):
        ent[key] = SA.evar(k, x0.get(key, 0.0))
    for k, key in enumerate(ovars):
        ent[key] = SA.ogen(k)
    def U(blk, i, j):
        return ent[(blk, i, j)]
    Sp = [[U('+', i, j) + (a if i == j else 0) for j in range(len(par_p))] for i in range(len(par_p))]
    Sm = [[U('-', i, j) + (b if i == j else 0) for j in range(len(par_m))] for i in range(len(par_m))]
    return evars, ovars, ent, Sp, Sm


def L_matrix(Sp, Sm, par_p, par_m, sign_minus=1):
    idx = [(i, al) for i in range(len(par_p)) for al in range(len(par_m))]
    parH = [(par_p[i] + par_m[al]) % 2 for (i, al) in idx]
    n = len(idx)
    M = [[SA.const(0) for _ in range(n)] for _ in range(n)]
    for r, (i, al) in enumerate(idx):
        for c, (j, be) in enumerate(idx):
            v = SA.const(0)
            if al == be:
                v = v + Sp[i][j]
            if i == j:
                sgn = (-1) ** (((par_p[i] + par_m[be]) * (par_m[be] + par_m[al])) % 2)
                v = v - Sm[be][al] * sgn
            M[r][c] = v
    return M, parH


def laplacian(f, evars, ovars, par_p, par_m, kappa):
    res = SA.const(0)
    eidx = {key: k for k, key in enumerate(evars)}
    oidx = {key: k for k, key in enumerate(ovars)}
    for blk, par in (('+', par_p), ('-', par_m)):
        n = len(par)
        for i in range(n):
            for j in range(n):
                if par[i] == par[j]:
                    sgn = 1 if par[i] == 0 else -1
                    res = res + f.d_even(eidx[(blk, j, i)]).d_even(eidx[(blk, i, j)]) * sgn
                elif par[i] == 0:  # i even, j odd: odd pair (i,j),(j,i)
                    res = res + f.d_odd(oidx[(blk, i, j)]).d_odd(oidx[(blk, j, i)]) * kappa
    return res


def str_of(M, par):
    return sum((M[i][i] * (1 if par[i] == 0 else -1) for i in range(len(par))), SA.const(0))


def find_kappa(par):
    """Fix kappa by |grad Q|^2 = 4Q on the odd part of Q = Str X^2 for gl(par)."""
    evars, ovars = setup(par, [])
    ent = {}
    for k, key in enumerate(evars):
        ent[key] = SA.evar(k)
    for k, key in enumerate(ovars):
        ent[key] = SA.ogen(k)
    n = len(par)
    X = [[ent[('+', i, j)] for j in range(n)] for i in range(n)]
    Qodd = SA.const(0)
    Q = SA.const(0)
    for i in range(n):
        for j in range(n):
            term = X[i][j] * X[j][i] * (1 if par[i] == 0 else -1)
            Q = Q + term
            if par[i] != par[j]:
                Qodd = Qodd + term
    out = {}
    for kappa in (1.0, -1.0, 2.0, -2.0, 0.5, -0.5):
        lap = lambda f: laplacian(f, evars, ovars, par, [], kappa)
        grad2 = (lap(Qodd * Qodd) - Qodd * lap(Qodd) * 2) * 0.5
        err = (grad2 - Qodd * 4).maxabs()
        out[kappa] = err
    kappa = min(out, key=out.get)
    lap = lambda f: laplacian(f, evars, ovars, par, [], kappa)
    m = par.count(0); nn = par.count(1)
    LQ = lap(Q).at_base()
    return kappa, out[kappa], LQ.body(), 2 * (m - nn) ** 2, (LQ - LQ.body()).maxabs()


def run(par_p, par_m, s=0.75, x0=None, kappa=-2.0, label=''):
    a, b = -1j * s, 1j * s
    gamma = 1 / (a - b)
    x0 = x0 or {}
    evars, ovars, ent, Sp, Sm = build(par_p, par_m, a, b, x0)
    M, parH = L_matrix(Sp, Sm, par_p, par_m)
    F = ber(M, parH)
    lap = lambda f: laplacian(f, evars, ovars, par_p, par_m, kappa)
    LF = lap(F).at_base()
    LF2 = lap(F * F).at_base()
    Fb = F.at_base()
    Fi = Fb.inv()
    VF = LF * Fi  # V_F at base point, full Grassmann polynomial
    grad2_logF = (LF2 - Fb * LF * 2) * (Fi * Fi) * 0.5
    lap_logF = VF - grad2_logF
    u = par_p.count(0) - par_p.count(1)
    v = par_m.count(0) - par_m.count(1)
    pred = gamma ** 2 * u * v * (u + v)
    print(f'{label} V+ = {par_p.count(0)}|{par_p.count(1)}, V- = {par_m.count(0)}|{par_m.count(1)}, '
          f'u = {u}, v = {v}, F(base) body = {Fb.body():.6g}')
    if not x0:
        print(f'   F(0) = {Fb.body():.6g}   predicted (-2is)^(uv) = {(a - b) ** (u * v):.6g}')
        print(f'   Delta log F(0) = {lap_logF.body():.10g}   predicted {-pred:.10g}')
        print(f'   |grad log F(0)|^2 = {grad2_logF.body():.10g}   predicted {pred:.10g}')
    print(f'   max |coefficient| of V_F at base point (all Grassmann monomials): {VF.maxabs():.3g}')
    return F, evars, ovars, ent


def taylor_check(par_p, par_m):
    """Check Str L_U = v Str U+ - u Str U-, Str L_U^2 = v Str U+^2 + u Str U-^2 - 2 Str U+ Str U-."""
    evars, ovars, ent, Up, Um = build(par_p, par_m, 0.0, 0.0, {})
    M, parH = L_matrix(Up, Um, par_p, par_m)
    u = par_p.count(0) - par_p.count(1)
    v = par_m.count(0) - par_m.count(1)
    strL = str_of(M, parH)
    strL2 = str_of(matmul(M, M), parH)
    sp = str_of(Up, par_p); sm = str_of(Um, par_m)
    sp2 = str_of(matmul(Up, Up), par_p); sm2 = str_of(matmul(Um, Um), par_m)
    e1 = (strL - (sp * v - sm * u)).maxabs()
    e2 = (strL2 - (sp2 * v + sm2 * u - sp * sm * 2)).maxabs()
    print(f'   Taylor identities: |Str L_U - (v StrU+ - u StrU-)| = {e1:.2g},'
          f' |Str L_U^2 - (...)| = {e2:.2g}')


if __name__ == '__main__':
    for par in ([0, 1], [0, 0, 1], [0, 1, 1], [0, 0, 1, 1]):
        k, err, LQ, pred, rest = find_kappa(par)
        print(f'gl({par.count(0)}|{par.count(1)}): kappa = {k} (|grad Q|^2-4Q err {err:.2g}); '
              f'Delta Str X^2 = {LQ.real:.6g}, 2 sdim(gl) = {pred}, nonconstant part {rest:.2g}')
    print()
    cases = [
        ([0], [0, 1, 1], 'k=1, p=0:'),
        ([0, 1], [0, 1], 'k=1, p=1:'),
        ([0, 1, 1], [0], 'k=1, p=2:'),
        ([0, 0, 1], [0], 'general (u+v != 0):'),
        ([0, 1, 1], [0, 0], 'general (u+v != 0):'),
        ([0, 0], [0, 1], 'general (u+v != 0):'),
    ]
    for pp, pm, lab in cases:
        run(pp, pm, label=lab)
        taylor_check(pp, pm)
    print()
    random.seed(3)
    for pp, pm, lab in [([0, 1], [0, 1], 'random base point'), ([0], [0, 1, 1], 'random base point'),
                        ([0, 0, 1], [0], 'random base point')]:
        evars, ovars = setup(pp, pm)
        x0 = {key: complex(random.uniform(-.2, .2), random.uniform(-.2, .2)) for key in evars}
        run(pp, pm, x0=x0, label=lab)