#!/usr/bin/env python3
"""Stage 2: turn per-page OCR lines into a clean, structured manuscript.

Handles running-header/footer removal, scan-speckle noise, per-page margin
drift, paragraph detection via first-line indent, centered chapter headings,
rule-delimited footnotes, block quotations, verse, cross-page paragraph
joining, and corpus-driven dehyphenation.
"""

import os as _os
# resolve paths relative to this archive rather than the machine it was built on
_ROOT = _os.path.dirname(_os.path.dirname(_os.path.abspath(__file__)))
ARCHIVE = _ROOT
_DATA = _os.path.join(_ROOT, 'data')
import json, re, statistics
from collections import Counter

LINES = json.load(open('/tmp/lines.json'))
RULES = {int(k): v for k, v in json.load(open('/tmp/rules.json')).items()}
_ov = json.load(open(_os.path.join(_DATA, 'overrides.json')))
OVERRIDES = {int(k): v for k, v in _ov.items() if k.isdigit()}

FRONT_MATTER_PAGES = {1, 2, 3}
PUNCT_ONLY = re.compile(r'^[^0-9A-Za-z]*$')
HEADER_RE = re.compile(r'(walton|rev\b|revision)', re.I)
PAGENUM_RE = re.compile(r'^[^0-9]{0,3}\d{1,3}[^0-9]{0,3}$')


# an alphabetic run of 4+ characters containing a vowel — i.e. something that
# could plausibly be a real word rather than scanner speckle
WORDLIKE = re.compile(r'[A-Za-zÀ-ÿ]*[aeiouyAEIOUY][A-Za-zÀ-ÿ]*')


def is_noise(l):
    t = l['text'].strip()
    if not t or PUNCT_ONLY.match(t):
        return True
    # the typist's page-break marker, not manuscript content
    if re.fullmatch(r'\(?continued\)?[.:]?', t, re.I):
        return True
    # the hand-written folio number at the foot of every page. Body text never
    # reaches this band in so few words, and these read as low-confidence
    # pseudo-words ("kore", "Qazds"), so they must be caught before the
    # word-like guard below would rescue them.
    if l['top'] > 3100 and l['nwords'] <= 2 and l['conf'] < 80:
        return True
    # never discard a line that contains a plausible word, however faint: a
    # short, badly-read line may still be real text (e.g. a one-word line
    # completing a sentence carried over from the previous page)
    if any(len(m.group(0)) >= 4 for m in WORDLIKE.finditer(t)):
        return False
    letters = sum(c.isalpha() for c in t)
    if l['conf'] < 50 and l['nwords'] <= 6 and letters < 12:
        return True
    if l['conf'] < 30 and l['nwords'] <= 10:
        return True
    if letters <= 2 and l['nwords'] <= 3:
        return True
    # short, low-confidence fragments: pen marks and strike-through debris
    if l['conf'] < 75 and l['nwords'] <= 3 and letters <= 7:
        return True
    return False


def classify_page(p):
    pno = p['page']
    lines = [l for l in p['lines'] if not is_noise(l)]
    if not lines:
        return None

    body = []
    for l in lines:
        if l['top'] < 210 and (HEADER_RE.search(l['text']) or l['nwords'] <= 6):
            continue                                    # running header
        # a few headers sit slightly lower; match them by their distinctive
        # "<name> Nth rev." shape rather than by position alone
        if l['top'] < 420 and l['nwords'] <= 6 \
                and re.search(r'walton', l['text'], re.I):
            continue
        body.append(l)
    body = [l for l in body
            if not (l['top'] > 3050 and l['nwords'] <= 4
                    and (PAGENUM_RE.match(l['text'].strip())
                         or sum(c.isalpha() for c in l['text']) <= 3))]
    if not body:
        return None

    longs = [l for l in body if l['nwords'] >= 8]
    margin = statistics.median([l['left'] for l in longs]) if longs \
        else min(l['left'] for l in body)
    body_right = statistics.median([l['right'] for l in longs]) if longs \
        else max(l['right'] for l in body)

    # ---- footnotes: everything below the detected separator rule ----
    # The rule often overlaps the first footnote line's bounding box (they
    # nearly touch), so allow a tolerance rather than a strict top > rule_y.
    notes = []
    rule_y = RULES.get(pno, {}).get('rule_y')
    if rule_y:
        TOL = 45
        notes = [l for l in body if l['top'] > rule_y - TOL]
        body = [l for l in body if l['top'] <= rule_y - TOL]
    else:
        # no rule found: fall back on an asterisk-led single-spaced tail
        for i, l in enumerate(body):
            if i == 0 or l['top'] < 2200:
                continue
            if re.match(r'^\s*[\*\+¥]', l['text']) and l['nwords'] >= 3 \
                    and l['left'] < margin + 180:
                notes, body = body[i:], body[:i]
                break

    return {'page': pno, 'body': body, 'notes': notes, 'margin': margin,
            'body_right': body_right}


