"""Reduce each scene to one generic sentence about the SITUATION, not the story.

The 2026-09-10 probe found that prose embeddings cluster by story identity --
names, setting, narrator voice -- rather than by what kind of situation a scene
is. This strips exactly those signals before embedding.

  python3 summarise.py --limit 40      # pilot: time it and read the output
  python3 summarise.py                 # everything, resuming from summaries.json

Model from STP_CHAT_MODEL (default qwen2.5:3b-instruct): the requirement is a
small local model, so the prompt is written for one.
"""
import argparse, json, os, pathlib, re, time
import ollama

ap = argparse.ArgumentParser()
ap.add_argument('--limit', type=int, default=None)
args = ap.parse_args()
MODEL = os.environ.get('STP_CHAT_MODEL', 'qwen2.5:3b-instruct')

# v1 gave examples of roles ("a clerk", "a landlady", "a rich man", "a city
# park"). A 3B model copied them into scenes that contain none of those people,
# which would make unrelated scenes cluster together and fake the very
# recurrence this probe measures. v2 names no example role at all.
SYSTEM = """Describe the social situation in a passage of fiction in ONE sentence.

Rules:
- Exactly one sentence, at most 25 words.
- Never use anyone's name. Refer to each person only by the role the passage gives them: their job, their family tie, or their position.
- Never name a place, business, street or city.
- Say what is happening between the people: who wants what from whom, who holds power, what is at stake.
- Do not describe the writing style, the narrator, or the mood.
- Output only the sentence."""

# The v2 pilot also tried one retry with a reminder whenever a name slipped
# through. It fired 35 times in 40, nearly doubled the run, and 27 summaries
# still carried names -- so there is no retry. Names are removed afterwards by
# strip_names.py, deterministically.

# A full stop after these is not the end of a sentence ("Mr. Peters", "St. Bernard").
END = re.compile(r'(?<!\bMr)(?<!\bMrs)(?<!\bMs)(?<!\bDr)(?<!\bSt)(?<!\bNo)[.!?](?:\s|$)')

def first_sentence(text):
    text = re.sub(r'\s+', ' ', text).strip().strip('"')
    m = END.search(text)
    return text[:m.end()].strip() if m else text

chunks = json.loads(pathlib.Path('chunks.json').read_text(encoding='utf-8'))
if args.limit: chunks = chunks[:args.limit]
out = pathlib.Path('summaries.json')
state = json.loads(out.read_text()) if out.exists() else {'model': MODEL, 'prompt': SYSTEM, 'summaries': []}
if state['model'] != MODEL or state['prompt'] != SYSTEM:
    raise SystemExit('summaries.json was made with a different model or prompt; move it aside first')
done = state['summaries']
print(f'{len(chunks)} scenes, resuming at {len(done)}, model {MODEL}', flush=True)

t0, gen_ns, pre_ns = time.time(), 0, 0
for i in range(len(done), len(chunks)):
    text, g, p = ollama.chat(MODEL, SYSTEM, chunks[i]['text'])
    gen_ns += g; pre_ns += p
    done.append({'chunk': i, 'story': chunks[i]['story'],
                 'summary': first_sentence(text), 'model_output': text})
    out.write_text(json.dumps(state, indent=1))
    if (i + 1) % 10 == 0 or i + 1 == len(chunks):
        el = time.time() - t0
        print(f'  {i + 1}/{len(chunks)}  {el:.0f}s  prefill {pre_ns / 1e9:.0f}s  gen {gen_ns / 1e9:.0f}s', flush=True)
print(f'done: {len(done)} summaries, {time.time() - t0:.0f}s this session')