Kansho/backend/journal_shape.py

"""Local journal shape after the model. Does not invent content."""
from __future__ import annotations

import re

NEW_PHASE = re.compile(
    r"^(?:Heute|Gestern|Morgens|Mittags|Abends|Nachts|Danach|Später|Dann|Zuerst|"
    r"Schließlich|Irgendwann|Unterwegs|Zurück|"
    r"Um\s+\d|Nach\s+(?:ca\.|einer|einem|dem|der|kurzer|weiteren)|"
    r"Als\s+(?:wir|ich)\b)",
    re.I,
)
_LETTER = r"A-Za-zÄÖÜäöüß"
_WORD = re.compile(rf"[{_LETTER}]+")
DATIVE_PREPS = {
    "mit", "bei", "von", "zu", "nach", "aus", "außer", "gegenüber", "gemäß",
}
NOM_CUES = {
    "dass", "daß", "ob", "weil", "wenn", "als", "während", "obwohl", "damit",
    "nachdem", "bevor", "sodass", "sofern",
}
THING_GOVERNORS = {
    "esse", "essen", "isst", "aß", "aßest", "gegessen",
    "koche", "kochen", "kochte", "gekocht",
    "trinke", "trinken", "trank", "getrunken",
    "kaufe", "kaufen", "kaufte", "gekauft",
    "bestelle", "bestellen", "bestellte",
    "hole", "holen", "holte",
}
FEM_HINTS = {"sie", "ihr", "ihre", "ihrem", "ihren", "ihrer", "ihres"}
MASC_HINTS = {"er", "ihn", "ihm", "seine", "seinem", "seinen", "seiner", "seines"}
KEEP_NAMES_PER_PARAGRAPH = 2


def _strip_roles(text: str) -> str:
    lines = [re.sub(r"^(?:user|assistant):\s*", "", line, flags=re.I) for line in (text or "").splitlines()]
    return "\n".join(lines).strip()


def _norm(text: str) -> str:
    return re.sub(r"[^\w\s]+", " ", re.sub(r"\s+", " ", _strip_roles(text).lower())).strip()


def _turns(user_bodies: list[str]) -> list[str]:
    return [re.sub(r"\s+", " ", item).strip() for item in user_bodies if (item or "").strip()]


def is_verbatim_join(body: str, user_bodies: list[str]) -> bool:
    turns = _turns(user_bodies)
    if not turns or not (body or "").strip():
        return False
    got = _norm(body)
    return got == _norm(" ".join(turns)) or got == _norm("\n\n".join(turns))


def paragraphize(body: str) -> str:
    text = (body or "").strip()
    if "\n\n" in text:
        return text
    sentences = [item.strip() for item in re.split(r"(?<=[.!?])\s+", text) if item.strip()]
    if len(sentences) < 4:
        return text
    chunks: list[str] = []
    current: list[str] = []
    for sentence in sentences:
        split = bool(current) and (NEW_PHASE.match(sentence) or len(current) >= 3)
        if split:
            chunks.append(" ".join(current))
            current = [sentence]
        else:
            current.append(sentence)
    if current:
        chunks.append(" ".join(current))
    return "\n\n".join(chunks)


def _words_before(text: str, index: int, n: int = 4) -> list[str]:
    return [word.lower() for word in _WORD.findall(text[:index])[-n:]]


def _words_after(text: str, index: int, n: int = 3) -> list[str]:
    return [word.lower() for word in _WORD.findall(text[index:])[:n]]


def _is_thing_mention(text: str, start: int, end: int) -> bool:
    prev = _words_before(text, start)
    nxt = _words_after(text, end)
    return any(word in THING_GOVERNORS for word in prev + nxt)


def _sentence_start(text: str, index: int) -> bool:
    prefix = text[:index].rstrip()
    return not prefix or prefix[-1] in ".!?\n"


def infer_person_gender(label: str, source_texts: list[str]) -> str | None:
    blob = " ".join(source_texts or [])
    if not label or not blob:
        return None
    pattern = re.compile(rf"(?<![{_LETTER}]){re.escape(label)}(?![{_LETTER}])", re.I)
    first = pattern.search(blob)
    if not first:
        return None
    after = [word.lower() for word in _WORD.findall(blob[first.start():])]
    fem = sum(1 for word in after if word in FEM_HINTS)
    masc = sum(1 for word in after if word in MASC_HINTS)
    if fem > masc:
        return "f"
    if masc > fem:
        return "m"
    return None


def _pronoun(gender: str, prev: str, sentence_start: bool) -> str:
    dative = prev in DATIVE_PREPS
    if gender == "f":
        form = "ihr" if dative else "sie"
    elif sentence_start or prev in NOM_CUES:
        form = "er"
    elif dative:
        form = "ihm"
    else:
        form = "ihn"
    if sentence_start:
        return form[:1].upper() + form[1:]
    return form


def naturalize_person_mentions(body: str, labels: list[str], source_texts: list[str]) -> str:
    """After demask: keep the name once or twice per passage, then pronouns. No new facts."""
    text = body or ""
    people = []
    for label in sorted({item.strip() for item in (labels or []) if (item or "").strip()}, key=len, reverse=True):
        gender = infer_person_gender(label, source_texts)
        if gender:
            people.append((label, gender))
    if not people:
        return text
    paragraphs = re.split(r"(\n\n+)", text)
    out: list[str] = []
    for block in paragraphs:
        if not block or block.startswith("\n"):
            out.append(block)
            continue
        piece = block
        for label, gender in people:
            pattern = re.compile(rf"(?<![{_LETTER}]){re.escape(label)}(?![{_LETTER}])", re.I)
            matches = [
                match
                for match in pattern.finditer(piece)
                if not _is_thing_mention(piece, match.start(), match.end())
            ]
            if len(matches) <= KEEP_NAMES_PER_PARAGRAPH:
                continue
            replacements: list[tuple[int, int, str]] = []
            kept = 0
            for match in matches:
                kept += 1
                if kept <= KEEP_NAMES_PER_PARAGRAPH:
                    continue
                prev = (_words_before(piece, match.start(), 1) or [""])[0]
                form = _pronoun(gender, prev, _sentence_start(piece, match.start()))
                replacements.append((match.start(), match.end(), form))
            for start, end, form in reversed(replacements):
                piece = piece[:start] + form + piece[end:]
        out.append(piece)
    return "".join(out)


def shape_journal(
    title: str,
    body: str,
    user_bodies: list[str],
    person_labels: list[str] | None = None,
) -> tuple[str, str]:
    turns = _turns(user_bodies)
    text = (body or "").strip()
    heading = (title or "").strip()
    text = _strip_roles(text)
    if is_verbatim_join(text, turns) and len(turns) > 1:
        text = "\n\n".join(turns)
    else:
        text = paragraphize(text)
    if person_labels:
        text = naturalize_person_mentions(text, person_labels, turns)
    return heading, text