"""Undeformed round-sphere ABJM, compact U(N)_k x U(N)_(-k), k>0.

Default: the regulated direct N=1 integral (AB13 preparation).
--spectral: kernel traces and N=2 (requires the AB18 spectral lesson).
Python 3.13.7, mpmath 1.3.0. No repository imports or network calls.
"""
import argparse
from functools import lru_cache
import json
import mpmath as mp


def z1_regulated(k, epsilon):
    C=mp.pi**2*k*k+epsilon**2
    width=mp.sqrt(epsilon/C)
    average=2*mp.quad(lambda u:mp.exp(-u*u/2)/(4*mp.cosh(mp.pi*width*u)**2),
                      [0,1,4,mp.inf])/mp.sqrt(2*mp.pi)
    return mp.pi*average/mp.sqrt(C)


def z2_positive(k, cutoff=mp.inf):
    integrand=lambda d: d*mp.tanh(d/(2*k))**2/mp.sinh(d/2) if d else mp.mpf(0)
    points=[0,1,4,cutoff] if cutoff>4 else [0,cutoff]
    return mp.quad(integrand,points)/(32*mp.pi**2*k*k)


def z2_tail(k, cutoff):
    # tanh^2 <= 1 and csch(d/2) <= 2 exp(-d/2)/(1-exp(-cutoff)).
    return (2*cutoff+4)*mp.exp(-cutoff/2)/(16*mp.pi**2*k*k*(1-mp.exp(-cutoff)))


def center_convolution(d, L):
    if d==0:return mp.tanh(L/2)
    a=d/2
    return 2*mp.atanh(mp.tanh(a/2)*mp.tanh((L-a)/2))/mp.sinh(a)


def z2_box(k,L):
    f=lambda d:center_convolution(d,L)*mp.tanh(d/(2*k))**2
    return mp.quad(f,[0,1,L,2*L])/(16*mp.pi**2*k*k)


def diagonal_tail(k,L):
    return mp.atan(mp.exp(-L/2))/(mp.pi*k)


@lru_cache(maxsize=None)
def legendre(n,dps):
    with mp.workdps(dps):
        x,w=mp.gauss_quadrature(n,'legendre')
        return tuple(x),tuple(w)


def kernel_traces(k,L,n):
    nodes,weights=legendre(n,mp.mp.dps)
    x=[L*a for a in nodes]
    # Diagonal of the symmetric weighted kernel sqrt(w_i w_j) rho(x_i,x_j).
    diagonal=[L*w/(8*mp.pi*k*mp.cosh(a/2)) for a,w in zip(x,weights)]
    t1=mp.fsum(diagonal)
    t2=mp.fsum(a*a for a in diagonal)
    pair_sum=mp.mpf(0)
    for i in range(n):
        for j in range(i):
            arg=(x[i]-x[j])/(2*k)
            product=diagonal[i]*diagonal[j]
            t2+=2*product/mp.cosh(arg)**2
            pair_sum+=product*mp.tanh(arg)**2
    trace_difference=(t1*t1-t2)/2
    return t1,t2,pair_sum,abs(trace_difference-pair_sum)


def number(value):
    return mp.nstr(value,22)


def run(k=1,spectral=False,omit_exchange=False,dps=50):
    mp.mp.dps=dps
    t1=mp.mpf(1)/(4*k)
    epsilon_rows=[]
    for eps in map(mp.mpf,['.1','.03','.01']):
        value=z1_regulated(k,eps)
        bias=1-value/t1
        bound=eps/(k*k)+eps*eps/(2*mp.pi**2*k*k)
        assert 0<=bias<=bound
        epsilon_rows.append({'epsilon':str(eps),'Z1_regulated':number(value),
                             'relative_regulator_bias':number(bias),'bias_upper_bound':number(bound)})
    low=z1_regulated(k,mp.mpf('.01'))
    with mp.workdps(dps+20): high=z1_regulated(k,mp.mpf('.01'))
    precision_error=abs(low-high)
    assert precision_error<mp.power(10,-dps+10)
    output={'theory':'compact U(N)_k x U(N)_(-k), round, undeformed, fixed finite N and k',
            'level':k,'digits':dps,'prescription':'positive Gaussian Abel damping, then epsilon -> 0',
            'direct_N1':{'limit':number(t1),'finite_epsilon':epsilon_rows,
                         'precision_repeat_difference':number(precision_error)},
            'spectral_checks':None}
    if spectral:
        value=z2_positive(k)
        with mp.workdps(dps+20): high=z2_positive(k)
        response_error=abs(value-high)
        assert response_error<mp.power(10,-dps+10)
        analytic={1:1/(16*mp.pi),2:1/(32*mp.pi**2)}.get(k)
        if analytic is not None:
            assert abs(value-analytic)<mp.power(10,-dps+10)
        cutoff_rows=[]
        for D in [16,32,64]:
            truncated=z2_positive(k,D)
            tail=z2_tail(k,D)
            assert 0<=value-truncated<=tail
            cutoff_rows.append({'D':D,'truncated_Z2':number(truncated),
                                'actual_omission':number(value-truncated),'tail_bound':number(tail)})
        boxes=[]
        for L in [16,24,32]:
            box=z2_box(k,L)
            delta=diagonal_tail(k,L)
            bound=t1*delta-delta*delta/2
            assert 0<=value-box<=bound
            for n in [32,64,128]:
                trace1,trace2,pairs,roundoff=kernel_traces(k,L,n)
                boxes.append({'L':L,'nodes':n,'trace1':number(trace1),'trace2':number(trace2),
                              'kernel_Z2':number(pairs),'quadrature_error_vs_same_box':number(abs(pairs-box)),
                              'box_Z2':number(box),'box_omission':number(value-box),
                              'box_omission_bound':number(bound),'trace_vs_pair_roundoff':number(roundoff)})
        wrong=t1*t1/2
        detected=omit_exchange and abs(wrong-value)>mp.mpf('1e-6')
        assert not omit_exchange or detected
        output['spectral_checks']={
            'measure':'L2(R, dx), hbar=2*pi*k; rho is not trace normalized',
            'trace1_exact':number(t1),'trace2':number(t1*t1-2*value),
            'Z2_positive_integral':number(value),'analytic_Z2':None if analytic is None else number(analytic),
            'precision_repeat_difference':number(response_error),'difference_cutoffs':cutoff_rows,
            'finite_box_kernel_checks':boxes,
            'mode':'omitted exchange diagnostic' if omit_exchange else 'fermionic exchange included',
            'reported_Z2':number(wrong if omit_exchange else value),
            'intentional_error_detected':bool(detected)}
    output['checks_pass']=True
    return output


if __name__=='__main__':
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--level',type=int,choices=[1,2,3,4],default=1)
    parser.add_argument('--spectral',action='store_true')
    parser.add_argument('--omit-exchange',action='store_true')
    parser.add_argument('--dps',type=int,choices=range(40,71),default=50)
    args=parser.parse_args()
    if args.omit_exchange and not args.spectral:
        parser.error('--omit-exchange requires --spectral and its AB18 preparation.')
    print(json.dumps(run(args.level,args.spectral,args.omit_exchange,args.dps),indent=2))