"""Exact flat 3d N=2 chiral closure, with explicitly ordered odd symbols.

Run: python chiral_closure.py
Diagnostic: python chiral_closure.py --wrong-parameter-sign

No network, data file, spatial lattice or floating-point approximation.
The deliberate diagnostic changes ONLY rho2*rho1=-rho1*rho2 to a plus.
It is not a second valid supersymmetry convention.
"""
import argparse
from itertools import product
import json
import sympy as s

I, rt = s.I, s.sqrt(2)
C = s.Matrix([[0, -1], [1, 0]])
gamma = [s.diag(1, -1), s.Matrix([[0, -1], [-1, 0]]),
         s.Matrix([[0, I], [-I, 0]])]
zero = (0, 0, 0)
multi = sorted([a for a in product(range(3), repeat=3) if sum(a) <= 2],
               key=lambda a: (sum(a), a))


def jetname(name, a):
    prefix = ''.join(str(mu+1)*a[mu] for mu in range(3))
    return ('d'+prefix+'_' if prefix else '')+name


class Algebra:
    """An ordered odd word -> commuting SymPy coefficient dictionary.

    Odd word order: rho1 < rho2 < all psi1 jets < all psi2 jets.
    Products explicitly reorder words; a repeated odd generator is zero.
    """
    def __init__(self, wrong=False):
        self.wrong = wrong
        self.even = {(f, a):s.Symbol(jetname(f, a))
                     for f in ['phi', 'F'] for a in multi}
        self.odd = [('rho', 0, zero), ('rho', 1, zero)]
        self.odd += [('psi', j, a) for j in range(2) for a in multi]
        self.index = {spec:n for n,spec in enumerate(self.odd)}
        self.names = ['rho1', 'rho2']+[jetname('psi'+str(j+1), a)
                     for j in range(2) for a in multi]

    def clean(self, p):
        return {w:s.expand(c) for w,c in p.items() if s.expand(c) != 0}

    def add(self, *ps):
        result = {}
        for p in ps:
            for word, coefficient in p.items():
                result[word] = result.get(word, 0)+coefficient
        return self.clean(result)

    def scale(self, c, p):
        return self.clean({w:c*v for w,v in p.items()})

    def mul(self, a, b):
        result = {}
        for left, x in a.items():
            for right, y in b.items():
                if set(left) & set(right):
                    continue
                sign = 1
                for i in left:
                    for j in right:
                        if i > j and not (self.wrong and {i, j} == {0, 1}):
                            sign = -sign
                word = tuple(sorted(left+right))
                result[word] = result.get(word, 0)+sign*x*y
        return self.clean(result)

    def word(self, word):
        return {tuple(word):s.Integer(1)}

    def field(self, name, a=zero, component=0):
        if name == 'psi':
            return self.word([self.index[('psi', component, a)]])
        return {():self.even[(name, a)]}

    def rho(self, j):
        return self.word([j])

    def derivative(self, p, mu):
        def raisejet(a):
            b = list(a); b[mu] += 1
            if sum(b) > 2:
                raise ValueError('This calculation needs a jet above order two.')
            return tuple(b)
        def even_image(f, a):
            return self.field(f, raisejet(a))
        def odd_image(spec):
            f, j, a = spec
            return {} if f == 'rho' else self.field(f, raisejet(a), j)
        return self.derive(p, even_image, odd_image)

    def derive(self, p, even_image, odd_image):
        """Ordinary Leibniz rule: this operation is EVEN.

        Signs enter when each substituted odd word is put back in order.
        Inserting an additional odd-Leibniz sign here would be wrong.
        """
        answer = {}
        for word, coefficient in p.items():
            for spec, symbol in self.even.items():
                dc = s.diff(coefficient, symbol)
                if dc:
                    term = self.mul(even_image(*spec), self.word(word))
                    answer = self.add(answer, self.scale(dc, term))
            for position, index in enumerate(word):
                image = odd_image(self.odd[index])
                term = self.mul(self.word(word[:position]), image)
                term = self.mul(term, self.word(word[position+1:]))
                answer = self.add(answer, self.scale(coefficient, term))
        return answer

    def transform(self, p, family, label, parameter, auxiliary=True):
        rho = self.rho(parameter)
        def base(f, component=0):
            if family == 'Q':
                if f == 'phi':
                    row = label.T*C
                    out = self.add(*(self.scale(rt*row[j], self.field('psi', component=j))
                                     for j in range(2)))
                elif f == 'psi' and auxiliary:
                    out = self.scale(rt*label[component], self.field('F'))
                else:
                    out = {}
            else:
                if f == 'psi':
                    out = self.add(*(self.scale(-rt*I*(gamma[mu]*label)[component],
                            self.derivative(self.field('phi'), mu)) for mu in range(3)))
                elif f == 'F' and auxiliary:
                    out = self.add(*(self.scale(-rt*I*(label.T*C*gamma[mu])[j],
                            self.derivative(self.field('psi', component=j), mu))
                            for mu in range(3) for j in range(2)))
                else:
                    out = {}
            return self.mul(rho, out)
        cache = {}
        def image(f, a, j=0):
            key = (f, a, j)
            if key not in cache:
                out = base(f, j)
                for mu in range(3):
                    for _ in range(a[mu]):
                        out = self.derivative(out, mu)
                cache[key] = out
            return cache[key]
        return self.derive(p, lambda f,a:image(f,a),
                           lambda spec:{} if spec[0]=='rho' else image(spec[0],spec[2],spec[1]))

    def readable(self, p):
        return {' '.join(self.names[j] for j in word) or '1':str(c)
                for word,c in sorted(self.clean(p).items())}


