"""Flatten a LaTeX working paper and resolve numbering from main.aux so that
pandoc produces a Word file with correct section/table/figure/equation numbers."""
import re, sys, pathlib

PAPER = pathlib.Path(sys.argv[1]).resolve()
OUT = pathlib.Path(sys.argv[2])
warn = lambda *a: print("WARN:", *a, file=sys.stderr)

def strip_comments(txt):
    out = []
    for line in txt.splitlines():
        if re.match(r'^\s*%', line):
            continue
        # drop trailing unescaped % comments
        line = re.sub(r'(?<!\\)%.*$', '', line)
        out.append(line)
    return "\n".join(out) + "\n"

def inline(path):
    txt = strip_comments(path.read_text())
    def rep(m):
        name = m.group(1)
        p = PAPER / (name if name.endswith('.tex') else name + '.tex')
        if not p.exists():
            warn("missing input", p); return m.group(0)
        return inline(p)
    return re.sub(r'\\input\{([^}]+)\}', rep, txt)

def brace_arg(s, i):
    """s[i] == '{' -> return (content, index after closing brace)"""
    assert s[i] == '{', s[i:i+20]
    depth = 0
    for j in range(i, len(s)):
        if s[j] == '{' and (j == 0 or s[j-1] != '\\'): depth += 1
        elif s[j] == '}' and s[j-1] != '\\':
            depth -= 1
            if depth == 0:
                return s[i+1:j], j+1
    raise ValueError("unbalanced braces")

# ---------------------------------------------------------------- aux
aux = (PAPER / 'main.aux').read_text()
labels = {}
for m in re.finditer(r'\\newlabel\{([^}]+)\}\{', aux):
    outer, _ = brace_arg(aux, m.end()-1)
    num, _ = brace_arg(outer, 0)
    labels[m.group(1)] = num.strip()
toc = []
for m in re.finditer(r'\\contentsline \{(section|subsection|subsubsection)\}\{\\numberline \{([^}]*)\}(.*?)\}\{\d+\}\{', aux):
    toc.append((m.group(1), m.group(2), m.group(3)))
print(f"aux: {len(labels)} labels, {len(toc)} numbered headings")

# ---------------------------------------------------------------- flatten
tex = inline(PAPER / 'main.tex')

# titlepage: pandoc cannot evaluate \ifx -> keep the \else branch only
tex = re.sub(r'\\ifx\\WPsubtitle\\empty\\else(.*?)\\fi', r'\1', tex, flags=re.S)

# appendix environment / counters that pandoc does not understand
tex = re.sub(r'\\(begin|end)\{appendices\}', '', tex)
tex = re.sub(r'\\renewcommand\{\\the(table|figure)\}\{[^{}]*(\{[^{}]*\})*[^{}]*\}', '', tex)
tex = re.sub(r'\\setcounter\{(table|figure)\}\{0\}', '', tex)

# ---------------------------------------------------------------- headings
pos = 0; out = []; ti = 0
pat = re.compile(r'\\(section|subsection|subsubsection)\{')
while True:
    m = pat.search(tex, pos)
    if not m: break
    title, end = brace_arg(tex, m.end()-1)
    out.append(tex[pos:m.start()])
    lvl = m.group(1)
    if ti < len(toc) and toc[ti][0] == lvl:
        num = toc[ti][1]; ti += 1
        if re.sub(r'\W', '', toc[ti-1][2])[:12] != re.sub(r'\W', '', title)[:12]:
            warn("heading mismatch:", repr(toc[ti-1][2][:40]), "vs", repr(title[:40]))
        out.append(f"\\{lvl}{{{num}~{title}}}")
    else:
        warn("unnumbered heading kept:", lvl, title[:40])
        out.append(tex[m.start():end])
    pos = end
out.append(tex[pos:]); tex = "".join(out)
if ti != len(toc): warn(f"only {ti}/{len(toc)} toc entries consumed")

# ---------------------------------------------------------------- floats: caption prefix
def fix_float(m):
    kind, body = m.group(1), m.group(2)
    labs = [l for l in re.findall(r'\\label\{([^}]+)\}', body) if l in labels]
    if not labs:
        warn("float without known label", body[:60]); return m.group(0)
    num = labels[labs[0]]
    word = 'Table' if kind == 'table' else 'Figure'
    cm = re.search(r'\\caption(\[[^\]]*\])?\{', body)
    if not cm: return m.group(0)
    cap, end = brace_arg(body, cm.end()-1)
    body = body[:cm.start()] + f"\\caption{{{word} {num}: {cap}}}" + body[end:]
    return f"\\begin{{{kind}}}{body}\\end{{{kind}}}"
tex = re.sub(r'\\begin\{(table|figure)\}(.*?)\\end\{\1\}', fix_float, tex, flags=re.S)

# ---------------------------------------------------------------- equations: number tag
def fix_eq(m):
    body = m.group(1)
    lm = re.search(r'\\label\{([^}]+)\}', body)
    if lm and lm.group(1) in labels:
        body = body.replace(lm.group(0), '')
        body = body.rstrip() + f"\\qquad\\text{{({labels[lm.group(1)]})}}\n"
    return "\\begin{equation}" + body + "\\end{equation}"
tex = re.sub(r'\\begin\{equation\}(.*?)\\end\{equation\}', fix_eq, tex, flags=re.S)

# ---------------------------------------------------------------- cross-references
missing = set()
def ref(m):
    k = m.group(2)
    if k not in labels:
        missing.add(k); return '??'
    return f"({labels[k]})" if m.group(1) == 'eqref' else labels[k]
tex = re.sub(r'\\(eqref|ref)\{([^}]+)\}', ref, tex)
if missing: warn("unresolved labels:", sorted(missing))
tex = re.sub(r'\\label\{[^}]+\}', '', tex)

# booktabs partial rules have no Word equivalent and leak as text
tex = re.sub(r'\\cmidrule(\([^)]*\))?\{[^}]*\}', '', tex)

# ---------------------------------------------------------------- pandoc shims
shims = r"""
\newcommand{\num}[1]{#1}
\renewcommand{\tablelegend}[1]{\par\textit{Notes:} #1}
\renewcommand{\tstat}[1]{(\textit{t}=#1)}
\newcommand{\textperiodcentered}{·}
"""
tex = tex.replace("\\begin{document}", shims + "\\begin{document}", 1)
tex = re.sub(r'(\\newcommand\{\\WPauthor\}\{[^}]*)\}', lambda m: m.group(1).replace('\\quad', '~~~~') + '}', tex)
OUT.write_text(tex)
print("wrote", OUT, len(tex), "chars")