"""Split the Project Gutenberg plain text of Anton Chekhov, *The Schoolmistress, and Other Stories*
(Constance Garnett translation, "Tales of Chekhov, Volume 9", eBook #1732), into stories, then cut
scenes with the same rule as ../chunk.py, ../jacobs/split_se.py and ../wharton/split_gutenberg.py
(paragraph boundaries, ~320 words, tail >= 80).
Deterministic: titles are the 21 contents entries, matched as whole upper-case lines after the
contents; roman-numeral section headings and the "A Wayfarer's Story" subtitle inside Champagne are
dropped; the Gutenberg licence after END is dropped."""
import json, pathlib, re

TITLES = ['The Schoolmistress', 'A Nervous Breakdown', 'Misery', 'Champagne', 'After the Theatre',
          "A Lady's Story", 'In Exile', 'The Cattle-Dealers', 'Sorrow', 'On Official Duty',
          'The First-Class Passenger', 'A Tragic Actor', 'A Transgression', 'Small Fry',
          'The Requiem', 'In the Coach-House', 'Panic Fears', 'The Bet',
          "The Head-Gardener's Story", 'The Beauties', 'The Shoemaker and the Devil']
TARGET = 320
VOLUME = 'The Schoolmistress, and Other Stories (Anton Chekhov, tr. Constance Garnett, 1918)'


def norm(s):
    return s.replace('’', "'").strip()


raw = pathlib.Path('corpus/schoolmistress-and-other-stories.txt').read_text(encoding='utf-8').replace('\r', '')
raw = raw[:raw.index('*** END OF THE PROJECT GUTENBERG')]
lines = raw.split('\n')
start = next(i for i, l in enumerate(lines) if norm(l) == 'CONTENTS')
heads = []
for t in TITLES:
    hits = [i for i, l in enumerate(lines) if i > start + 22 and norm(l) == t.upper()]
    assert len(hits) == 1, (t, hits)
    heads.append(hits[0])
assert heads == sorted(heads)

stories = []
for n, (t, a) in enumerate(zip(TITLES, heads)):
    b = heads[n + 1] if n + 1 < len(heads) else len(lines)
    body = '\n'.join(lines[a + 1:b])
    paras = []
    for p in re.split(r'\n\s*\n', body):
        p = re.sub(r'\s+', ' ', p).strip()
        if not p or re.fullmatch(r'[IVX]+\.?', p): continue
        if norm(p).upper() == "A WAYFARER'S STORY": continue
        paras.append(p)
    stories.append({'volume': VOLUME, 'title': t, 'words': sum(len(p.split()) for p in paras),
                    'text': '\n\n'.join(paras)})

chunks = []
for si, s in enumerate(stories):
    buf, n = [], 0
    for p in s['text'].split('\n\n'):
        w = len(p.split())
        if n + w > TARGET and buf:
            chunks.append({'story': si, 'title': s['title'], 'volume': s['volume'], 'text': ' '.join(buf), 'words': n})
            buf, n = [], 0
        buf.append(p); n += w
    if buf and n >= 80:
        chunks.append({'story': si, 'title': s['title'], 'volume': s['volume'], 'text': ' '.join(buf), 'words': n})
    elif buf:
        # tail too short to stand alone: fold into the previous scene rather than drop it
        if chunks and chunks[-1]['story'] == si:
            chunks[-1]['text'] += ' ' + ' '.join(buf)
            chunks[-1]['words'] += n
        else:
            chunks.append({'story': si, 'title': s['title'], 'volume': s['volume'], 'text': ' '.join(buf), 'words': n})

pathlib.Path('stories.json').write_text(json.dumps(stories, indent=1, ensure_ascii=False), encoding='utf-8')
pathlib.Path('chunks.json').write_text(json.dumps(chunks, ensure_ascii=False), encoding='utf-8')
for si, s in enumerate(stories):
    print(f"{si:2} {s['title']:28} {s['words']:6} words  {sum(c['story']==si for c in chunks):3} scenes")
ws = sorted(c['words'] for c in chunks)
print(f'stories {len(stories)}  scenes {len(chunks)}  words/scene min {ws[0]} median {ws[len(ws)//2]} max {ws[-1]}  total {sum(ws):,}')