#!/usr/bin/env python3
"""Assemble boundary_constants_PX2.json: definitions of the 31 master integrals that support the ten
flow-free directions of the PX2 top sectors S_1, S_2, S_3, the values of the twelve that were computed
directly (reference point gamma = 5/4, q^2 = -1), and the ten flow-free directions (left null vectors of the
response matrices) with their weights.  All numbers are copied from the computation records named on the
command line; their sha256 are recorded under "source_sha256" with descriptive names.
usage: build_boundary_constants.py --targets .. --families .. --s1 .. --bank .. --r5 .. --r4 .. --r6 .. --r7 ..
       --r5closure .. --r67closure .. --r4closure .. --s3 .. --s3corr .. --s3pslq .. --kv .. --kb .. --fpa2 .. --out .."""
import sys, json, re, argparse, math
import mpmath as mp
from common import sha256_file, dump

ap = argparse.ArgumentParser()
for k in ['targets','families','s1','bank','r5','r4','r6','r7','r5closure','r67closure','r4closure','s3','s3corr','s3pslq','kv','kb','fpa2','out','nullout']:
    ap.add_argument('--'+k, required=True)
a = ap.parse_args()
mp.mp.dps = 60
src = {}
def load(tag, name, path):
    src[tag] = {'description': name, 'sha256': sha256_file(path)}
    return json.load(open(path)) if path.endswith('.json') else open(path).read()

fam = load('families', 'families_dissipative.json (this package)', a.families)
props = fam['families']['PX2']['propagators']
tg = load('targets', 'table of the 42 target integrals of the three top sectors with their Kira index vectors', a.targets)
S1 = load('S1_values', 'direct evaluation of the S_1 corner integrals I_0, I_2, I_3 (two-frequency Bessel-K cores, closed order-0 forms, numerical order-1 moments)', a.s1)
BANK = load('S2_values', 'assembled S_2 boundary values I_r0..I_r7 with exact forms and routing structure', a.bank)
R5 = load('I_r5', 'static-slice evaluation of I_r5 (Weber-Schafheitlin and Bessel-K Mellin closed forms)', a.r5)
R4 = load('I_r4', 'evaluation of the dotted corner integral I_r4', a.r4)
R6 = load('I_r6', 'static-slice evaluation of I_r6', a.r6)
R7 = load('I_r7', 'static-slice evaluation of I_r7', a.r7)
C5 = load('I_r5_closed', 'integer-relation identification of the I_r5 coefficients (two precisions, re-verified below 1e-150)', a.r5closure)
C67 = load('I_r67_closed', 'integer-relation identification of the I_r6, I_r7 coefficients', a.r67closure)
C4 = load('I_r4_closed', 'integer-relation identification of the I_r4 coefficients', a.r4closure)
S3 = load('S3_values', 'S_3 boundary values (sigma images of the S_1 corner integrals; static-slice closed forms)', a.s3)
S3C = load('S3_static_corrected', 'independent closed-form evaluation of the S_3 static-slice integrals Itilde_5,6,7', a.s3corr)
S3P = load('S3_static_pslq', 'integer-relation identification of the Itilde_5,6,7 coefficients (re-verified below 1e-200)', a.s3pslq)
KV = load('S2_nullvectors', 'left null vectors of the S_2 response matrix at twelve values of epsilon (120-digit extraction)', a.kv)
KB = load('S3_nullvectors', 'left null vectors of the S_3 response matrix at twelve values of epsilon (two extractions at 60 and 120 digits)', a.kb)
FPA2 = load('S2_response_matrix', 'response matrix A(epsilon) of S_2 (source of the epsilon grid, common to S_1, S_2, S_3)', a.fpa2)

# ---- the 31 support rows -------------------------------------------------------------------------
rows = []
for line in tg.split('\n'):
    m = re.match(r'^\|\s*\d+\s*\|\s*(S\d) \(top (\d+)\)\s*\|\s*(\d+)\s*\|\s*(.*?)\s*\|\s*`([^`]+)`\s*\|\s*(\d+)\s*\|\s*(\d+)\s*\|\s*(\d+)\s*\|', line)
    if m:
        S, top, row, pname, kira, sec, dots, nums = m.groups()
        rows.append(dict(S=S, top=int(top), row=int(row), pname=pname, kira=kira, sector=int(sec), dots=int(dots), nums=int(nums)))
