"""Finite-rank ABJM gas versus its perturbative Airy answer.

Compact U(N)_k x U(N)_(-k), k=1 or 2; F=-log Z. No repository imports.
The kernel/eigensolver uses float64; --dps changes only mpmath calculations.
The resolution indicator is empirical, not a certified quadrature bound.
"""
import argparse
import json
import math
from pathlib import Path
import numpy as np
import mpmath as mp


def airy_data(k, N):
    C = 2/(mp.pi**2*k)
    B = mp.mpf(k)/24 + mp.mpf(1)/(3*k)
    A = mp.log(2)/4-mp.zeta(3)/(8*mp.pi**2) if k == 1 else -mp.zeta(3)/(2*mp.pi**2)
    lp = A-mp.log(C)/3+mp.log(mp.airyai((N-B)/mp.root(C, 3)))
    return C, B, A, lp


def exact_small(k):
    p = mp.pi
    if k == 1:
        return {1: mp.mpf(1)/4, 2: 1/(16*p), 4: (10-p*p)/(1024*p*p)}
    return {1: mp.mpf(1)/8, 2: 1/(32*p*p),
            4: (24-32*p*p+3*p**4)/(49152*p**4)}


def log_coefficients(eigenvalues, nmax):
    """Each eigenvalue is used at most once: vector RHS uses the old array."""
    logs = np.full(nmax+1, -np.inf)
    logs[0] = 0.
    for value in eigenvalues:
        if value > 0:
            logs[1:] = np.logaddexp(logs[1:], math.log(value)+logs[:-1])
    return logs


def spectral_run(k, L, nodes, nmax):
    x, w = np.polynomial.legendre.leggauss(nodes)
    x, w = L*x, L*w
    a = np.sqrt(w/(2*np.cosh(x/2)))
    matrix = a[:, None]*a[None, :]/(4*np.pi*k*np.cosh((x[:, None]-x[None, :])/(2*k)))
    eigenvalues = np.linalg.eigvalsh(matrix)
    neg = float(-np.sum(eigenvalues[eigenvalues < 0]))
    # A visibly negative eigenvalue is a failure, not something to erase.
    if float(eigenvalues[0]) < -2e-14:
        raise ArithmeticError("Negative eigenvalue exceeds the float64 diagnostic threshold")
    logs = log_coefficients(eigenvalues, nmax)
    return {"L": L, "nodes": nodes, "negative_eigenvalue_mass": neg,
            "minimum_eigenvalue": float(eigenvalues[0]),
            "trace_error": float(abs(mp.mpf(float(np.trace(matrix)))-mp.mpf(1)/(4*k))),
            "log_Z": logs.tolist()}


def relative_tail(logs, N, trace_tail):
    """Exact coefficient bound IF logs describe the exact box compression.

    Here evaluated on numerical coefficients: it does not bound quadrature.
    """
    return mp.fsum(trace_tail**j/mp.factorial(j)*mp.exp(mp.mpf(logs[N-j])-logs[N])
                   for j in range(1, N+1))


