#!/usr/bin/env python3
"""Write a first-stage master-integral list with a descriptive header.
usage: build_masters.py RAW_MASTERS KIRA_LOG OUT --family NAME --paper NAME --ref NAME --tops N --seeds TEXT
       [--classified JSON]   (PX2 only: adds parity and the uncut linearized propagators per master)
The integral lines are copied from the Kira 'masters' file; the master count and the symmetry lines are
read from the Kira log of the same run and quoted."""
import sys, re, json, argparse
from common import sha256_file
ap = argparse.ArgumentParser()
ap.add_argument('raw'); ap.add_argument('log'); ap.add_argument('out')
ap.add_argument('--family'); ap.add_argument('--paper'); ap.add_argument('--ref'); ap.add_argument('--tops')
ap.add_argument('--seeds'); ap.add_argument('--classified', default=None)
a = ap.parse_args()
raw = [l.rstrip('\n') for l in open(a.raw) if l.strip()]
log = open(a.log).read()
nrel = re.search(r'Sector relations: (\d+)', log).group(1)
nsym = re.search(r'Sector symmetries: (\d+)', log).group(1)
nmas = re.search(r'Number of master integrals: (\d+)', log).group(1)
assert int(nmas) == len(raw), (nmas, len(raw))
cls = None
if a.classified:
    cls = {tuple(e['indices']): e for e in json.load(open(a.classified))}
hdr = [f"# First-stage master integrals of the dissipative 5PM-2SF family {a.paper} (Kira family {a.family}; called {a.ref} in arXiv:2601.16256).",
       f"# Reduction: Kira, integration-by-parts identities seeded with {a.seeds}, on the {a.tops} leading top sectors of the family,",
       f"# with Kira's automatic sector relations and sector symmetries included (run log: 'Sector relations: {nrel}', 'Sector symmetries: {nsym}').",
       f"# Result: {nmas} master integrals ('Number of master integrals: {nmas}' in the run log).  This is a count at the stated seeding, not a minimal basis.",
       f"# One integral per line: {a.family}[n_1,...,n_22] with n_i the exponent of propagator D_i (order of kira/integralfamilies.yaml and of",
       f"# families_dissipative.json; D_1..D_4 are cut), followed by '#' and the Kira sector number (sum of 2^(i-1) over the D_i with n_i > 0)."]
if cls:
    hdr.append("# For PX2 two further columns give the parity under u_i -> -u_i (even/odd) and the uncut linearized worldline propagators")
    hdr.append("# D_5..D_8 present in the integral's sector.")
    ne = sum(1 for e in cls.values() if e['parity'] == 'even'); no = sum(1 for e in cls.values() if e['parity'] == 'odd')
    hdr.append(f"# Parity split: {ne} even + {no} odd.")
out = list(hdr)
for l in raw:
    m = re.match(r'^(\S+)\[([-\d,]+)\]\s*#\s*(\d+)\s*$', l)
    assert m, l
    fam, idx, sec = m.groups()
    v = tuple(int(x) for x in idx.split(','))
    assert len(v) == 22
    assert sum(2**i for i, n in enumerate(v) if n > 0) == int(sec), l
    line = f"{fam}[{idx}]  # {sec}"
    if cls:
        e = cls[v]
        wl = ','.join(f'D{s}' for s in e['uncut_wl_slots']) or '-'
        line += f"  {e['parity']:4s}  uncut linearized: {wl}"
    out.append(line)
open(a.out, 'w').write('\n'.join(out) + '\n')
print(f"wrote {a.out.split('/')[-1]}: {nmas} masters (raw sha256 {sha256_file(a.raw)[:16]}, log sha256 {sha256_file(a.log)[:16]})")