PAGES = [x for x in (classify_page(p) for p in LINES) if x]

# ------------------------------------------------- dehyphenation vocabulary
tokens = Counter()
for p in PAGES:
    for l in p['body'] + p['notes']:
        for w in re.findall(r"[A-Za-z][A-Za-z'\-]*", l['text']):
            tokens[w.lower()] += 1


# The scanner renders a line-end hyphen variously as '-', '—', '-—', '-~', '¬'.
HYPH_END = re.compile(r'[A-Za-zÀ-ÿ][-‐-―~¬]+\s*$')


def dehyphen(first, second):
    """Rejoin a word split across a line break."""
    a = re.sub(r'[-‐-―~¬]+\s*$', '', first)
    b = second.lstrip()
    ma = re.search(r"[A-Za-zÀ-ÿ]+$", a)          # trailing word of the first part
    mb = re.match(r"[A-Za-zÀ-ÿ]+", b)            # leading word of the second
    if not ma or not mb:
        return a + b
    stem, head, rest = ma.group(0), mb.group(0), b[mb.end():]
    prefix = a[:ma.start()]

    def n(w):
        return tokens.get(w.lower(), 0)

    # a capital mid-word is an artefact of the split ("neces-" + "Sary");
    # n() is case-insensitive, so test the raw casing here
    if head[:1].isupper() and stem[-1:].islower() \
            and (stem + head) not in tokens and n(stem + head) > 0:
        head = head.lower()
    if n(stem + '-' + head) > n(stem + head):
        return prefix + stem + '-' + head + rest
    if n(stem + head) > 0:
        return prefix + stem + head + rest
    if head[:1].isupper() or stem[-1:].isdigit():
        return prefix + stem + '-' + head + rest
    return prefix + stem + head + rest


def join_lines(ls):
    buf = ''
    for l in ls:
        t = l['text'].strip()
        if not buf:
            buf = t
        elif HYPH_END.search(buf):
            buf = dehyphen(buf, t)
        else:
            buf += ' ' + t
    return re.sub(r'\s+', ' ', buf).strip()


def segment(lines, margin, body_right):
    """Split body lines into runs: ('body'|'quote'|'verse', [lines]).

    Paragraph first-line indent varies by page (roughly 150-280px), so indent
    alone can't identify quoted matter. Quoted blocks are also *narrower* than
    the body measure. A candidate run is accepted as a quote only if it has
    several lines or is clearly narrow — which rejects faint page-bottom lines
    whose left edge the OCR mis-measured.
    """
    cand, cur, kind = [], [], 'body'
    for i, l in enumerate(lines):
        indented = l['left'] > margin + 190
        k = 'quote' if indented else 'body'
        if k != kind and cur:
            cand.append((kind, cur))
            cur = []
        kind = k
        cur.append(l)
    if cur:
        cand.append((kind, cur))

    out = []
    for k, ls in cand:
        if not ls:
            continue
        if k == 'quote':
            rgap = statistics.median([body_right - x['right'] for x in ls])
            if not (len(ls) >= 3 or rgap > 200):
                k = 'body'                     # mis-measured faint body line
            elif statistics.mean(x['nwords'] for x in ls) <= 4.5:
                k = 'verse'
        if out and out[-1][0] == k == 'body':
            out[-1] = ('body', out[-1][1] + ls)
        else:
            out.append((k, ls))
    return out