assert len(rows) == 42, len(rows)
SUPPORT = {'S1': [0, 2, 3], 'S2': [0,1,2,3,4,5,6,7,9,10,18,19], 'S3': [0,1,2,3,4,5,6,7,9,10,19,20,23,24,31,32]}
assert KV['cert_support'] == SUPPORT['S2'] and KB['support'] == SUPPORT['S3']
def name_of(S, row):
    return {'S1': {0: 'I_0', 2: 'I_2', 3: 'I_3'}}.get(S, {}).get(row) or (f'I_r{row}' if S == 'S2' else f'Itilde_{row}')
def idxvec(kira):
    return [int(x) for x in kira.split('[')[1].rstrip(']').split(',')]
def describe(v):
    den = [f'D{i+1}' for i, n in enumerate(v) if n > 0 and i >= 4]
    lin = [f'D{i+1}' for i in range(4, 8) if v[i] > 0]
    dot = [f'D{i+1}' for i, n in enumerate(v) if n > 1]
    num = [f'D{i+1}' + (f'^{-n}' if n < -1 else '') for i, n in enumerate(v) if n < 0]
    return dict(uncut_denominators=den, uncut_linearized_worldline_propagators=lin, dotted=dot, numerators=num)
CLASS = {}
for S, r, c in [('S1',0,'radiative corner'),('S1',2,'radiative corner, numerator'),('S1',3,'radiative corner, numerator'),
                ('S2',0,'radiative corner (relabeling of I_0)'),('S2',1,'odd numerator, no linearized denominator: vanishes'),
                ('S2',2,'radiative corner, numerator (relabeling of I_3)'),('S2',3,'radiative corner, numerator (relabeling of I_2)'),
                ('S2',4,'radiative corner, dotted graviton propagator'),('S2',5,'static slice (one uncut linearized propagator)'),
                ('S2',6,'static slice, numerator'),('S2',7,'static slice, numerator'),
                ('S2',9,'one uncut linearized propagator D7, odd numerator'),('S2',10,'one uncut linearized propagator D7, numerator'),
                ('S2',18,'two uncut linearized propagators D5, D7'),('S2',19,'two uncut linearized propagators D5, D7, numerator'),
                ('S3',0,'radiative corner (image of I_0)'),('S3',1,'odd numerator, no linearized denominator: vanishes'),
                ('S3',2,'radiative corner, numerator (image of I_3)'),('S3',3,'radiative corner, numerator (image of I_2)'),
                ('S3',4,'radiative corner, dotted graviton propagator (image of I_r4)'),('S3',5,'static slice (image of I_r5)'),
                ('S3',6,'static slice, numerator (image of I_r6)'),('S3',7,'static slice, numerator (image of I_r7)'),
                ('S3',9,'one uncut linearized propagator D7, odd numerator'),('S3',10,'one uncut linearized propagator D7, numerator'),
                ('S3',19,'one uncut linearized propagator D8, odd numerator'),('S3',20,'one uncut linearized propagator D8, numerator'),
                ('S3',23,'two uncut linearized propagators D5, D7'),('S3',24,'two uncut linearized propagators D5, D7, numerator'),
                ('S3',31,'two uncut linearized propagators D5, D8'),('S3',32,'two uncut linearized propagators D5, D8, numerator')]:
    CLASS[(S, r)] = c

def flo(x):  # floor of an 'honest digits' float
    return int(math.floor(float(x)))
def laur(d, digits=None):
    o = {}
    for k in sorted(d, key=lambda s: int(s)):
        e = d[k]; oo = {}
        for f in ('re', 'im', 'exact_re', 'exact_im', 'exact'):
            if f in e: oo[f] = e[f]
        if digits and k in digits: oo['correct_digits'] = flo(digits[k])
        o[k] = oo
    return o