def run(wrong=False):
    alg = Algebra(wrong)
    checks = {'ordered_algebra':0, 'off_shell_components':0,
              'off_shell_products':0, 'same_family':0, 'reduced_remainder':0,
              'polynomial_profile':0}
    def check(a, b, group):
        assert not alg.add(a, alg.scale(-1,b)), (group,alg.readable(a),alg.readable(b))
        checks[group] += 1
    z1,z2,t1,t2 = s.symbols('z1 z2 t1 t2')
    z, t = s.Matrix([z1,z2]), s.Matrix([t1,t2])
    K = [(z.T*C*g*t)[0] for g in gamma]
    r12 = alg.mul(alg.rho(0), alg.rho(1))
    phi, F = alg.field('phi'), alg.field('F')
    psi = [alg.field('psi', component=j) for j in range(2)]
    fields = [phi, *psi, F]
    def comm(p, auxiliary=True):
        first = alg.transform(alg.transform(p,'T',t,1,auxiliary),'Q',z,0,auxiliary)
        second = alg.transform(alg.transform(p,'Q',z,0,auxiliary),'T',t,1,auxiliary)
        return alg.add(first,alg.scale(-1,second))
    def translation(p):
        return alg.mul(r12, alg.add(*(alg.scale(2*I*K[mu],alg.derivative(p,mu))
                                     for mu in range(3))))

    if wrong:
        differences = [alg.add(comm(p),alg.scale(-1,translation(p))) for p in fields]
        assert not differences[0] and differences[1] and differences[2] and differences[3]
        return {'mode':'deliberately wrong parameter exchange',
                'rho2_rho1_equals_minus_rho1_rho2':False,
                'closure_matches_translation':dict(zip(['phi','psi1','psi2','F'],
                                                        [not p for p in differences])),
                'fermion_residues':[alg.readable(p) for p in differences[1:3]],
                'interpretation':'Expected diagnostic: the scalar alone misses the wrong parameter sign. Not a valid alternate SUSY algebra.'}

    for j in [0,1,alg.index[('psi',0,zero)],alg.index[('psi',1,zero)]]:
        check(alg.mul(alg.word([j]),alg.word([j])),{},'ordered_algebra')
    check(alg.mul(alg.rho(1),alg.rho(0)),alg.scale(-1,r12),'ordered_algebra')
    check(alg.mul(psi[0],alg.rho(0)),alg.scale(-1,alg.mul(alg.rho(0),psi[0])), 'ordered_algebra')
    for p in fields:
        check(comm(p),translation(p),'off_shell_components')
    products = [alg.mul(phi,psi[0]),alg.mul(psi[0],psi[1]),alg.mul(F,psi[1])]
    for p in products:
        check(comm(p),translation(p),'off_shell_products')
    # Both same-family brackets with independent odd parameters and labels.
    for family in ['Q','T']:
        for p in fields:
            first=alg.transform(alg.transform(p,family,t,1),family,z,0)
            second=alg.transform(alg.transform(p,family,z,0),family,t,1)
            check(first,second,'same_family')
    E = [alg.add(*(alg.scale(gamma[mu][a,j],alg.derivative(psi[j],mu))
                   for mu in range(3) for j in range(2))) for a in range(2)]
    projection = alg.add(*(alg.scale((t.T*C)[j],E[j]) for j in range(2)))
    R = [alg.mul(r12,alg.scale(-2*I*z[j],projection)) for j in range(2)]
    for j in range(2):
        check(alg.add(comm(psi[j],False),alg.scale(-1,translation(psi[j]))),
              R[j],'reduced_remainder')
    check(comm(phi,False),translation(phi),'reduced_remainder')
    check(alg.add(comm(products[1],False),alg.scale(-1,translation(products[1]))),
          alg.add(alg.mul(R[0],psi[1]),alg.mul(psi[0],R[1])), 'reduced_remainder')

    # Evaluate AFTER composing: psi=(theta1*u1^2,theta2*u1*u2).
    # Distinct theta1,theta2 are represented by the two zeroth-order psi
    # generators only for this final output; jets were independent above.
    u = s.symbols('u1:4')
    fs = [u[0]**2,u[0]*u[1]]
    label_values = {z1:1,z2:0,t1:0,t2:1}
    def profile(p):
        result = {}
        for word, coefficient in p.items():
            out = {():coefficient.subs(label_values)}
            for index in word:
                kind,j,a=alg.odd[index]
                if kind=='rho':
                    image=alg.rho(j)
                else:
                    f=fs[j]
                    for mu in range(3):
                        f=s.diff(f,u[mu],a[mu])
                    image=alg.scale(f,psi[j])
                out=alg.mul(out,image)
            result=alg.add(result,out)
        return result
    full = [profile(comm(p)) for p in psi]
    reduced = [profile(comm(p,False)) for p in psi]
    residue = [profile(r) for r in R]
    expected_full = [alg.mul(r12,alg.scale(4*I*u[0],psi[0])),
                     alg.mul(r12,alg.scale(2*I*u[1],psi[1]))]
    for a,b in zip(full,expected_full):check(a,b,'polynomial_profile')
    check(profile(comm(products[1])),alg.mul(r12,alg.scale(6*I*u[0]**2*u[1],products[1])), 'polynomial_profile')
    check(profile(comm(products[1],False)),alg.mul(r12,alg.scale(2*I*u[0]**2*u[1],products[1])), 'polynomial_profile')
    return {'mode':'correct ordered algebra','arithmetic':'exact symbolic polynomials',
            'checks':checks,'total':sum(checks.values()),
            'off_shell_closure':True,'reduced_fermion_closes_for_arbitrary_jets':False,
            'remainder':'-2*I*rho1*rho2*zeta*(tilde_zeta^T*C*gamma^mu*d_mu psi)',
            'profile':{'coordinates':'u1,u2,u3; reference scale set to one',
                       'psi':['theta1*u1**2','theta2*u1*u2'],
                       'word_dictionary':'psi1,psi2 in these output words denote theta1,theta2',
                       'full':[alg.readable(p) for p in full],
                       'reduced':[alg.readable(p) for p in reduced],
                       'reduced_minus_translation':[alg.readable(p) for p in residue]},
            'scope':'Flat free neutral chiral; constant labels; local jets. No action, boundary, interacting or curved-space test.'}


if __name__=='__main__':
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--wrong-parameter-sign',action='store_true')
    args=parser.parse_args()
    print(json.dumps(run(args.wrong_parameter_sign),indent=2))