SENT_END = re.compile(r'[.!?][\'"”’)\]]?$|[\'"”’]$')


def paragraphs(lines, margin):
    """Group body lines into paragraphs by first-line indent.

    Indent alone is unreliable: on faint pages the OCR sometimes captures only
    the right-hand part of a line, so its left edge looks like an indent. A new
    paragraph can only follow a completed sentence, so an indented line whose
    predecessor doesn't end a sentence is treated as a continuation.
    """
    groups, cur = [], []
    for l in lines:
        indented = l['left'] > margin + 85
        if indented and cur and SENT_END.search(cur[-1]['text'].strip()):
            groups.append(cur)
            cur = []
        cur.append(l)
    if cur:
        groups.append(cur)
    return groups


# ------------------------------------------------ chapter headings from TOC
TOC = [
    ('Prologue', 'AFTERNOON TEA IN AFRICA'),
    ('Chapter I', 'AN IDEA IS BORN'),
    ('Chapter II', 'BAD NEWS FROM LONDON'),
    ('Chapter III', "PROFESSOR ZUPITZA'S IDEA"),
    ('Chapter IV', 'THE ARMY SAYS "NO"!'),
    ('Chapter V', 'THE CHINA AFFAIR'),
    ('Chapter VI', 'KAPITÄNLEUTNANT BOCKHOLT'),
    ('Chapter VII', 'A QUESTION OF JUDGMENT'),
    ('Chapter VIII', 'A CLOSE CALL'),
    ('Chapter IX', 'ORDEAL OVER TURKEY'),
    ('Chapter X', 'EAVESDROPPER EWING AND HIS CREW'),
    ('Chapter XI', '"CHINA" AT LAST'),
    ('Chapter XII', 'RETREAT'),
    ('Chapter XIII', 'THE END OF "L 59"'),
    ('Chapter XIV', 'EPILOGUE'),
]


def norm(s):
    return re.sub(r'[^a-z]', '', s.lower())


TOC_BY_TITLE = {norm(t): (lab, t) for lab, t in TOC}


def match_toc(cand):
    """Fuzzy-match an OCR'd heading against the table of contents."""
    c = norm(cand)
    if not c:
        return None
    best, score = None, 0.0
    for key, val in TOC_BY_TITLE.items():
        if c == key:
            return val
        # character-bigram overlap
        A = {key[i:i+2] for i in range(len(key)-1)}
        B = {c[i:i+2] for i in range(len(c)-1)}
        if not A or not B:
            continue
        s = 2 * len(A & B) / (len(A) + len(B))
        if s > score:
            best, score = val, s
    return best if score >= 0.62 else None


# --------------------------------------------------------------- assemble
# Paragraphs get a stable id (pid). A paragraph that runs over a page break is
# held in `carry` and merged into the next page's opening paragraph, keeping
# its original pid — so footnotes anchored to it still resolve afterwards.
doc = []
carry = None                     # {'text','pages','pid'} or None
used_toc = set()
_pid = [0]


def new_para(text, pages):
    _pid[0] += 1
    return {'type': 'para', 'text': text, 'pages': pages, 'pid': _pid[0]}


def flush_carry():
    global carry
    if carry:
        doc.append(carry)
        carry = None