masters = []
computed_names = []
for rr in rows:
    S, row = rr['S'], rr['row']
    if row not in SUPPORT[S]: continue
    v = idxvec(rr['kira'])
    ent = dict(name=name_of(S, row), sector=S, top_sector_label=rr['top'], response_matrix_row=row,
               kira_integral=rr['kira'], index_vector=v, kira_sector=rr['sector'], **describe(v), **{'class': CLASS[(S,row)]})
    nm = ent['name']
    # ---- values ----
    if S == 'S1':
        key = rr['kira'].replace('pm5_2SF_PX2_diss[', 'pm5_2SF_PX2_diss|').rstrip(']').replace(',', '|')
        vals = {}
        for rt in ('fey', 'conj', 'pv'):
            e = S1['routings'][rt][key]
            vals[rt] = laur(e['laurent'], e['honest_digits'])
        ent['status'] = 'computed'
        ent['values'] = {'graviton_i0': vals,
                         'note': 'orders eps^-1 (exact and numerical) and eps^0 (numerical); correct_digits = number of leading significant digits established by two independent evaluations'}
        computed_names.append(nm)
    elif S == 'S2' and row in (0, 2, 3):
        ent['status'] = 'computed (identical to an S_1 integral by the relabeling k1 <-> k2 of loop momenta)'
        ent['equal_to'] = {0: 'I_0', 2: 'I_3', 3: 'I_2'}[row]
    elif (S, row) in (('S2', 1), ('S3', 1)):
        ent['status'] = 'computed'
        ent['values'] = {'exact': '0', 'note': 'vanishes identically, all orders in eps, any i0 assignment (reflection of the two uncut worldline frequency components)'}
        computed_names.append(nm)
    elif S == 'S2' and row == 4:
        m4 = BANK['masters']['4']
        ent['status'] = 'computed'
        ent['values'] = {'graviton_i0': {'fey': laur(R4['laurent_fey']),
                                         'conj': 'complex conjugate of fey (only the radiation-cone layer is complex)',
                                         'pv': 'real part of fey'},
                         'exact': {'-1': 'pi^4', '0': 'pi^4*(4 + (20/3)*ln2 - 4*gammaE) - (2/3)*i*pi^5'},
                         'exact_identification': {'-1': f"integer relation over (pi^4, pi^4 ln2, pi^5), re-verified to {C4['cm1']['reverify_resid']}",
                                                  '0': f"integer relation at 40 digits (re parts over pi^4(1, ln2, gammaE): {C4['c0_re']['coeffs']}; im part over pi^5: {C4['c0_im']['coeffs'][0]}), residuals {C4['c0_re']['reverify_resid_40d']}, {C4['c0_im']['reverify_resid_40d']}"},
                         'correct_digits': {'-1': 'exact', '0': 38},
                         'note': 'order eps^1 not computed'}
        computed_names.append(nm)
    elif S == 'S2' and row in (5, 6, 7):
        srcv = {5: R5, 6: R6, 7: R7}[row]
        cl = {5: C5, 6: C67['r6'], 7: C67['r7']}[row]
        ent['status'] = 'computed'
        ent['values'] = {'linearized_i0': 'value at sigma = +1 for the uncut linearized propagator 1/(D5 + i sigma 0); sigma = -1 reverses the sign; the principal-value part is zero',
                         'graviton_i0': 'independent of the graviton i0 assignment (static slice, radiation cone not crossed)',
                         'exact': {'-1': '0', '0': cl['c0']['exact_im'].replace('*pi^5', '*i*pi^5'), '1': cl['c1']['exact_im'].replace('*pi^5', '*i*pi^5').replace('log(2)', 'ln2')},
                         'exact_identification': {'0': f"integer relation, re-verified to {cl['c0']['reverify_resid']}", '1': f"integer relation over i pi^5 (1, ln2, gammaE), re-verified to {cl['c1']['reverify_resid']}"},
                         'numerical_sigma_plus': laur({k: v for k, v in srcv['laurent_fey'].items() if k in ('0', '1', '2')}),
                         'note': 'purely imaginary; the listed real parts are numerical zeros of the evaluation; order eps^2 numerical only'}
        computed_names.append(nm)
    elif S == 'S3' and row in (0, 2, 3):
        ent['status'] = 'computed (identical to an S_1 integral by a relabeling of loop momenta)'
        ent['equal_to'] = {0: 'I_0', 2: 'I_3', 3: 'I_2'}[row]
    elif S == 'S3' and row == 4:
        ent['status'] = 'not computed directly; the reduction identities in boundary_ibp_identities_PX2.m equate it to I_r4'
    elif S == 'S3' and row in (5, 6, 7):
        rk = f'row{row}'
        e0 = S3C['exact_eps0_sigL_plus'][rk]; e1 = S3C['exact_eps1_sigL_plus'][rk]
        pr = S3P['results'][rk]
        ent['status'] = 'computed'
        ent['values'] = {'linearized_i0': 'value at sigma = +1 for the uncut linearized propagator 1/(D5 + i sigma 0); sigma = -1 reverses the sign; the principal-value part is zero',
                         'graviton_i0': 'independent of the graviton i0 assignment',
                         'exact': {'-1': '0', '0': e0.replace('I*', 'i*').replace('**', '^'), '1': e1.replace('I*', 'i*').replace('**', '^').replace('log(2)', 'ln2').replace('EulerGamma', 'gammaE')},
                         'exact_identification': {'0': f"integer relation, re-verified (relative residual {pr['eps0_reverify']['relresid']})", '1': f"integer relation over i pi^5 (1, ln2, gammaE), re-verified (relative residual {pr['eps1_reverify']['relresid']})"},
                         'equal_to': f"I_r{row} (same values through order eps; independent evaluation in the S_3 variables)"}
        computed_names.append(nm)
    else:
        ent['status'] = 'not computed'
    masters.append(ent)
