"""Gaussian Wilson-loop lab, susy.qft.org; Python 3.13.7 tested.

Independent Hermitian samples, ordinary fundamental trace / N, lambda=g^2 N.
No Markov chain. Standard errors are statistical estimates, not rigorous bounds.
Run with --help. JSON goes to stdout; importing this file does not run a sample.
"""
import argparse
import json
import math
import platform
import mpmath as mp
import numpy as np


def exact_loop(n, coupling, traceless=False):
    """Positive finite Laguerre sum; no subtraction of alternating large terms."""
    lam = mp.mpf(coupling)
    term = mp.mpf(1)
    terms = [term]
    for j in range(n-1):
        term *= lam*(n-j-1)/(4*n*(j+1)*(j+2))
        terms.append(term)
    answer = mp.exp(lam/(8*n))*mp.fsum(terms)
    return answer*mp.exp(-lam/(8*n*n)) if traceless else answer


def planar_loop(coupling):
    lam = mp.mpf(coupling)
    return 2*mp.besseli(1, mp.sqrt(lam))/mp.sqrt(lam) if lam else mp.mpf(1)


def analytic_checks():
    """Independent radial integration at N=2 and elementary Abelian checks."""
    residuals = []
    tolerance = mp.power(10, -(mp.mp.dps-12))
    def compare(a, b):
        error = abs(a-b)/max(mp.mpf(1), abs(a), abs(b))
        assert error < tolerance, mp.nstr(error)
        residuals.append(error)
    for n in (1, 2, 4, 10):
        compare(exact_loop(n, 0), 1)
        compare(exact_loop(n, 0, True), 1)
    for lam in map(mp.mpf, ('1', '4', '16')):
        compare(exact_loop(1, lam), mp.exp(lam/8))
        compare(exact_loop(1, lam, True), 1)  # trivial traceless 1x1 matrix
        b = 8/lam
        denom = mp.quad(lambda r:r*r*mp.exp(-b*r*r), [0, mp.inf])
        num = mp.quad(lambda r:r*r*mp.exp(-b*r*r)*mp.cosh(r), [0, mp.inf])
        radial = num/denom
        compare(radial, exact_loop(2, lam, True))
        compare(radial, (1+lam/16)*mp.exp(lam/32))
        # Independently integrate the U(2) center, variance lambda/16.
        center = mp.quad(lambda c:mp.exp(c-8*c*c/lam), [-mp.inf, 0, mp.inf])
        center /= mp.sqrt(mp.pi*lam/8)
        compare(radial*center, exact_loop(2, lam))
    compare(planar_loop(0), 1)
    return {'comparisons':len(residuals), 'decimal_digits':mp.mp.dps,
            'maximum_scaled_residual':mp.nstr(max(residuals), 12),
            'tolerance':mp.nstr(tolerance, 5),
            'method':'N=2 radial quadrature plus independent center integral; N=1 and zero coupling'}


def statistics(values, target=None):
    mean = float(np.mean(values))
    se = float(np.std(values, ddof=1)/math.sqrt(len(values)))
    result = {'mean':mean, 'standard_error':se}
    if target is not None:
        result['target'] = float(target)
        result['difference'] = mean-float(target)
        result['difference_in_standard_errors'] = (mean-float(target))/se if se else None
    return result


def sample(n, coupling, count, seed, diagonal_only=False):
    """Direct samples of exp[-Tr(M^2)/(2v)], v=lambda/(4N).

    A_ij has independent real/imaginary N(0,1) entries. Symmetrization
    gives diagonal variance v and variance v/2 for each off-diagonal part.
    Returning every scalar result allows honest paired-error estimates.
    """
    lam = float(coupling)
    v = lam/(4*n)
    rng = np.random.Generator(np.random.PCG64(seed))
    rows = []
    for start in range(0, count, 2048):
        size = min(2048, count-start)
        a = rng.standard_normal((size,n,n)) + 1j*rng.standard_normal((size,n,n))
        matrices = (a+a.conj().swapaxes(1,2))*(math.sqrt(v)/2)
        if diagonal_only:
            matrices *= np.eye(n)[None,:,:]
        centers = np.trace(matrices,axis1=1,axis2=2).real/n
        eigenvalues = np.linalg.eigvalsh(matrices)
        # Subtracting c*I shifts every eigenvalue by -c exactly.
        u = np.exp(eigenvalues).mean(axis=1)
        su = np.exp(eigenvalues-centers[:,None]).mean(axis=1)
        moment2 = np.square(eigenvalues).mean(axis=1)
        rows.append(np.column_stack([u,su,centers,centers**2,moment2]))
    result = np.concatenate(rows)
    if not np.all(np.isfinite(result)):
        raise ValueError('Nonfinite sample; reduce the coupling or use a different numerical method.')
    return result


