"""Convert HPD-Parsing predictions (JSON) into per-page markdown for OmniDocBench.

Input : JSON, a list of ``{img_path, pred}`` (pred is the ``<BLOCK> <type> [bbox]
        <CHILD> <content>`` stream from ``document parsing with fork.``).
Output: a folder of ``<image_stem>.md`` files matching the OmniDocBench GT paths.

    python hpd_to_markdown.py --input preds.json --out-md pred_md/ \
        --simplify-left-paren --clean-formula-tail --norm-formula-flag --wrap-cjk-arith
"""

import argparse
import json
import os
import re
from pathlib import Path


_TALL = re.compile(
    r'\\d?frac|\\tfrac|\\cfrac|\\binom|\\sqrt'
    r'|\\sum|\\prod|\\coprod|\\int|\\iint|\\iiint|\\oint'
    r'|\\bigcup|\\bigcap|\\bigoplus|\\bigotimes|\\bigsqcup'
    r'|\\begin\{'
    r'|\\overbrace|\\underbrace|\\overset|\\underset|\\stackrel'
    r'|\\substack|\\atop|\\\\'
)


def _scan_delims(s):
    out = []
    for m in re.finditer(r'\\(left|right)\s*', s):
        dm = re.match(r'\\[a-zA-Z]+|\\.|.', s[m.end():])
        if not dm:
            continue
        out.append({'kind': m.group(1), 'delim': dm.group(0),
                    'start': m.start(), 'end': m.end() + dm.end()})
    return out


def simplify_left_right(s: str) -> str:
    """Downgrade `\\left( ... \\right)` with no tall inner structure to plain `( )`."""
    if '\\left' not in s:
        return s
    stack, pairs = [], []
    for d in _scan_delims(s):
        if d['kind'] == 'left':
            stack.append(d)
        elif stack:
            pairs.append((stack.pop(), d))
    edits = []
    for L, R in pairs:
        if L['delim'] == '(' and R['delim'] == ')' and not _TALL.search(s[L['end']:R['start']]):
            edits.append((L['start'], L['end'], '('))
            edits.append((R['start'], R['end'], ')'))
    for st, en, rep in sorted(edits, key=lambda x: x[0], reverse=True):
        s = s[:st] + rep + s[en:]
    return s


_ELLIPSIS = r'(?:\\dots|\\cdots|\\ldots|\\dotsb|\\dotsc)'
_CLOSER = r'(?:\\right\s*[.\}\]\)]|\\end\s*\{(?:array|matrix|cases|bmatrix|pmatrix|vmatrix|smallmatrix)\})'
_TAIL_WRAP = re.compile(r'^(?P<core>.*?)(?P<wrap>\s*(?:\\\]|\\\)|\$\$))?\s*$', re.DOTALL)


def clean_formula_tail(s: str) -> str:
    """Strip degenerate formula tails (repeated/dangling ellipses, stray `\\quad`)."""
    if not s:
        return s
    m = _TAIL_WRAP.match(s)
    core, wrap = m.group('core'), m.group('wrap') or ''
    prev = None
    while prev != core:
        prev = core
        core = re.sub(r'(' + _ELLIPSIS + r')(?:\s*' + _ELLIPSIS + r')+', r'\1', core)
        core = re.sub(r'(?P<keep>' + _CLOSER + r')\s*(?:\\q?quad\s*)*' + _ELLIPSIS + r'\s*$',
                      lambda mm: mm.group('keep'), core)
        core = re.sub(r'(?:\s*\\q?quad)+\s*' + _ELLIPSIS + r'\s*$', '', core)
        core = re.sub(r'(?:\s*\\q?quad)+\s*$', '', core)
        core = core.rstrip()
    return core + wrap


_OP_MAP = {
    '≈': r'\approx', '≠': r'\neq', '≤': r'\leq', '≥': r'\geq', '×': r'\times',
    '÷': r'\div', '±': r'\pm', '∓': r'\mp', '·': r'\cdot', '∙': r'\cdot',
    '⋅': r'\cdot', '∗': '*', '−': '-', '≡': r'\equiv', '∝': r'\propto',
    '∞': r'\infty', '√': r'\sqrt', '→': r'\to', '≪': r'\ll', '≫': r'\gg',
}
_ARITH_ALLOWED = re.compile(r'^[0-9A-Za-z\s=+\-*/^_().,:;<>|%!\u4e00-\u9fff' + ''.join(_OP_MAP.keys()) + r']+$')
_ARITH_HASOP = re.compile(r'[=+\-*/' + ''.join(_OP_MAP.keys()) + r']')
_KNOWN_FUNCS = {'sin', 'cos', 'tan', 'cot', 'sec', 'csc', 'log', 'ln', 'exp',
                'lim', 'max', 'min', 'det', 'mod', 'arcsin', 'arccos', 'arctan', 'sqrt'}
_CJK_RUN = re.compile(r'[\u4e00-\u9fff]+')
_MATH_SPAN = re.compile(r'(\\\[.*?\\\]|\$\$.*?\$\$|\\\(.*?\\\)|\$.*?\$)', re.DOTALL)