assert len(masters) == 31, len(masters)
distinct = 31 - 6  # I_r0,I_r2,I_r3, Itilde_0,2,3 are relabelings of the S_1 triple
# ---- relations (left null vectors) ------------------------------------------------------------
def cstr(s):
    """split a Python-style complex string 'a+bj' / '(a+-bj)' / '1.0' into decimal strings (re, im)."""
    s = s.replace(' ', '').strip('()')
    if not s.endswith('j'):
        return s, '0'
    m = re.match(r'^([+-]?[0-9.]+(?:e[+-]?[0-9]+)?)\+?(-?[0-9.]+(?:e[+-]?[0-9]+)?)j$', s)
    assert m, s
    return m.group(1), m.group(2)
eps2 = []
for e in FPA2['eps']:
    m = re.match(r'^\(\[([0-9.e+-]+) \+/- ([0-9.e+-]+)\]', e)
    eps2.append(mp.nstr(mp.mpf(m.group(1)), 40, strip_zeros=False))
rel = {'note': ('Left null vectors v of the response matrices A(eps) of Eq. (27) (v^T A = 0), i.e. the flow-free directions of Table 3, '
                'normalized to 1 on one support row.  The relation among master integrals is sum_row weight[row] * M[row] = 0 at each eps.  '
                'Rows are response-matrix rows (field response_matrix_row of the masters above).  Weights that are rational and '
                'independent of eps are given exactly; the others are given at the twelve eps values of the computation.'),
       'S1': [{'support': ['I_0', 'I_2', 'I_3'], 'rows': [0, 2, 3], 'weights_exact': ['1', '16/3', '-6'], 'eps_dependence': 'none (checked to 4e-84 over the twelve eps values)',
               'relation': 'I_0 + (16/3) I_2 - 6 I_3 = 0, Eq. (36) of the paper'}],
       'S2': [], 'S3': []}
# S2
v0 = BANK['flow_free_directions']['v0']
rel['S2'].append({'support': ['I_r0', 'I_r2', 'I_r3'], 'rows': [0, 2, 3], 'weights_exact': [v0['weights']['0'], v0['weights']['2'], v0['weights']['3']],
                  'eps_dependence': 'none (spread 2.4e-83 over the twelve eps values)', 'relation': '-(1/6) I_r0 + I_r2 - (8/9) I_r3 = 0 (the same relation as in S_1 after relabeling)'})