def report(n, coupling, count, seed, digits, diagonal_only=False):
    mp.mp.dps = digits
    checks = analytic_checks()
    lam = mp.mpf(coupling)
    full_u = exact_loop(n,lam)
    full_su = exact_loop(n,lam,True)
    factor = mp.exp(lam/(8*n*n))
    w0 = planar_loop(lam)
    # Repeat the finite sum with extra precision; this is a stability probe.
    with mp.workdps(digits+20):
        refined = exact_loop(n,lam)
    stability = abs(refined-full_u)/max(1,abs(refined))
    data = sample(n,lam,count,seed,diagonal_only)
    u,su,c,c2,m2 = data.T
    target_u = mp.exp(lam/(8*n)) if diagonal_only else full_u
    target_su = mp.exp(lam*(n-1)/(8*n*n)) if diagonal_only else full_su
    prefixes = []
    for size in sorted(set([max(2,count//4),count])):
        prefixes.append({'samples':size,'U':statistics(u[:size],target_u),
                         'SU':statistics(su[:size],target_su)})
    return {
        'environment':{'python':platform.python_version(),'numpy':np.__version__,'mpmath':mp.__version__},
        'inputs':{'N':n,'lambda':str(coupling),'samples':count,'seed':seed,
                  'rng':'PCG64','batch_size':2048,'decimal_digits':digits,
                  'measure':'diagonal-only negative control' if diagonal_only else 'full Hermitian Gaussian'},
        'analytic_checks':checks,
        'full_matrix_targets':{'U':mp.nstr(full_u,25),'SU':mp.nstr(full_su,25),
                               'central_factor_U_over_SU':mp.nstr(factor,25),
                               'planar':mp.nstr(w0,25),
                               'U_over_planar_minus_one':mp.nstr(full_u/w0-1,18),
                               'SU_over_planar_minus_one':mp.nstr(full_su/w0-1,18),
                               'extra_precision_scaled_change':mp.nstr(stability,8)},
        'sampling':{'U':statistics(u,target_u),'SU':statistics(su,target_su),
                    'paired_U_minus_factor_SU':statistics(u-float(factor)*su,0),
                    'center':statistics(c,0),
                    'center_squared':statistics(c2,lam/(4*n*n)),
                    'normalized_trace_M_squared':statistics(m2,lam/(4*n) if diagonal_only else lam/4)},
        'nested_prefixes':prefixes,
        'interpretation':'Standard errors estimate sampling noise for independent matrices. Prefix errors need not decrease monotonically. Exact/planar differences are approximation effects. No statistical test proves localization.'
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--n',type=int,default=4,help='color number, default4; SU rank N-1')
    parser.add_argument('--coupling',default='4',help='nonnegative dimensionless lambda=g^2 N')
    parser.add_argument('--samples',type=int,default=65536,help='independent matrices, default65536')
    parser.add_argument('--seed',type=int,default=20261020)
    parser.add_argument('--dps',type=int,default=60,help='precision of exact/quadrature checks only')
    parser.add_argument('--diagonal-only',action='store_true',help='deliberately wrong full-matrix measure; labeled negative control')
    args = parser.parse_args()
    try:lam = mp.mpf(args.coupling)
    except ValueError:parser.error('--coupling must be a real number')
    if not mp.isfinite(lam) or lam<0:parser.error('--coupling must be finite and nonnegative')
    if not 1<=args.n<=64:parser.error('--n must be between1 and64')
    if not 2<=args.samples<=2000000:parser.error('--samples must be between2 and2000000')
    if not 30<=args.dps<=150:parser.error('--dps must be between30 and150')
    if args.seed<0:parser.error('--seed must be nonnegative')
    print(json.dumps(report(args.n,args.coupling,args.samples,args.seed,args.dps,args.diagonal_only),indent=2,allow_nan=False))


if __name__=='__main__':
    main()