def calculate(k=2, nmax=8, coarse=False):
    # A 2x2 grid varies window and node count separately.
    widths, orders = ((48, 64), (128, 192)) if coarse else ((80, 96), (768, 1024))
    runs = [spectral_run(k, L, n, nmax) for L in widths for n in orders]
    best = runs[-1]
    logs = best["log_Z"]
    trace_tail = mp.atan(mp.exp(-mp.mpf(best["L"])/2))/(mp.pi*k)
    exact = exact_small(k)
    rows = []
    for N in range(1, nmax+1):
        C, B, A, lp = airy_data(k, N)
        logz = mp.mpf(logs[N])
        F, Fp = -logz, -lp
        delta = mp.expm1(logz-lp)
        node_change = max(abs(runs[0]["log_Z"][N]-runs[1]["log_Z"][N]),
                          abs(runs[2]["log_Z"][N]-runs[3]["log_Z"][N]))
        window_change = max(abs(runs[0]["log_Z"][N]-runs[2]["log_Z"][N]),
                            abs(runs[1]["log_Z"][N]-runs[3]["log_Z"][N]))
        # A diagnostic sensitivity scale, not a bound on eigensolver error.
        sign_scale = max(r["negative_eigenvalue_mass"] *
                         math.exp(r["log_Z"][N-1]-r["log_Z"][N]) for r in runs)
        floor = 32*N*np.finfo(float).eps
        tail = relative_tail(logs, N, trace_tail)
        indicator = 10*max(node_change, window_change, sign_scale, floor)+float(mp.log1p(tail))
        exact_error = float(abs(mp.expm1(logz-mp.log(exact[N])))) if N in exact else None
        rows.append({"N": N, "F": float(F), "F_over_N32": float(F/N**mp.mpf('1.5')),
                     "F_pert": float(Fp), "F_pert_over_N32": float(Fp/N**mp.mpf('1.5')),
                     "relative_Airy_residual": float(delta), "F_minus_F_pert": float(F-Fp),
                     "node_change_in_F": node_change, "window_change_in_F": window_change,
                     "negative_mass_sensitivity": sign_scale,
                     "box_relative_bound_evaluated": float(tail),
                     "resolution_indicator_in_F": indicator,
                     "resolved_by_indicator": bool(abs(F-Fp) > indicator),
                     "exact_calibration_relative_error": exact_error})
    return {"model": "compact equal-rank U(N)_k x U(N)_(-k), undeformed round sphere",
            "k": k, "nmax": nmax, "fixed_k": True, "coarse_negative_control": coarse,
            "kernel_precision": "float64", "airy_decimal_precision": mp.mp.dps,
            "leading_coefficient": float(mp.pi*mp.sqrt(2*k)/3),
            "A": float(airy_data(k, 1)[2]), "trace_tail_at_final_L": float(trace_tail),
            "indicator_is_certified": False,
            "runs": runs, "rows": rows}