s2dirs = [(1, [1, 5, 6], 6, '(-7/54, 8/27, 1)'), (2, [1, 5, 7], 7, '(-2/27, 2/27, 1)'), (3, [0, 3, 4, 9, 10, 18, 19], 19, None)]
for vi, sup, piv, lim in s2dirs:
    per = []
    for k, pe in enumerate(KV['per_eps']):
        vec = pe['vectors'][vi]
        w = {}
        for rrow in sup:
            re_, im_ = cstr(vec[str(rrow)])
            w[name_of('S2', rrow)] = {'re': re_, 'im': im_}
        per.append({'eps': eps2[k], 'weights': w, 'max_off_support_weight': f"{float(pe['worst_off_support']):.2e}", 'residual_vT_A': f"{float(pe['worst_residual']):.2e}"})
    ent = {'support': [name_of('S2', x) for x in sup], 'rows': sup, 'normalized_row': piv, 'per_eps': per,
           'weight_digits': 'extracted at 120-digit working precision from response-matrix entries with radii below 1e-107; at least 70 digits'}
    if lim: ent['weights_at_eps_0'] = lim + ' (rational values reached in the limit eps -> 0; Table 5 of the paper)'
    rel['S2'].append(ent)
# S3
DIRSUP = KB['per_eps'][0]['dir_support']
assert all(pe['dir_support'] == DIRSUP for pe in KB['per_eps'])
for di, sup in enumerate(DIRSUP):
    per = []
    for pe in KB['per_eps']:
        w = {}
        for ri, rrow in enumerate(KB['support']):
            if rrow in sup:
                c = pe['rref'][ri][di]
                w[name_of('S3', rrow)] = {'re': c['re'], 'im': c['im']}
        k3 = KB['per_eps'].index(pe)
        assert abs(mp.mpf(pe['eps'])/mp.mpf(eps2[k3]) - 1) < mp.mpf('1e-14')
        per.append({'eps': eps2[k3], 'weights': w, 'two_extraction_agreement_digits': int(math.floor(pe['ker_2leg_digits'])),
                    'max_off_support_weight': f"{float(pe['off_support_max']):.2e}", 'residual_vT_A': f"{float(pe['legB_max_resid']):.2e}"})
    ent = {'support': [name_of('S3', x) for x in sup], 'rows': sup, 'normalized_row': KB['pivot_rows'][di], 'per_eps': per,
           'weight_digits': 'two extractions at 60 and 120 digits agree to the number of digits listed per eps (at least 45)'}
    if sup == [1, 5, 7]: ent['weights_at_eps_0'] = '(-2/27, 2/27, 1) (Table 5 of the paper)'
    if sup == [1, 5, 6]: ent['weights_at_eps_0'] = '(-7/54, 8/27, 1) (Table 5 of the paper)'
    rel['S3'].append(ent)
# check: the S_3 image of Eq. (36), (1, -6, 16/3) on (Itilde_0, Itilde_2, Itilde_3), lies in the span of the two directions normalized on rows 0 and 2
pe = KB['per_eps'][0]
def wcol(di):
    return {rrow: mp.mpc(mp.mpf(pe['rref'][ri][di]['re']), mp.mpf(pe['rref'][ri][di]['im'])) for ri, rrow in enumerate(KB['support'])}
d0 = KB['pivot_rows'].index(0); d2 = KB['pivot_rows'].index(2)
u0, u2 = wcol(d0), wcol(d2)
comb = {rrow: u0[rrow] - 6*u2[rrow] for rrow in KB['support']}
resid = max(abs(comb[rrow]) for rrow in KB['support'] if rrow not in (0, 2, 3))
chk37 = abs(comb[3] - mp.mpf(16)/3)
rel['S3_check'] = (f"the combination (direction normalized on row 0) - 6 (direction normalized on row 2) has weight {mp.nstr(comb[3].real, 20)} on Itilde_3 "
                   f"(16/3 to {mp.nstr(chk37, 2)}) and weights below {mp.nstr(resid, 2)} on all other rows: the relation Itilde_0 - 6 Itilde_2 + (16/3) Itilde_3 = 0, "
                   f"the S_3 image of Eq. (36), lies in the span")