for p in PAGES:
    pno = p['page']
    if pno in FRONT_MATTER_PAGES:
        continue

    page_pids = []

    # ---- manually transcribed pages (faint carbon copies) ----
    if pno in OVERRIDES:
        ov = OVERRIDES[pno]
        if ov.get('heading'):
            flush_carry()
            lab, title = ov['heading']
            used_toc.add(norm(title))
            doc.append({'type': 'heading', 'label': lab, 'title': title,
                        'page': pno})
        made = []
        for i, t in enumerate(ov['paras']):
            if i == 0 and ov.get('continues') and carry:
                carry['text'] = re.sub(r'\s+', ' ', carry['text'] + ' ' + t)
                carry['pages'] = carry['pages'] + [pno]
                b = carry
                carry = None
            else:
                flush_carry()
                b = new_para(t, [pno])
            b['manual'] = True
            doc.append(b)
            made.append(b)
            page_pids.append(b['pid'])
        if ov.get('footnote'):
            anchor = next((b['pid'] for b in reversed(made) if '*' in b['text']),
                          made[-1]['pid'] if made else None)
            doc.append({'type': 'footnote', 'text': ov['footnote'],
                        'page': pno, 'anchor': anchor, 'manual': True})
        if ov.get('open_tail') and made:
            carry = made[-1]
            doc.remove(carry)
        continue

    margin, body = p['margin'], p['body']

    # centered heading lines at the head of a chapter page
    head_lines, rest = [], []
    for l in body:
        centered = l['left'] > margin + 330 and l['nwords'] <= 9 and l['top'] < 900
        if centered and not rest:
            head_lines.append(l)
        else:
            rest.append(l)

    hit = None
    if head_lines:
        hit = match_toc(' '.join(h['text'] for h in head_lines))
        if not hit:                       # maybe only the title line is centered
            hit = match_toc(head_lines[-1]['text'])
        if not hit:
            rest = head_lines + rest      # not a heading after all (e.g. verse)
            head_lines = []

    started_chapter = False
    if hit and norm(hit[1]) not in used_toc:
        used_toc.add(norm(hit[1]))
        flush_carry()
        doc.append({'type': 'heading', 'label': hit[0], 'title': hit[1],
                    'page': pno})
        started_chapter = True
    elif head_lines:
        rest = head_lines + rest

    for kind, ls in segment(rest, margin, p["body_right"]):
        if kind == 'body':
            for gi, g in enumerate(paragraphs(ls, margin)):
                text = join_lines(g)
                first_indented = g[0]['left'] > margin + 85
                if carry and gi == 0 and not first_indented \
                        and not started_chapter:
                    carry['text'] = dehyphen(carry['text'], text) \
                        if HYPH_END.search(carry['text']) \
                        else re.sub(r'\s+', ' ', carry['text'] + ' ' + text)
                    carry['pages'] = carry['pages'] + [pno]
                    b = carry
                    carry = None
                else:
                    flush_carry()
                    b = new_para(text, [pno])
                doc.append(b)
                page_pids.append(b['pid'])
        else:
            flush_carry()
            blk = {'type': kind, 'lines': [x['text'].strip() for x in ls],
                   'page': pno}
            if kind == 'quote':
                # prose quotations reflow like body text (so line-break
                # hyphens rejoin); verse keeps its line breaks
                blk['text'] = join_lines(ls)
            doc.append(blk)

    if p['notes']:
        ntext = join_lines(p['notes'])
        ntext = re.sub(r'^\s*[\*\+¥®]+\s*', '', ntext)
        anchor = None
        by_pid = {b['pid']: b for b in doc if b['type'] == 'para'}
        for pid in reversed(page_pids):
            if '*' in by_pid[pid]['text']:
                anchor = pid
                break
        if anchor is None and page_pids:
            anchor = page_pids[-1]
        doc.append({'type': 'footnote', 'text': ntext, 'page': pno,
                    'anchor': anchor})

    # the last paragraph on the page may continue onto the next one
    if page_pids and doc and doc[-1]['type'] == 'para' \
            and doc[-1]['pid'] == page_pids[-1]:
        carry = doc.pop()
    elif page_pids and p['notes']:
        last = next((b for b in reversed(doc)
                     if b['type'] == 'para' and b['pid'] == page_pids[-1]), None)
        if last is not None:
            carry = last
            doc.remove(last)

flush_carry()

# --------------------------------------------- OCR corrections (verified)
FIX = json.load(open(_os.path.join(_DATA, 'fixes.json')))
TOKEN_RE = re.compile(
    r"(?<![A-Za-z&])(" + '|'.join(sorted((re.escape(k) for k in FIX['tokens']),
                                        key=len, reverse=True)) + r")(?![A-Za-z])")
n_phrase = n_token = 0
# record which manuscript page each correction landed on, for the audit log
FIX_LOG = {'phrases': [], 'tokens': []}
_cur_pages = []


