"""LAB09: compact one-brane ABJM quotient, not a full quantum partition function.

Run from this directory: python one_brane_quotient.py
Changed input: --level 5. Deliberate bad identification: --wrong-independent-phases.
Exact cyclotomic polynomial arithmetic determines equivalence and invariance;
mpmath evaluates distances and geometric Gaussian integrals. No repository
import/network/random input.
"""
import argparse
import json
import sympy as sp
import mpmath as mp


def calculate(level, dps=50, wrong=False):
    mp.mp.dps = dps
    w = sp.Symbol('w')
    # When 4 divides k, i is already w^(k/4) in the chosen primitive embedding.
    # Otherwise Q(i) and Q(w) have trivial intersection over Q.
    domain = sp.QQ if level % 4 == 0 else sp.QQ_I
    cyclotomic = sp.Poly(sp.cyclotomic_poly(level, w), w, domain=domain)
    checks = 0

    def reduce(expr):
        if level % 4 == 0:
            expr = sp.sympify(expr).xreplace({sp.I: w**(level//4)})
        return sp.rem(sp.Poly(sp.expand(expr), w, domain=domain), cyclotomic).as_expr()

    def equal(a, b):
        nonlocal checks
        assert reduce(a-b) == 0, (a, b)
        checks += 1

    # The original compact basis is integral, not a half-sum convention.
    transform = sp.Matrix([[1, 1], [0, 1]])  # (a,ahat)=T(b,c)
    assert transform.det() == 1
    assert transform.T*sp.diag(level, -level)*transform == sp.Matrix([[level, level], [level, 0]])
    checks += 2
    for n1 in range(-2, 3):
        for n2 in range(-2, 3):
            assert transform*sp.Matrix([n1-n2, n2]) == sp.Matrix([n1, n2])
    checks += 1
    # Exact compact group projector, computed as a polynomial remainder rather
    # than by testing whether floating-point complex sums are close to zero.
    projected = []
    for charge in range(-2*level, 2*level+1):
        value = reduce(sum(w**((j*charge) % level) for j in range(level))/level)
        equal(value, int(charge % level == 0))
        if value == 1:
            projected.append(charge)

    X = [sp.Integer(1), sp.Integer(2), 1+sp.I, sp.Integer(0)]
    Z = [X[0], w*X[1], X[2], X[3]]  # only the second coordinate rotates
    orbit = [[reduce(w**j*x) for x in X] for j in range(level)]
    assert len({tuple(v) for v in orbit}) == level
    checks += 1
    stabilizer = sum(all(reduce((w**j-1)*x) == 0 for x in X) for j in range(level))
    assert stabilizer == 1
    checks += 1
    equal(w**level, 1)
    # Pure powers miss this comparison; one mixed invariant repairs it.
    assert all(reduce(z**level-x**level) == 0 for x, z in zip(X, Z))
    checks += 1
    mixed_X = reduce(X[0]**(level-1)*X[1])
    mixed_Z = reduce(Z[0]**(level-1)*Z[1])
    equivalent = any(all(reduce(z-w**j*x) == 0 for x, z in zip(X, Z)) for j in range(level))
    assert equivalent == (level == 1)
    assert (reduce(mixed_Z-mixed_X) == 0) == equivalent
    checks += 2
    for j in range(level):
        equal((w**j*X[0])**(level-1)*(w**j*X[1]), mixed_X)

    omega = mp.exp(2j*mp.pi/level)
    xn = [mp.mpc(1), mp.mpc(2), mp.mpc(1, 1), mp.mpc(0)]
    zn = [xn[0], omega*xn[1], xn[2], xn[3]]
    distances = [sum(abs(x-omega**j*z)**2 for x, z in zip(xn, zn)) for j in range(level)]
    minimum = min(distances)
    # Independent phases are an intentional wrong gauge group. This distance
    # is explicitly a diagnostic and is never substituted for the true one.
    wrong_distance = sum(min(abs(x-omega**j*z)**2 for j in range(level)) for x, z in zip(xn, zn))
    wrong_equivalent = all(any(reduce(z-w**j*x) == 0 for j in range(level))
                           for x, z in zip(X, Z))
    detected = not equivalent and wrong_equivalent
    if wrong:
        assert detected, 'this diagnostic requires level > 1'
        assert minimum > 0 and wrong_distance < mp.power(10, -dps+15)

    # The monomial-map patch uses W=X1^k, Rj=Xj/X1. Its four complex
    # coordinates separate orbits where X1 != 0; mixed entries encode Rj.
    patch = [reduce(X[0]**level)] + [reduce(x/X[0]) for x in X[1:]]
    flux_examples = [{'node_flux': [n1, n2], 'b_flux': n1-n2, 'c_flux': n2,
                      'sum_flux': n1+n2} for n1, n2 in [(1, 0), (0, 1), (1, 1)]]
    # An independent target-space measure check. Use a fundamental angular
    # sector in X1, of width 2*pi/k, not separate sectors in all four X's.
    # Ordinary convergent one-dimensional radial quadrature supplies the
    # Gaussian integrals; it is not a quantum ABJM partition function.
    def radial(power):
        return 2*mp.pi*mp.quad(lambda r: r**(2*power+1)*mp.exp(-r*r), [0, 1, mp.inf])
    g0, g1, gp = radial(0), radial(1), radial(level-1)
    gaussian_volume = g0**4/level
    mixed_moment = (gp*g1*g0**2/level)/gaussian_volume
    gaussian_errors = [abs(gaussian_volume-mp.pi**4/level)/max(1, mp.pi**4/level),
                       abs(mixed_moment-mp.factorial(level-1))/max(1, mp.factorial(level-1))]
    assert max(gaussian_errors) < mp.power(10, -dps+15)
    return {
        'model': 'compact U(1)_k x U(1)_(-k), no common-center quotient, zero masses/FI',
        'level': level, 'decimal_precision': dps,
        'exact_grouped_assertions': checks,
        'dual_period': '2*pi', 'residual_angle_over_2pi': [str(sp.Rational(j, level)) for j in range(level)],
        'original_flux_examples': flux_examples,
        'projected_charges_in_test_window': projected,
        'nonzero_orbit_size': len(orbit), 'nonzero_stabilizer_order': stabilizer,
        'origin_stabilizer_order': level,
        'pure_kth_powers_agree': True, 'common_phase_equivalent': equivalent,
        'mixed_invariant_X': str(mixed_X), 'mixed_invariant_Z': str(mixed_Z),
        'patch_coordinates_X': [str(z) for z in patch],
        'covering_distance_squared_by_image': [mp.nstr(x, 14) for x in distances],
        'quotient_distance_squared': mp.nstr(minimum, 14),
        'quotient_distance': mp.nstr(mp.sqrt(minimum), 14),
        'mode': 'deliberate wrong independent phases' if wrong else 'physical common phase',
        'wrong_independent_phases_detected': bool(detected) if wrong else None,
        'target_gaussian_integral': mp.nstr(gaussian_volume, 14),
        'normalized_mixed_invariant_squared': mp.nstr(mixed_moment, 14),
        'gaussian_scaled_error': mp.nstr(max(gaussian_errors), 8),
        'interpretation': 'Exact orbit test decides equivalence; rounded distance is only geometric output.'
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--level', type=int, default=3, help='positive integer k (default 3)')
    parser.add_argument('--dps', type=int, default=50, help='distance precision, at least 30')
    parser.add_argument('--wrong-independent-phases', action='store_true')
    args = parser.parse_args()
    if args.level < 1 or args.dps < 30 or (args.wrong_independent_phases and args.level == 1):
        parser.error('require k >= 1, dps >= 30; the deliberate failure requires k > 1')
    print(json.dumps(calculate(args.level, args.dps, args.wrong_independent_phases), indent=2))


if __name__ == '__main__':
    main()