out = {
 'title': 'Boundary master integrals of the dissipative crossed family PX2 at 5PM-2SF: definitions, computed values, and the flow-free relations',
 'paper': 'M. D. Schwartz, Memory back-reaction in black-hole scattering at fifth post-Minkowskian order: Sec. 5 (response matrices, boundary constants of PX2, relations among boundary constants), Tables 3, 4 and 5; equation numbers below as in the revised manuscript that this package accompanies',
 'conventions': {
   'family': 'pm5_2SF_PX2_diss (kira/integralfamilies.yaml, families_dissipative.json)',
   'propagators_D1_to_D22': props,
   'cut': 'D1..D4 are cut: 2 pi delta(2 k_i.u) each',
   'measure': 'prod_i d^D k_i / (i pi^(D/2)), D = 4 - 2 eps, no exp(gammaE eps) factor',
   'kinematic_point': {'gamma': '5/4', 'q^2': '-1', 'u1^2': '1', 'u2^2': '1', 'u_i.q': '0'},
   'index_vector': 'n_i = exponent of D_i; negative n_i = numerator power; kira_sector = sum of 2^(i-1) over D_i with n_i > 0',
   'sectors': {'S1': 1077759, 'S2': 1091071, 'S3': 599551},
   'graviton_i0': "keys fey / conj / pv: Feynman (+i0) on every graviton propagator / the complex-conjugate assignment / their average.  All values in this file are boundary data with Feynman-type graviton propagators; no in-in (retarded) boundary values are given.",
   'linearized_i0': 'for integrals with an uncut linearized worldline propagator 1/(D_j + i sigma 0) the value is given at sigma = +1; which sign of sigma corresponds to retarded boundary conditions is not fixed here',
   'laurent': "keys '-2','-1','0','1','2' = coefficient of eps^k; 're','im' decimal strings as produced (digits beyond 'correct_digits' are not significant); 'exact' closed forms use pi, ln2 = log 2, gammaE = Euler's constant, i",
 },
 'counts': {'support_rows': 31, 'distinct_integrals_at_most': distinct,
            'computed_distinct': sorted(set(computed_names)), 'n_computed_distinct': len(set(computed_names)),
            'identically_zero': ['I_r1', 'Itilde_1'],
            'note': 'I_r0, I_r2, I_r3 and Itilde_0, Itilde_2, Itilde_3 are relabelings of I_0, I_3, I_2; Itilde_5,6,7 equal I_r5,6,7 in value'},
 'masters': masters,
 'flow_free_relations': rel,
 'source_sha256': src,
}
dump(out, a.out)
# the same null vectors, one file per sector, next to the response matrices
import os
for S in ('S1', 'S2', 'S3'):
    dump({'sector': S, 'note': rel['note'], 'eps_grid': eps2 if S != 'S1' else 'weights independent of eps',
          'left_null_vectors': rel[S], **({'check': rel['S3_check']} if S == 'S3' else {}),
          'source_sha256': {k: v for k, v in src.items() if k in ({'S1': ['targets'], 'S2': ['S2_nullvectors', 'S2_response_matrix', 'S2_values'], 'S3': ['S3_nullvectors', 'S2_response_matrix']}[S])}},
         os.path.join(a.nullout, f'null_vectors_{S}.json'))
print(f"wrote {a.out.split('/')[-1]}: {len(masters)} support rows, {len(set(computed_names))} computed distinct "
      f"({sorted(set(computed_names))}); relations S1 {len(rel['S1'])} S2 {len(rel['S2'])} S3 {len(rel['S3'])}; S3 Eq.(36)-image check: 16/3 to {mp.nstr(chk37,2)}, off-support {mp.nstr(resid,2)}")