"""Post-processing utilities for HPD-Parsing demo: turns the raw
``<BLOCK> <type> [bbox] <CHILD> <content>`` stream into (a) clean markdown text
and (b) a list of typed bounding boxes for visualization.

Formula-cleaning helpers (``simplify_left_right`` / ``clean_formula_tail`` /
``normalize_arith``) are copied verbatim from the model repo's
``eval/hpd_to_markdown.py`` so the produced markdown matches the official
OmniDocBench post-processing.
"""

import re

from PIL import ImageDraw, ImageFont

# --- Formula-cleaning helpers (copied from eval/hpd_to_markdown.py) --------

_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 clean_text(text: str, simplify_left_paren=True, clean_formula_tail_flag=True,
               norm_formula_flag=True) -> str:
    """Apply the same per-block cleaning steps as the official hpd_to_markdown.py."""
    text = 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]':
        return ''
    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)
    return text


# --- Block parsing (bbox-preserving, unlike the official script) -----------

# type + [bbox], usually followed by <FORK|CHILD|BLOCK>. Container blocks such
# as `list [x1,y1,x2,y2]` may have no trailing tag at all -- after splitting on
# `<BLOCK>` the header is simply the entire segment -- so the tag is optional.
_BLOCK_HEADER = re.compile(
    r'([a-zA-Z_]+)\s*\[\s*([-\d.,\s]+)\]\s*(?:<(?:FORK|CHILD|BLOCK)>)?'
)
_NO_CONTENT_TYPES = {'chart', 'seal'}


def parse_blocks(raw_text: str):
    """Parse the ``<BLOCK> <type> [bbox] <CHILD> <content>`` stream.

    Returns a list of dicts: ``{"type": str, "bbox": [x1,y1,x2,y2] | None, "text": str}``.
    ``bbox`` values are in the model's native 0-1000 normalized coordinate space
    (confirmed via real inference: max observed values ~922/589 against a
    1240x1754 source image). ``text`` is the cleaned markdown for that block;
    it is empty for container blocks (e.g. ``list``, ``table`` wrapper) and for
    ``chart``/``seal`` blocks, matching the official markdown-export behavior.
    """
    blocks = []
    segments = raw_text.split('<BLOCK>')[1:]
    for seg in segments:
        header_m = _BLOCK_HEADER.match(seg.strip())
        # segments always start right after a <BLOCK>, so the header should be
        # at the very start; fall back to a bare type token if bbox is absent.
        type_m = re.match(r'\s*([a-zA-Z_]+)', seg)
        block_type = type_m.group(1) if type_m else 'unknown'

        bbox = None
        if header_m:
            block_type = header_m.group(1)
            nums = [float(x) for x in re.split(r'[,\s]+', header_m.group(2).strip()) if x]
            if len(nums) == 4:
                bbox = nums

        text = ''
        if block_type.lower() not in _NO_CONTENT_TYPES:
            content_m = re.search(r'<CHILD>(.*)', seg, re.DOTALL)
            if content_m:
                # stop at the next control tag if any leaked through
                raw_content = re.split(r'<(?:FORK|CHILD|BLOCK)>', content_m.group(1))[0]
                text = clean_text(raw_content)

        blocks.append({'type': block_type, 'bbox': bbox, 'text': text})
    return blocks


def blocks_to_markdown(blocks) -> str:
    """Join the non-empty block texts in reading order into one markdown string."""
    return '\n\n'.join(b['text'] for b in blocks if b['text']).strip()


# --- Bounding-box visualization ---------------------------------------------

_TYPE_COLORS = {
    'header': '#e74c3c',
    'title': '#e67e22',
    'text': '#2980b9',
    'list': '#8e44ad',
    'table': '#27ae60',
    'table_caption': '#16a085',
    'figure': '#2c3e50',
    'figure_caption': '#34495e',
    'chart': '#f39c12',
    'seal': '#c0392b',
    'formula': '#d35400',
}
_DEFAULT_COLOR = '#7f8c8d'


def draw_boxes_on_image(image, blocks):
    """Draw typed, numbered bounding boxes on a copy of ``image``.

    ``blocks`` bbox values are 0-1000 normalized coordinates in ``[x1, y1, x2, y2]``
    order; they are rescaled to ``image``'s actual pixel size before drawing.
    Blocks with ``bbox is None`` are skipped (still counted, just not drawn).
    Boxes are numbered in reading order (the order they appear in ``blocks``,
    i.e. the model's own reading-order stream) to visualize the parse order.
    """
    if image is None:
        return None
    canvas = image.convert('RGB').copy()
    draw = ImageDraw.Draw(canvas)
    width, height = canvas.size

    try:
        font = ImageFont.load_default()
    except Exception:
        font = None

    order = 0
    for block in blocks:
        bbox = block.get('bbox')
        if not bbox:
            continue
        order += 1
        x1, y1, x2, y2 = bbox
        px1, py1 = x1 / 1000.0 * width, y1 / 1000.0 * height
        px2, py2 = x2 / 1000.0 * width, y2 / 1000.0 * height
        color = _TYPE_COLORS.get(block['type'].lower(), _DEFAULT_COLOR)
        draw.rectangle([px1, py1, px2, py2], outline=color, width=2)
        label = f"{order}. {block['type']}"
        text_y = max(0, py1 - 12)
        if font is not None:
            text_w = draw.textlength(label, font=font)
        else:
            text_w = len(label) * 6
        draw.rectangle([px1, text_y, px1 + text_w + 2, text_y + 11], fill=color)
        if font is not None:
            draw.text((px1 + 1, text_y), label, fill='white', font=font)
        else:
            draw.text((px1 + 1, text_y), label, fill='white')

    return canvas


if __name__ == "__main__":
    print("Module self-test requires a probe output file; see git history for the "
          "original probe script used during development.")