def write_svg(report, path):
    """Static, labeled scientific figure; no plotting package or remote assets."""
    rows = report["rows"]
    out = ['<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 980 610" role="img" aria-labelledby="title desc">',
           '<title id="title">Finite-rank ABJM and the Airy approximation</title>',
           '<desc id="desc">Upper panel: F divided by N to the three-halves, finite kernel points and perturbative Airy curve, with the leading fixed-level constant. Lower panel: signed relative Airy residual plotted by its logarithmic magnitude, alongside an empirical resolution indicator. A plus or minus next to each point gives the sign; open points are unresolved. The indicator is not a certified error bound.</desc>',
           '<rect width="980" height="610" fill="white"/>',
           '<g font-family="system-ui,sans-serif" fill="#172231" font-size="14">',
           f'<text x="80" y="28" font-size="20">ABJM, fixed k = {report["k"]}; F = −log Z</text>']
    nmax = report['nmax']
    X = lambda n: 80 + 820*(n-1)/max(1, nmax-1)
    c = report['leading_coefficient']
    ylow = min([r['F_over_N32'] for r in rows]+[c])-.035
    yhigh = max([r['F_over_N32'] for r in rows]+[c])+.035
    Y = lambda y: 255-180*(y-ylow)/(yhigh-ylow)
    for y in np.linspace(ylow, yhigh, 5):
        out += [f'<path d="M80 {Y(y):.2f} H900" stroke="#dde2e7"/>',
                f'<text x="72" y="{Y(y)+5:.2f}" text-anchor="end">{y:.3f}</text>']
    out += ['<text x="80" y="58">F / N^(3/2)</text>',
            f'<path d="M80 {Y(c):.2f} H900" stroke="#923c00" stroke-dasharray="7 5"/>']
    points = ' '.join(f'{X(r["N"]):.2f},{Y(r["F_pert_over_N32"]):.2f}' for r in rows)
    out += [f'<polyline points="{points}" fill="none" stroke="#0868ac" stroke-width="2"/>']
    for r in rows:
        out += [f'<circle cx="{X(r["N"]):.2f}" cy="{Y(r["F_over_N32"]):.2f}" r="4" fill="#172231"/>']
    out += ['<text x="80" y="282">● finite kernel    blue: Airy (nearly overlapping)    dashed: π√(2k)/3</text>']
    data = [(r, math.log10(max(abs(r['relative_Airy_residual']), 1e-18)),
             math.log10(max(math.expm1(r['resolution_indicator_in_F'])*math.exp(-r['F_minus_F_pert']),1e-18))) for r in rows]
    low = math.floor(min(min(d, e) for _,d,e in data))-1
    high = math.ceil(max(max(d, e) for _,d,e in data))+1
    Q = lambda y: 530-175*(y-low)/(high-low)
    out += ['<text x="80" y="328">log₁₀ magnitude of Z / Z_Airy − 1 (point labels show sign)</text>']
    for y in range(low, high+1, max(1, math.ceil((high-low)/5))):
        out += [f'<path d="M80 {Q(y):.2f} H900" stroke="#dde2e7"/>',
                f'<text x="72" y="{Q(y)+5:.2f}" text-anchor="end">{y}</text>']
    pts = ' '.join(f'{X(r["N"]):.2f},{Q(e):.2f}' for r,_,e in data)
    out += [f'<polyline points="{pts}" fill="none" stroke="#9b2c87" stroke-width="2" stroke-dasharray="3 4"/>']
    for r,d,e in data:
        color = '#0868ac' if r['resolved_by_indicator'] else 'white'
        sign = '+' if r['relative_Airy_residual'] >= 0 else '−'
        out += [f'<circle cx="{X(r["N"]):.2f}" cy="{Q(d):.2f}" r="4" fill="{color}" stroke="#0868ac" stroke-width="2"/>',
                f'<text x="{X(r["N"])+7:.2f}" y="{Q(d)-5:.2f}">{sign}</text>',
                f'<text x="{X(r["N"]):.2f}" y="552" text-anchor="middle">{r["N"]}</text>']
    out += ['<text x="920" y="552">N</text>',
            '<text x="80" y="582">Dotted: empirical resolution indicator, not a certified bound. Open points: unresolved.</text>',
            '</g></svg>']
    Path(path).write_text('\n'.join(out)+'\n', encoding='utf-8')


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--k', type=int, choices=[1, 2], default=2)
    parser.add_argument('--nmax', type=int, default=8)
    parser.add_argument('--dps', type=int, default=60)
    parser.add_argument('--coarse', action='store_true', help='deliberately under-resolved negative control')
    parser.add_argument('--json', type=Path, help='save the complete report')
    parser.add_argument('--svg', type=Path, help='save the two-panel static plot')
    parser.add_argument('--check-expected', type=Path, help='check default physical columns within the saved tolerance')
    args = parser.parse_args()
    if not 1 <= args.nmax <= 12:
        parser.error('This bounded lab supports 1 <= nmax <= 12')
    if not 40 <= args.dps <= 100:
        parser.error('Use 40 <= dps <= 100; the kernel remains float64')
    mp.mp.dps = args.dps
    report = calculate(args.k, args.nmax, args.coarse)
    if args.check_expected:
        expected = json.loads(args.check_expected.read_text(encoding='utf-8'))
        if (args.k, args.nmax, args.coarse) != (2, 8, False):
            parser.error('Expected-output comparison is for the default k=2, nmax=8 run')
        for row, reference in zip(report['rows'], expected['rows'], strict=True):
            for key in ['F', 'F_pert', 'relative_Airy_residual']:
                if abs(row[key]-reference[key]) > expected['absolute_tolerance']:
                    raise ArithmeticError(f'Expected-output discrepancy: N={row["N"]}, {key}')
        print('Expected physical columns agree within the declared absolute tolerance.')
    print('k N F/N^(3/2) relative_Airy_residual indicator_in_F resolved')
    for row in report['rows']:
        print(args.k, row['N'], f'{row["F_over_N32"]:.10f}',
              f'{row["relative_Airy_residual"]:+.7e}',
              f'{row["resolution_indicator_in_F"]:.3e}', row['resolved_by_indicator'])
    print('Resolution indicator is empirical. Kernel arithmetic is float64.')
    if args.json:
        args.json.write_text(json.dumps(report, indent=2)+'\n', encoding='utf-8')
    if args.svg:
        write_svg(report, args.svg)


if __name__ == '__main__':
    main()