WRAP_CJK_IN_ARITH = True


def _convert_unicode_ops(s: str) -> str:
    for k, v in _OP_MAP.items():
        s = s.replace(k, (v + ' ') if v.startswith('\\') else v)
    if WRAP_CJK_IN_ARITH:
        s = _CJK_RUN.sub(lambda m: r'\text{' + m.group(0) + '}', s)
    return re.sub(r'[ \t]{2,}', ' ', s)


def _is_pure_arith_line(line: str) -> bool:
    t = line.strip()
    if not t or '\\(' in t or '\\[' in t or '$' in t or '<' in t:
        return False
    if not WRAP_CJK_IN_ARITH and re.search(r'[\u4e00-\u9fff]', t):
        return False
    if not _ARITH_ALLOWED.match(t) or not _ARITH_HASOP.search(t):
        return False
    return all(w.lower() in _KNOWN_FUNCS for w in re.findall(r'[A-Za-z]{2,}', t))


def normalize_arith(text: str) -> str:
    """Normalize Unicode operators to LaTeX and wrap pure-arithmetic lines as `\\( .. \\)`."""
    if not text:
        return text
    text = _MATH_SPAN.sub(lambda m: _convert_unicode_ops(m.group(0)), text)
    out = []
    for line in text.split('\n'):
        if _is_pure_arith_line(line):
            out.append('\\( ' + _convert_unicode_ops(line.strip()) + ' \\)')
        else:
            out.append(line)
    return '\n'.join(out)


def remove_block_fork_tags(result, simplify_left_paren=True, clean_formula_tail_flag=True,
                           norm_formula_flag=True):
    """Split on `<BLOCK>`, keep the text after each `<CHILD>`, and join in reading order."""
    seg_pattern = re.compile(r'[^<]*<CHILD>(.*)', re.DOTALL)
    lines = []
    for seg in result.split('<BLOCK>')[1:]:
        cat_m = re.match(r'\s*([a-zA-Z_]+)', seg)
        if cat_m and cat_m.group(1).lower() in ['chart', 'seal']:
            continue
        m = seg_pattern.match(seg)
        if not m:
            continue
        text = m.group(1).strip()
        text = re.sub(r'\b\w+\s*\[\s*[-\d.,\s]+\]\s*<(?:FORK|CHILD|BLOCK)>', '', text)
        text = re.sub(r'<(?:FORK|CHILD|BLOCK)>', '', text).strip()
        text = text.replace('The image is too blurry to recognize any text content.', '').strip()
        text = text.replace("The image contains no text or characters. It is a graphical element (a horizontal line with a vertical line) and does not contain any chart, graph, or data points that can be extracted. Therefore, the correct OCR output is an empty string.", "").strip()
        if not text or text == '[Non-Text]':
            continue
        if text.startswith('\\[') and not text.endswith('\n\\]'):
            text += '\n\\]'
        if text.startswith('<table>') and not text.endswith('</table>'):
            text += '</table>'
        if '\\[\n' in text and '\\\\' not in text:
            text = text.replace('\\[\n', '\\(').replace('\n\\]', '\\)')
        text = text.replace('\\) \\(', '\\)\n\n\\(')
        if '÷' in text and '\\(' not in text:
            text = '\\( ' + text + ' \\)'
        text = re.sub(r'\\tag\s*\{[^{}]*\}', '', text)
        text = text.replace('\\supset', '\\sqsupset')
        if simplify_left_paren:
            text = simplify_left_right(text)
        if clean_formula_tail_flag:
            text = clean_formula_tail(text)
        if norm_formula_flag:
            text = normalize_arith(text)
        lines.append(text)
    return '\n\n'.join(lines).strip()


def basename_to_md_name(img_path: str) -> str:
    return os.path.splitext(os.path.basename(img_path))[0] + ".md"


def convert_json(in_path, out_md_dir, simplify_left_paren=True,
                 clean_formula_tail_flag=True, norm_formula_flag=True) -> int:
    with open(in_path, "r", encoding="utf-8") as f:
        rows = json.load(f)
    os.makedirs(out_md_dir, exist_ok=True)
    n = 0
    for row in rows:
        img_path = row.get("img_path") or row.get("image_path")
        pred = row.get("pred") or row.get("prediction") or ""
        if not img_path:
            continue
        md = remove_block_fork_tags(pred, simplify_left_paren, clean_formula_tail_flag, norm_formula_flag)
        with open(os.path.join(out_md_dir, basename_to_md_name(img_path)), "w", encoding="utf-8") as f:
            f.write(md)
        n += 1
    return n


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", required=True, help="json path (list of {img_path, pred})")
    ap.add_argument("--out-md", required=True, help="output markdown folder")
    args = ap.parse_args()

    if Path(args.input).suffix.lower() != ".json":
        raise SystemExit(f"unsupported extension: {Path(args.input).suffix} (expects .json)")
    n = convert_json(args.input, args.out_md)
    print(f"[ok] wrote {n} markdown files -> {args.out_md}")


if __name__ == "__main__":
    main()