def fix_text(s):
    global n_phrase, n_token
    # Normalise whitespace and quotes FIRST. The manuscript was typed, so every
    # quote is straight; tesseract emits a mix of curly forms, sometimes doubles
    # the opening mark, and leaves ragged spacing. Doing this before the phrase
    # pass means the phrases in fixes.json can be written plainly and still match.
    s = s.replace('“', '"').replace('”', '"').replace('‘', "'").replace('’', "'")
    s = re.sub(r'"\s*"', '"', s)
    s = re.sub(r"(?<![A-Za-z])'\s*'(?![A-Za-z])", "'", s)
    # stray pipes / underscores are scan speckle, not typed characters
    s = re.sub(r'(?<![A-Za-z0-9])[|_]+(?![A-Za-z0-9])', ' ', s)
    # spacing artefacts left by the typewriter/OCR
    s = re.sub(r'\s+([,.;:!?])', r'\1', s)
    s = re.sub(r'\(\s+', '(', s)
    s = re.sub(r'\s+\)', ')', s)
    s = re.sub(r'\s+', ' ', s).strip()

    for a, b in FIX['phrases']:
        if a in s:
            n = s.count(a)
            n_phrase += n
            FIX_LOG['phrases'].append({'from': a, 'to': b, 'n': n,
                                       'pages': list(_cur_pages)})
            s = s.replace(a, b)

    def sub(m):
        global n_token
        n_token += 1
        FIX_LOG['tokens'].append({'from': m.group(1),
                                  'to': FIX['tokens'][m.group(1)],
                                  'pages': list(_cur_pages)})
        return FIX['tokens'][m.group(1)]
    s = TOKEN_RE.sub(sub, s)
    # stray pipes / underscores are scan speckle, not typed characters
    s = re.sub(r'(?<![A-Za-z0-9])[|_]+(?![A-Za-z0-9])', ' ', s)
    # tidy spacing artefacts left by the typewriter/OCR
    s = re.sub(r'\s+([,.;:!?])', r'\1', s)
    s = re.sub(r'\(\s+', '(', s)
    s = re.sub(r'\s+\)', ')', s)
    s = re.sub(r'\s{2,}', ' ', s)
    return s.strip()


for b in doc:
    _cur_pages = b.get('pages') or ([b['page']] if 'page' in b else [])
    if 'text' in b:
        b['text'] = fix_text(b['text'])
    if 'lines' in b:
        b['lines'] = [fix_text(x) for x in b['lines']]

# A paragraph always opens a sentence; tesseract occasionally lower-cases the
# leading capital (verified on the scan for the cases this affects).
n_cap = 0
cap_log = []
for b in doc:
    if b['type'] == 'para' and b['text'][:1].islower():
        cap_log.append({'pages': b.get('pages', []),
                        'word': b['text'].split()[0]})
        b['text'] = b['text'][0].upper() + b['text'][1:]
        n_cap += 1

json.dump({'phrases': FIX_LOG['phrases'], 'tokens': FIX_LOG['tokens'],
           'capitalisations': cap_log},
          open('/tmp/fixlog.json', 'w'), ensure_ascii=False, indent=1)

json.dump({'toc': TOC, 'blocks': doc}, open('/tmp/doc.json', 'w'),
          ensure_ascii=False)

# ------------------------------------------------------------------ report
print(f"blocks: {len(doc)}   {Counter(b['type'] for b in doc)}")
print(f"words: {sum(len(b.get('text', '').split()) for b in doc):,}")
print(f"OCR fixes applied: {n_phrase} phrase, {n_token} token, {n_cap} capitalisation")
print(f"\nchapters matched: {len(used_toc)} / {len(TOC)}")
for b in doc:
    if b['type'] == 'heading':
        print(f"   p{b['page']:3d}  {b['label']:<13} {b['title']}")
missing = [t for _, t in TOC if norm(t) not in used_toc]
print("MISSING:", missing if missing else "none")
print("\nquote/verse blocks:")
for b in doc:
    if b['type'] in ('quote', 'verse'):
        print(f"   p{b['page']:3d} {b['type']:6} {' / '.join(b['lines'])[:78]}")
print(f"\nfootnotes: {sum(1 for b in doc if b['type']=='footnote')}")
