"""Cyclic two-impurity scalar words: fixed-length planar one-loop N=4 SYM.

Python 3.13.7; SymPy 1.13.3; mpmath 1.3.0. No repository imports.
H acts on the operator column. Coefficient columns use H.T.
The auxiliary adjacent-swap walk is a diagnostic, not physical time evolution.
"""
import argparse
from itertools import combinations
import json
from math import comb, gcd
import sympy as sp
import mpmath as mp


def words(length):
    # 0=Z, 1=W, with exactly two W fields.
    return [tuple(int(i in pair) for i in range(length))
            for pair in combinations(range(length), 2)]


def rotations(word):
    return {word[j:]+word[:j] for j in range(len(word))}


def representative(word):
    return min(rotations(word))


def swap(word, j):
    output = list(word)
    k = (j+1) % len(word)
    output[j], output[k] = output[k], output[j]
    return tuple(output)


def cyclic_data(length):
    basis = sorted({representative(w) for w in words(length)})
    index = {w:i for i,w in enumerate(basis)}
    sizes = [len(rotations(w)) for w in basis]
    H = sp.zeros(len(basis))
    for i,w in enumerate(basis):
        for j in range(length):
            H[i,i] += 1
            H[i,index[representative(swap(w,j))]] -= 1
    gram = sp.diag(*[length//d for d in sizes])
    return basis, sizes, H, gram


def burnside_count(length):
    total = 0
    for shift in range(length):
        cycles = gcd(length,shift)
        size = length//cycles
        if 2 % size == 0:
            total += comb(cycles,2//size)
    return total//length


def edge_laplacian(length, basis):
    """Count every labeled word, independently of the representative action."""
    index = {w:i for i,w in enumerate(basis)}
    conductance = sp.zeros(len(basis))
    for word in words(length):
        i = index[representative(word)]
        for j in range(length):
            k = index[representative(swap(word,j))]
            if k != i:
                conductance[i,k] += 1
    return sp.diag(*[sum(conductance.row(i)) for i in range(len(basis))])-conductance


def walk_row(length, start, basis, time, dps):
    """Poisson mixture of elementary swaps on labeled words, not exp(H)."""
    allwords = words(length)
    lookup = {w:i for i,w in enumerate(allwords)}
    targets = [[lookup[swap(w,j)] for j in range(length)] for w in allwords]
    orbit = rotations(basis[start])
    probability = [mp.mpf(int(w in orbit))/len(orbit) for w in allwords]
    rate = length*time
    weight = mp.exp(-rate)
    mixture = [weight*p for p in probability]
    step = 0
    while True:
        # Bound all omitted Poisson weights by a geometric series.
        next_weight = weight*rate/(step+1)
        if step+2 > rate:
            tail = next_weight/(1-rate/(step+2))
            if tail < mp.power(10,-dps-5):
                break
        new = [mp.mpf(0) for _ in allwords]
        for i,p in enumerate(probability):
            for target in targets[i]:
                new[target] += p/length
        probability = new
        step += 1
        weight *= rate/step
        for i,p in enumerate(probability):
            mixture[i] += weight*p
    result = [mp.fsum(mixture[i] for i,w in enumerate(allwords)
                     if representative(w)==b) for b in basis]
    return result, tail, step


def run(length=4, wrong_gram=False, dps=50):
    basis,sizes,H,G = cyclic_data(length)
    D = sp.diag(*sizes)
    n = len(basis)
    coefficients = sp.Matrix(sizes)
    common = sp.igcd(*sizes)
    primitive = coefficients/common
    E = sp.Symbol('E')
    polynomial = H.charpoly(E).as_expr()
    eigenvalues = sorted(H.eigenvals().items(),key=lambda pair:float(pair[0].evalf()))
    eigenvectors = []
    for value,_,vectors in sorted(H.T.eigenvects(),key=lambda item:float(item[0].evalf())):
        normalized=[]
        for vector in vectors:
            first=next(x for x in vector if x!=0)
            vector=vector.applyfunc(lambda x:sp.simplify(x/first))
            if all(x.is_Rational for x in vector):
                vector*=sp.ilcm(*[sp.denom(x) for x in vector])
            normalized.append(vector)
        eigenvectors.append((value,normalized))
    B = edge_laplacian(length,basis)
    tree = B.minor_submatrix(0,0).det()
    product = sp.Rational(sum(sizes),sp.prod(sizes))*tree
    checks = {
        'burnside_basis_count':burnside_count(length)==n,
        'orbit_sizes_partition_words':sum(sizes)==comb(length,2),
        'gram_hermiticity':H*G==G*H.T,
        'protected_coefficient_vector':H.T*coefficients==sp.zeros(n,1),
        'independent_edge_laplacian':B==D*H,
        'symmetric_edge_laplacian':B==B.T,
        'cofactor_product':sp.simplify(product-(-1)**(n-1)*sp.diff(polynomial,E).subs(E,0))==0,
        'connected_unique_zero_mode':H.rank()==n-1,
        'coefficient_eigenvectors':all((H.T*v-value*v).applyfunc(sp.simplify)==sp.zeros(n,1)
                                      for value,vectors in eigenvectors for v in vectors),
        'distinct_eigenvector_orthogonality':all(sp.simplify((v.T*G*w)[0])==0
             for value,vectors in eigenvectors for other,others in eigenvectors if value!=other
             for v in vectors for w in others),
    }
    used_G = sp.eye(n) if wrong_gram else G
    residual = H*used_G-used_G*H.T
    error_detected = wrong_gram and residual != sp.zeros(n)
    # All protected coefficients are kept; changing a metric cannot change H's spectrum.
    if wrong_gram and not error_detected:
        raise ValueError('This length has equal stabilizers: choose length 4 or 6 for this control.')

    mp.mp.dps=dps
    hm=mp.matrix([[int(H[i,j]) for j in range(n)] for i in range(n)])
    t=mp.mpf('0.4')
    reduced=mp.expm(-t*hm)
    errors=[]; tails=[]; steps=[]
    for i in range(n):
        actual,tail,count=walk_row(length,i,basis,t,dps)
        errors.append(max(abs(actual[j]-reduced[i,j]) for j in range(n)))
        tails.append(tail); steps.append(count)
    checks['independent_poisson_evolution']=max(errors)<mp.power(10,-dps+10)
    if not all(checks.values()):
        raise AssertionError(checks)

    return {
        'regime':'4d N=4 SU(N), fixed L before N -> infinity, planar one loop only',
        'length':length,'W_count':2,
        'basis':[''.join('W' if b else 'Z' for b in w) for w in basis],
        'orbit_sizes':sizes,'gram_diagonal':[int(G[i,i]) for i in range(n)],
        'operator_column_H':[[int(x) for x in H.row(i)] for i in range(n)],
        'coefficient_equation':'H.T * c = E * c',
        'characteristic_polynomial':str(sp.factor(polynomial)),
        'eigenvalues':[{'E':str(v),'multiplicity':m} for v,m in eigenvalues],
        'coefficient_eigenvectors':[{'E':str(value),'vectors':[[str(x) for x in v] for v in vectors]}
                                   for value,vectors in eigenvectors],
        'anomalous_dimension':'gamma = lambda * E / (8*pi**2) + O(lambda**2), planar',
        'protected_coefficients':[int(x) for x in primitive],
        'protected_free_norm':int((primitive.T*G*primitive)[0]),
        'graph_cofactor':int(tree),'nonzero_eigenvalue_product':str(product),
        'auxiliary_walk':{'time':'0.4','digits':dps,'max_error':mp.nstr(max(errors),12),
                          'poisson_tail_bound':mp.nstr(max(tails),12),'max_steps':max(steps)},
        'mode':'wrong Gram diagnostic' if wrong_gram else 'correct cyclic norm',
        'used_metric_residual':[[int(x) for x in residual.row(i)] for i in range(n)],
        'intentional_error_detected':bool(error_detected),
        'checks':checks,'all_checks_pass':all(checks.values()),
    }


if __name__=='__main__':
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--length',type=int,choices=range(4,8),default=4)
    parser.add_argument('--wrong-gram',action='store_true')
    parser.add_argument('--dps',type=int,choices=range(40,81),default=50)
    args=parser.parse_args()
    if args.wrong_gram and args.length % 2:
        parser.error('Equal stabilizers at this odd length: choose length 4 or 6 for --wrong-gram.')
    print(json.dumps(run(args.length,args.wrong_gram,args.dps),indent=2))
