app/pipeline/alignment.py

"""Paragraph-preserving sentence alignment for hybrid rewrite."""

from __future__ import annotations

from dataclasses import dataclass
from difflib import SequenceMatcher

from app.pipeline.nlp import get_nlp
from app.pipeline.normalize import split_paragraphs, split_sentences_regex


@dataclass
class AlignedUnit:
    """One aligned source/candidate span within a paragraph."""

    source: str
    candidate: str
    paragraph_index: int
    confidence: float
    kind: str = "1:1"  # 1:1 | 2:1 | 1:2 | para


def _split_sentences(text: str) -> list[str]:
    text = (text or "").strip()
    if not text:
        return []
    nlp = get_nlp()
    if nlp is not None:
        return [s.text.strip() for s in nlp(text).sents if s.text.strip()]
    return split_sentences_regex(text)


def iter_source_units(source: str) -> list[tuple[int, str]]:
    """Paragraph-preserving sentence units from the grammar-corrected source."""
    out: list[tuple[int, str]] = []
    for idx, para in enumerate(split_paragraphs(source)):
        sents = _split_sentences(para)
        if not sents and para.strip():
            out.append((idx, para.strip()))
            continue
        for sent in sents:
            if sent.strip():
                out.append((idx, sent.strip()))
    return out


def _sim(a: str, b: str) -> float:
    return SequenceMatcher(None, a.lower().strip(), b.lower().strip()).ratio()


def align_paragraph_sentences(
    source_para: str,
    candidate_para: str,
    *,
    paragraph_index: int = 0,
    min_confidence: float = 0.35,
) -> list[AlignedUnit]:
    """Monotonic alignment of sentences within one paragraph."""
    src = _split_sentences(source_para)
    cand = _split_sentences(candidate_para)
    if not src and not cand:
        return []
    if not src:
        return [
            AlignedUnit("", c, paragraph_index, 0.0, "1:1")
            for c in cand
        ]
    if not cand:
        return [
            AlignedUnit(s, "", paragraph_index, 0.0, "1:1")
            for s in src
        ]

    # Equal counts → order-based pairs (verify confidence)
    if len(src) == len(cand):
        units: list[AlignedUnit] = []
        for s, c in zip(src, cand):
            conf = _sim(s, c)
            if conf < min_confidence:
                # Fall back to whole-paragraph unit when pairing looks wrong
                return [
                    AlignedUnit(
                        source_para.strip(),
                        candidate_para.strip(),
                        paragraph_index,
                        _sim(source_para, candidate_para),
                        "para",
                    )
                ]
            units.append(AlignedUnit(s, c, paragraph_index, conf, "1:1"))
        return units

    # Greedy monotonic merge for unequal counts
    units = []
    i = j = 0
    while i < len(src) and j < len(cand):
        one = _sim(src[i], cand[j])
        merge_src = (
            _sim(" ".join(src[i : i + 2]), cand[j]) if i + 1 < len(src) else -1.0
        )
        merge_cand = (
            _sim(src[i], " ".join(cand[j : j + 2])) if j + 1 < len(cand) else -1.0
        )
        best = max(one, merge_src, merge_cand)
        if best < min_confidence:
            return [
                AlignedUnit(
                    source_para.strip(),
                    candidate_para.strip(),
                    paragraph_index,
                    _sim(source_para, candidate_para),
                    "para",
                )
            ]
        if best == merge_src and merge_src >= one and merge_src >= merge_cand:
            units.append(
                AlignedUnit(
                    " ".join(src[i : i + 2]),
                    cand[j],
                    paragraph_index,
                    merge_src,
                    "2:1",
                )
            )
            i += 2
            j += 1
        elif best == merge_cand and merge_cand >= one:
            units.append(
                AlignedUnit(
                    src[i],
                    " ".join(cand[j : j + 2]),
                    paragraph_index,
                    merge_cand,
                    "1:2",
                )
            )
            i += 1
            j += 2
        else:
            units.append(AlignedUnit(src[i], cand[j], paragraph_index, one, "1:1"))
            i += 1
            j += 1

    # Trailing leftovers → attach to last unit or emit unpaired
    while i < len(src):
        if units:
            last = units[-1]
            units[-1] = AlignedUnit(
                (last.source + " " + src[i]).strip(),
                last.candidate,
                paragraph_index,
                _sim(last.source + " " + src[i], last.candidate),
                "2:1",
            )
        else:
            units.append(AlignedUnit(src[i], "", paragraph_index, 0.0, "1:1"))
        i += 1
    while j < len(cand):
        if units:
            last = units[-1]
            units[-1] = AlignedUnit(
                last.source,
                (last.candidate + " " + cand[j]).strip(),
                paragraph_index,
                _sim(last.source, last.candidate + " " + cand[j]),
                "1:2",
            )
        else:
            units.append(AlignedUnit("", cand[j], paragraph_index, 0.0, "1:1"))
        j += 1
    return units


def align_documents(source: str, candidate: str) -> list[AlignedUnit]:
    """Align source and candidate paragraph-by-paragraph."""
    s_paras = split_paragraphs(source)
    c_paras = split_paragraphs(candidate)
    # Pad shorter side so paragraphs stay in order
    n = max(len(s_paras), len(c_paras))
    while len(s_paras) < n:
        s_paras.append("")
    while len(c_paras) < n:
        c_paras.append("")
    out: list[AlignedUnit] = []
    for idx, (sp, cp) in enumerate(zip(s_paras, c_paras)):
        if not sp.strip() and not cp.strip():
            continue
        if not sp.strip() or not cp.strip():
            out.append(
                AlignedUnit(sp.strip(), cp.strip(), idx, 0.0, "para")
            )
            continue
        out.extend(align_paragraph_sentences(sp, cp, paragraph_index=idx))
    return out