Build a legal document conflict detector

git clone --depth 1 https://github.com/isaacus-dev/cookbooks.git
cd cookbooks/cookbooks/conflict-detector
uv venv
# Install required packages for the conflict detector.
uv pip install ftfy isaacus msgspec numpy python-dotenv text2markdown flask
# Only needed for generation (section C) with an API model: install the one you use.
# uv pip install anthropic openai google-genai

import re

from functools import cache

import ftfy

from text2markdown import text2markdown

NON_STANDARD_CONTROL_CHARS_RE = re.compile(r"[" + re.escape("\a\b\f\r\v") + "]")
START_OF_TEXT_THAT_IS_ONLY_WHITESPACE_FOLLOWED_BY_NEWLINE_RE = re.compile(r"^\s*\n")
END_OF_TEXT_THAT_IS_ONLY_WHITESPACE_PRECEDED_BY_NEWLINE_RE = re.compile(r"\n\s*$")


def clean_text(text: str) -> str:
    text = ftfy.fix_text(text)
    text = text.replace("\xa0", " ")
    text = NON_STANDARD_CONTROL_CHARS_RE.sub("", text)
    text = "\n".join([re.sub(r"\s+$", "", line) for line in text.split("\n")])
    text = START_OF_TEXT_THAT_IS_ONLY_WHITESPACE_FOLLOWED_BY_NEWLINE_RE.sub("", text)

    return END_OF_TEXT_THAT_IS_ONLY_WHITESPACE_PRECEDED_BY_NEWLINE_RE.sub("", text)


# A document can be loaded more than once in a run; convert it only once
@cache
def to_markdown(text: str) -> str:
    # Skip cross-reference links: they don't help the models, and their HTML would stop ftfy decoding entities
    return text2markdown(text, link_xrefs=False)


def normalize_text(text: str, *, markdown: bool) -> str:
    if not markdown:
        text = to_markdown(text)

    return clean_text(text)

from pathlib import Path

path = Path("samples/03-common-paper-cloud-service-agreement.md")
raw = path.read_text(encoding="utf-8")
text = normalize_text(raw, markdown=True)

for line in [line for line in text.splitlines() if line][:6]:
    print(f"  {line[:100]}")

from pathlib import Path

import msgspec

from dotenv import load_dotenv

PROJECT_ROOT = Path.cwd()
SAMPLES_DIR = PROJECT_ROOT / "samples"

MARKDOWN_SUFFIXES = (".md", ".markdown")

ENRICHER_MODEL = "kanon-2-enricher"

load_dotenv(PROJECT_ROOT / ".env")


class Clause(msgspec.Struct):
    doc: str
    seg_id: str
    type: str | None
    text: str
    start: int | None = None
    end: int | None = None
    parent: str | None = None
    category: str | None = None
    heading: str | None = None
    order: int | None = None
    definition: bool = False


def load_text(path) -> str:
    path = Path(path)
    text = path.read_text(encoding="utf-8")

    return normalize_text(text, markdown=path.suffix.lower() in MARKDOWN_SUFFIXES)

import os

from isaacus import Isaacus

client = Isaacus(api_key=os.getenv("ISAACUS_API_KEY"))

import re


def enrich_document(client, text):
    return (
        client.enrichments.create(
            model=ENRICHER_MODEL,
            texts=text,
            overflow_strategy="auto",
        )
        .results[0]
        .document
    )


def document_sections(document):
    # The document's titled containers, saved with the index
    sections = []

    for seg in document.segments:
        if seg.kind != "container" or not seg.title:
            continue

        title = seg.title.decode(document.text).strip()

        if not title:
            continue

        sections.append(
            {
                "title": title,
                "level": seg.level,
                "start": seg.span.start,
                "end": seg.span.end,
            }
        )

    return sections


def own_title(document, seg):
    return seg.title.decode(document.text).strip() if seg.title else None


MARKDOWN_HEADING = re.compile(r"^#+\s+(.+)$", re.M)
HEADING_NUMBER = re.compile(r"^\d+(?:\\?\.\d+)*\\?\.?\s+")
DEFINITION_LEAD = re.compile(r"[\W\d_]*(?:(?:the|an?)\b[\W_]*)?", re.I)


def markdown_headings(text):
    found = []

    for match in MARKDOWN_HEADING.finditer(text):
        title = HEADING_NUMBER.sub("", match.group(1).replace("*", "")).strip()

        if title:
            found.append((match.start(), title))

    return found


def resolve_headings(document):
    markdown = markdown_headings(document.text)
    segments = {seg.id: seg for seg in document.segments}
    headings = {}

    for seg in document.segments:
        parent_id, heading = seg.parent, None

        while parent_id is not None and parent_id in segments:
            heading = own_title(document, segments[parent_id])

            if heading:
                break

            parent_id = segments[parent_id].parent

        # Title of the closest parent section that has one, else its own title, else the nearest markdown heading above
        headings[seg.id] = (
            heading
            or own_title(document, seg)
            or next((title for start, title in reversed(markdown) if start <= seg.span.start), None)
        )

    return headings

def opening(text):
    return text[DEFINITION_LEAD.match(text).end() :]


def defines_term(document, *, start, end):
    clause_opening = opening(document.text[start:end])

    # A definition opens with a term the enricher found, and that term is defined inside this clause
    return any(
        clause_opening.startswith(opening(term.name.decode(document.text)))
        and any(start <= pos < end for pos in (term.name.start, term.meaning.start))
        for term in document.terms
    )


def has_body(document, seg, *, start, end):
    # Spans that only label the segment: its type name ("Section"), code ("4.2") and title
    labels = [(span.start, span.end) for span in (seg.type_name, seg.code, seg.title) if span]

    # Markdown heading lines count as labels too
    pos = start

    for line in document.text[start:end].split("\n"):
        if line.lstrip().startswith("#"):
            labels.append((pos, pos + len(line)))

        pos += len(line) + 1

    # It has a body if any letter or digit sits outside every label
    return any(
        ch.isalnum() for n, ch in enumerate(document.text[start:end], start) if not any(a <= n < b for a, b in labels)
    )


def extract_clauses(doc_name, document):
    headings = resolve_headings(document)
    spans = {seg.id: seg.span for seg in document.segments}
    clauses = []

    for seg in document.segments:
        pieces, cut = [], seg.span.start

        # Cut the segment around its children's spans; each word belongs to one clause
        for child in sorted((spans[c] for c in seg.children), key=lambda span: span.start):
            pieces.append((cut, child.start))
            cut = max(cut, child.end)

        pieces.append((cut, seg.span.end))

        for n, (start, end) in enumerate(pieces):
            raw = document.text[start:end]
            clause_text = raw.strip()

            if not has_body(document, seg, start=start, end=end):
                continue

            # Offsets of the text once leading and trailing whitespace is trimmed
            clause_start, clause_end = start + len(raw) - len(raw.lstrip()), start + len(raw.rstrip())
            clauses.append(
                Clause(
                    doc=doc_name,
                    seg_id=seg.id if n == 0 else f"{seg.id}.{n}",
                    type=seg.type,
                    text=clause_text,
                    start=clause_start,
                    end=clause_end,
                    parent=seg.parent,
                    category=seg.category,
                    heading=headings.get(seg.id),
                    definition=defines_term(document, start=clause_start, end=clause_end),
                )
            )

    clauses.sort(key=lambda clause: clause.start)

    for order, clause in enumerate(clauses):
        clause.order = order

    return clauses

path = SAMPLES_DIR / "03-common-paper-cloud-service-agreement.md"
document = enrich_document(client, load_text(path))
clauses = extract_clauses(path.name, document)

print(f"{len(document.segments)} segments -> {len(clauses)} clauses")

for clause in clauses[:5]:
    print(f"  {clause.seg_id:<8} [{clause.heading}] {clause.text[:80]}")

import numpy as np

EMBEDDER_MODEL = "kanon-2-embedder"
EMBED_BATCH_SIZE = 96


def unit_normalize(rows) -> np.ndarray:
    vectors = np.array(rows, dtype=np.float32) if rows else np.zeros((0, 0), np.float32)

    if vectors.size == 0:
        return vectors

    norms = np.linalg.norm(vectors, axis=1, keepdims=True)

    # Divide each vector by its length; an all-zero vector stays zero rather than dividing by zero
    return (vectors / np.where(norms == 0, 1, norms)).astype(np.float32)


def embed_propositions(client, texts, *, task="retrieval/document"):
    rows = []

    for i in range(0, len(texts), EMBED_BATCH_SIZE):
        resp = client.embeddings.create(
            model=EMBEDDER_MODEL,
            texts=texts[i : i + EMBED_BATCH_SIZE],
            task=task,
        )
        rows.extend(e.embedding for e in resp.embeddings)

    return unit_normalize(rows)

path = SAMPLES_DIR / "03-common-paper-cloud-service-agreement.md"
document = enrich_document(client, load_text(path))
clauses = extract_clauses(path.name, document)
vectors = embed_propositions(client, [c.text for c in clauses])

print(f"{len(clauses)} clauses -> {len(vectors)} vectors")

n = next(n for n, clause in enumerate(clauses) if clause.heading == "Payment")
scores = vectors @ vectors[n]
print(f"closest to [{clauses[n].heading}] {clauses[n].text[:60]}")

for m in np.argsort(-scores)[:4]:
    print(f"  {scores[m]:.2f}  [{clauses[m].heading}] {clauses[m].text[:60]}")

import json
import hashlib

RESULTS_DIR = PROJECT_ROOT / "results"
DOC_CACHE_DIR = RESULTS_DIR / "cache"
DOCUMENTS_DIR = RESULTS_DIR / "documents"
DOCUMENT_SCHEMA = "conflict-document/1"

PIPELINE_VERSION = 19


class DocumentIndex(msgspec.Struct):
    # One document's clauses (called propositions here), their embeddings and its cleaned text
    propositions: list
    embeddings: np.ndarray
    manifest: dict | None = None
    text: str | None = None
    sections: list | None = None


def read_index(index_dir: Path) -> DocumentIndex | None:
    index_dir = Path(index_dir)

    try:
        rows = json.loads((index_dir / "propositions.json").read_text(encoding="utf-8"))
        matrix = np.load(index_dir / "embeddings.npy")
        manifest = json.loads((index_dir / "manifest.json").read_text(encoding="utf-8"))
        text = (index_dir / "text.txt").read_text(encoding="utf-8")
        sections = json.loads((index_dir / "sections.json").read_text(encoding="utf-8"))

    except (OSError, ValueError):
        return None

    if matrix.shape[0] != len(rows):
        return None

    return DocumentIndex(
        propositions=[Clause(**{f.name: r.get(f.name) for f in msgspec.structs.fields(Clause)}) for r in rows],
        embeddings=matrix,
        manifest=manifest,
        text=text,
        sections=sections,
    )


def write_index(index, index_dir):
    index_dir = Path(index_dir)
    index_dir.mkdir(parents=True, exist_ok=True)
    (index_dir / "manifest.json").write_text(json.dumps(index.manifest, indent=2), encoding="utf-8")
    (index_dir / "propositions.json").write_text(
        json.dumps([msgspec.structs.asdict(c) for c in index.propositions], indent=2),
        encoding="utf-8",
    )
    np.save(index_dir / "embeddings.npy", index.embeddings)
    (index_dir / "text.txt").write_text(index.text, encoding="utf-8")
    (index_dir / "sections.json").write_text(
        json.dumps(index.sections or [], indent=2),
        encoding="utf-8",
    )

    return index_dir


def document_identity(path, index, *, text=None) -> dict:
    path = Path(path)

    if text is None:
        text = index.text if index.text is not None else load_text(path)

    try:
        source = path.resolve().relative_to(PROJECT_ROOT).as_posix()

    except ValueError:
        source = str(path)

    return {
        "source": source,
        "source_sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
        "text_sha256": hashlib.sha256(text.encode("utf-8")).hexdigest(),
        "sections": index.sections or [],
    }


def document_artifact(path, index) -> dict:
    path = Path(path)
    text = load_text(path)
    assert index.text is None or index.text == text, f"{path.name}: clear results/cache"
    text = index.text if index.text is not None else text

    return {"schema": DOCUMENT_SCHEMA, **document_identity(path, index, text=text), "text": text}


def write_document_artifact(path, index, *, documents_dir=DOCUMENTS_DIR) -> Path:
    # The report only stores offsets, so the demo reads each document's text from here
    payload = document_artifact(path, index)
    dest = Path(documents_dir) / f"{payload['source_sha256']}.json"
    dest.parent.mkdir(parents=True, exist_ok=True)
    dest.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")

    return dest


def document_fingerprint(normalized_text):
    # Any change to the text, the pipeline version or a model re-indexes; delete results/cache/ to force it
    stamp = f"{PIPELINE_VERSION}|{ENRICHER_MODEL}|{EMBEDDER_MODEL}|{normalized_text}"

    return hashlib.sha256(stamp.encode("utf-8")).hexdigest()[:16]


def document_manifest(doc_name, path, entry):
    return {"doc_name": doc_name, "source_path": str(path), "files": {doc_name: entry}}


def index_document(client, path, *, use_cache=True) -> DocumentIndex:
    path = Path(path)
    text = load_text(path)
    cache_dir = DOC_CACHE_DIR / document_fingerprint(text)

    if use_cache:
        cached = rekey(read_index(cache_dir), path.name, path)

        if cached is not None:
            return cached

    document = enrich_document(client, text)
    propositions = extract_clauses(path.name, document)

    index = DocumentIndex(
        propositions=propositions,
        embeddings=embed_propositions(client, [p.text for p in propositions]),
        text=document.text,
        sections=document_sections(document),
        manifest=document_manifest(
            path.name, path, {"proposition_count": len(propositions), "clause_count": len(propositions)}
        ),
    )
    write_index(index, cache_dir)

    return index


def rekey(index, doc_name, path):
    if index is None or index.text is None:
        return None

    # A cached index may have been saved under another file name
    for clause in index.propositions:
        clause.doc = doc_name

    entry = next(iter((index.manifest or {}).get("files", {}).values()), {})
    index.manifest = document_manifest(doc_name, path, entry)

    return index

from time import perf_counter

path = SAMPLES_DIR / "03-common-paper-cloud-service-agreement.md"

for attempt in ("cold", "warm"):
    started = perf_counter()
    index = index_document(client, path)
    seconds = perf_counter() - started

    print(f"{attempt:<5} {seconds:6.2f}s  {len(index.propositions)} propositions")

reference = index_document(client, SAMPLES_DIR / "03-common-paper-cloud-service-agreement.md")
input_doc = index_document(client, SAMPLES_DIR / "04-mistral-ai-terms-of-use.md")

for name, index in (("reference", reference), ("input", input_doc)):
    print(f"{name:<10} {len(index.propositions):>4} propositions")

from functools import partial
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor

import numpy as np

CLASSIFIER_MODEL = "kanon-universal-classifier"

SHORTLIST_MIN_RELEVANCE = 0.5
SHORTLIST_TOP_K = 3

STRUCTURAL_CATEGORIES = ("front_matter", "back_matter", "annotation")

FILTER_QUERY = "{This is a footnote, citation, editorial note, heading or signature, not a contractual provision.}"
FILTER_THRESHOLD = 0.5


def pair_matrix(input_doc: DocumentIndex, reference: DocumentIndex) -> np.ndarray:
    if input_doc.embeddings.size == 0 or reference.embeddings.size == 0:
        return np.zeros((len(input_doc.propositions), len(reference.propositions)), np.float32)

    return input_doc.embeddings @ reference.embeddings.T


def non_operative(client, index: DocumentIndex) -> np.ndarray:
    scores = classify(client, FILTER_QUERY, [flatten(c.text) for c in index.propositions], is_iql=True)

    return np.array(
        [
            c.category in STRUCTURAL_CATEGORIES or (not c.definition and score >= FILTER_THRESHOLD)
            for c, score in zip(index.propositions, scores, strict=True)
        ],
        dtype=bool,
    )


def mask_non_operative(sim: np.ndarray, inc: np.ndarray, ref: np.ndarray) -> np.ndarray:
    masked = sim.copy()

    # A filtered clause scores -1 with everything, so it can't be shortlisted
    masked[inc[:, None] | ref[None, :]] = -1.0

    return masked


def shortlist_pairs(sim: np.ndarray, *, min_relevance: float, top_k: int) -> dict[int, list[tuple[int, float]]]:
    if sim.size == 0:
        return {}

    # Rank of each pair among its input clause's matches, and among its reference clause's (0 = closest)
    row_rank = np.argsort(np.argsort(-sim, axis=1, kind="stable"), axis=1)
    col_rank = np.argsort(np.argsort(-sim, axis=0, kind="stable"), axis=0)

    # Keep a pair that's similar enough and in the top k for either clause
    keep = (sim >= min_relevance) & ((row_rank < top_k) | (col_rank < top_k))

    shortlist: dict[int, list[tuple[int, float]]] = defaultdict(list)

    for i, j in zip(*np.nonzero(keep), strict=True):
        shortlist[int(i)].append((int(j), float(sim[i, j])))

    for row in shortlist.values():
        row.sort(key=lambda entry: -entry[1])

    return dict(shortlist)


def flatten(text: str) -> str:
    return " ".join(text.split())


def span(clause: Clause) -> dict | None:
    if clause.start is None or clause.end is None:
        return None

    return {"start": clause.start, "end": clause.end}


def clause_ref(clause: Clause) -> dict:
    return {"seg_id": clause.seg_id, "text": flatten(clause.text), "span": span(clause)}


CLASSIFY_MAX_PARALLEL_REQUESTS = 6
CLASSIFY_BATCH_SIZE = 50


def batched(score, texts: list[str]) -> list[float]:
    batches = [texts[n : n + CLASSIFY_BATCH_SIZE] for n in range(0, len(texts), CLASSIFY_BATCH_SIZE)]

    with ThreadPoolExecutor(max_workers=CLASSIFY_MAX_PARALLEL_REQUESTS) as pool:
        return [s for scores in pool.map(score, batches) for s in scores]


def classify_batch(client, batch: list[str], *, query: str, is_iql: bool) -> list[float]:
    resp = client.classifications.universal.create(
        model=CLASSIFIER_MODEL, query=query, texts=batch, is_iql=is_iql, scoring_method="auto"
    )

    return [c.score for c in sorted(resp.classifications, key=lambda c: c.index)]


def classify(client, query: str, texts: list[str], *, is_iql: bool = False) -> list[float]:
    return batched(partial(classify_batch, client, query=query, is_iql=is_iql), texts)

reference = index_document(client, SAMPLES_DIR / "03-common-paper-cloud-service-agreement.md")
input_doc = index_document(client, SAMPLES_DIR / "04-mistral-ai-terms-of-use.md")
input_filtered, reference_filtered = non_operative(client, input_doc), non_operative(client, reference)
sim = pair_matrix(input_doc, reference)

shortlist = shortlist_pairs(
    mask_non_operative(sim, input_filtered, reference_filtered),
    min_relevance=SHORTLIST_MIN_RELEVANCE,
    top_k=SHORTLIST_TOP_K,
)

pairs = sum(len(row) for row in shortlist.values())

for name, filtered in (("input", input_filtered), ("reference", reference_filtered)):
    print(f"{name:<10} {filtered.sum():>3} of {len(filtered)} clauses filtered")

print(f"candidate pairs: {pairs}")

import math

import msgspec

RERANKER_MODEL = "kanon-2-reranker"

EMPTY_QUERY = "routine boilerplate, a recital, a preamble, a heading, a signature block, an address or a definition"
CARE_STATEMENT = "A party would not care whether this clause or the following clause applied: {clause}"
AGREE_QUERY = "A lawyer would compare this clause with the following clause: {clause}"
MONEY_STATEMENT = "This states an amount of money, a percentage or a cap."
CLASSIFY_QUERY_MAX_CHARS = 2000
MATCH_TEMPERATURE = 0.06
SAME_SLOPE = 2.049
SAME_BIAS = 3.058

EMPTY_WEIGHT = -4.0
CARE_WEIGHT = -1.0
AGREE_WEIGHT = -1.0
MONEY_WEIGHT = 1.0
DIFFERS_SLOPE = 0.526
DIFFERS_BIAS = -12.312
CONFLICT_SAME_SLOPE = 0.576
CONFLICT_DIFFERS_SLOPE = 0.707
CONFLICT_BIAS = 1.663


def rerank_batch(client, batch: list[str], *, query: str) -> list[float]:
    resp = client.rerankings.create(model=RERANKER_MODEL, query=query, texts=batch)

    return [r.score for r in sorted(resp.results, key=lambda r: r.index)]


def rerank(client, query: str, texts: list[str]) -> list[float]:
    return batched(partial(rerank_batch, client, query=query), texts)


def logit(p: float) -> float:
    p = min(max(p, 1e-6), 1 - 1e-6)

    return math.log(p / (1 - p))


def sigmoid(x: float) -> float:
    return 1 / (1 + math.exp(-x))


def log_sigmoid(x: float) -> float:
    # log(sigmoid(x)), without rounding tiny probabilities to zero
    return -float(np.logaddexp(0.0, -x))


def mean(scores: dict, a, b) -> float:
    return (scores[a] + scores[b]) / 2


def score_query(b: str, *, texts: list[str], score, template: str, max_chars: int | None) -> dict:
    # Ask the model about clause b once, against every clause it's paired with
    return {(a, b): s for a, s in zip(texts, score(template.format(clause=b[:max_chars]), texts), strict=True)}


def pairwise(score, template: str, pairs: list[tuple[str, str]], *, max_chars: int | None = None) -> dict:
    by_query: dict[str, list[str]] = defaultdict(list)

    # Group pairs by their second clause, so each clause goes to the model once
    for a, b in dict.fromkeys(pairs):
        by_query[b].append(a)

    out = {}

    with ThreadPoolExecutor(max_workers=CLASSIFY_MAX_PARALLEL_REQUESTS) as pool:
        jobs = [
            pool.submit(score_query, b, texts=texts, score=score, template=template, max_chars=max_chars)
            for b, texts in by_query.items()
        ]

        for job in jobs:
            out.update(job.result())

    return out


class PairScore(msgspec.Struct):
    relevance: float
    match: float
    empty: float
    care: float
    agree: float
    money: float

    @property
    def same_matter(self) -> float:
        return sigmoid(SAME_SLOPE * self.match + SAME_BIAS)

    @property
    def materiality(self) -> float:
        return (
            EMPTY_WEIGHT * logit(self.empty)
            + CARE_WEIGHT * logit(self.care)
            + AGREE_WEIGHT * logit(self.agree)
            + MONEY_WEIGHT * logit(self.money)
        )

    @property
    def differs(self) -> float:
        return sigmoid(DIFFERS_SLOPE * self.materiality + DIFFERS_BIAS)

    @property
    def conflict(self) -> float:
        return sigmoid(
            CONFLICT_SAME_SLOPE * log_sigmoid(SAME_SLOPE * self.match + SAME_BIAS)
            + CONFLICT_DIFFERS_SLOPE * log_sigmoid(DIFFERS_SLOPE * self.materiality + DIFFERS_BIAS)
            + CONFLICT_BIAS
        )


def match_matrix(sim: np.ndarray) -> np.ndarray:
    logits = sim / MATCH_TEMPERATURE
    row = logits - np.logaddexp.reduce(logits, axis=1, keepdims=True)
    col = logits - np.logaddexp.reduce(logits, axis=0, keepdims=True)

    return (row + col) / 2


def score_pairs(
    client,
    input_doc: DocumentIndex,
    reference: DocumentIndex,
    *,
    shortlist: dict[int, list[tuple[int, float]]],
) -> dict[tuple[int, int], PairScore]:
    match = match_matrix(pair_matrix(input_doc, reference))
    texts = {
        (i, j): (flatten(input_doc.propositions[i].text), flatten(reference.propositions[j].text))
        for i, row in shortlist.items()
        for j, _ in row
    }
    unique = list(dict.fromkeys(t for pair in texts.values() for t in pair))
    both_ways = [p for a, b in texts.values() for p in ((a, b), (b, a))]

    empty = dict(zip(unique, rerank(client, EMPTY_QUERY, unique), strict=True))
    money = dict(zip(unique, classify(client, MONEY_STATEMENT, unique), strict=True))
    care = pairwise(partial(classify, client), CARE_STATEMENT, both_ways, max_chars=CLASSIFY_QUERY_MAX_CHARS)
    agree = pairwise(partial(rerank, client), AGREE_QUERY, both_ways)

    return {
        (i, j): PairScore(
            relevance=relevance,
            match=float(match[i, j]),
            empty=mean(empty, *texts[(i, j)]),
            money=mean(money, *texts[(i, j)]),
            care=mean(care, texts[(i, j)], texts[(i, j)][::-1]),
            agree=mean(agree, texts[(i, j)], texts[(i, j)][::-1]),
        )
        for i, candidates in shortlist.items()
        for j, relevance in candidates
    }

SAME_THRESHOLD = 0.3

CONFLICT_THRESHOLD = 0.728


def status(pairs: list[tuple[int, PairScore]], *, deviating: bool) -> dict:
    best = max(pairs, key=lambda e: e[1].relevance, default=None)

    return {
        "status": "deviation" if deviating else "covered" if pairs else "gap",
        "best_relevance": best[1].relevance if best else None,
        "nearest_index": best[0] if best else None,
    }


def side_status(index: DocumentIndex, pairs, deviating: set[int]) -> dict[int, dict]:
    return {n: status(pairs.get(n, []), deviating=n in deviating) for n in range(len(index.propositions))}


def partition(input_doc: DocumentIndex, reference: DocumentIndex, scored: dict[tuple[int, int], PairScore]):
    closest: dict[int, dict[int, tuple[tuple[int, int], PairScore]]] = {0: {}, 1: {}}

    # For each clause on each side, keep its best match that's about the same matter
    for pair, sc in scored.items():
        for side, best in closest.items():
            if sc.same_matter >= SAME_THRESHOLD and (pair[side] not in best or sc.match > best[pair[side]][1].match):
                best[pair[side]] = (pair, sc)

    # Both sides' picks together; a pair both clauses picked appears once
    selected = dict(entry for best in closest.values() for entry in best.values())

    # A picked pair is a deviation when its P(conflict) clears the threshold
    deviating_pairs = [(i, j, sc) for (i, j), sc in selected.items() if sc.conflict >= CONFLICT_THRESHOLD]

    inc_pairs: dict[int, list[tuple[int, PairScore]]] = defaultdict(list)
    ref_pairs: dict[int, list[tuple[int, PairScore]]] = defaultdict(list)

    for (i, j), score in scored.items():
        if score.same_matter >= SAME_THRESHOLD:
            inc_pairs[i].append((j, score))
            ref_pairs[j].append((i, score))

    return (
        deviating_pairs,
        side_status(input_doc, inc_pairs, {i for i, _, _ in deviating_pairs}),
        side_status(reference, ref_pairs, {j for _, j, _ in deviating_pairs}),
    )


def clause_row(clause: Clause) -> dict:
    return {
        "id": None,
        "seg_id": clause.seg_id,
        "heading": clause.heading,
        "text": flatten(clause.text),
        "span": span(clause),
        "order": clause.order,
        "category": clause.category,
    }


def numbered(rows: list[dict], sort_key) -> list[dict]:
    # Number rows in document order, so the report reads top to bottom like the documents
    rows.sort(key=sort_key)

    for n, row in enumerate(rows, start=1):
        row["id"] = n

    return rows


def roll_up_deviations(deviating_pairs, input_doc: DocumentIndex, reference: DocumentIndex) -> list[dict]:
    by_input: dict[int, list[tuple[int, PairScore]]] = defaultdict(list)

    for i, j, score in deviating_pairs:
        by_input[i].append((j, score))

    rows = []

    for i, pairs in by_input.items():
        pairs.sort(key=lambda e: reference.propositions[e[0]].order or 0)
        rows.append(
            {
                **clause_row(input_doc.propositions[i]),
                "relevance": pairs[0][1].relevance,
                "match_count": len(pairs),
                "matches": [
                    {
                        "input": clause_ref(input_doc.propositions[i]),
                        "reference": {
                            **clause_ref(reference.propositions[j]),
                            "clause_text": flatten(reference.propositions[j].text),
                        },
                        "relevance": score.relevance,
                    }
                    for j, score in pairs
                ],
            }
        )

    return numbered(rows, lambda r: r["order"] or 0)


def roll_up_gaps(
    index: DocumentIndex, status: dict[int, dict], *, other: DocumentIndex, filtered: np.ndarray
) -> list[dict]:
    rows = []

    for i, clause in enumerate(index.propositions):
        info = status[i]

        if filtered[i] or info["status"] != "gap" or clause.definition:
            continue

        near = info["nearest_index"]
        rows.append(
            {
                **clause_row(clause),
                "propositions": {"total": 1, "matched_count": 0, "uncovered_count": 1},
                "gap_score": 1.0 - (info["best_relevance"] or 0.0),
                "uncovered_propositions": [
                    {
                        **clause_ref(clause),
                        **{k: info[k] for k in ("status", "best_relevance")},
                        "nearest": clause_ref(other.propositions[near]) if near is not None else None,
                    }
                ],
            }
        )

    return numbered(rows, lambda r: r["order"] or 0)

import re
import json
import shutil
import hashlib

from pathlib import Path
from datetime import UTC, datetime

REPORTS_DIR = RESULTS_DIR / "reports"

SCHEMA = "doc-compare/2"


def coverage_counts(status: dict[int, dict], filtered: np.ndarray) -> dict:
    counts = {"total": len(status), "covered": 0, "deviation": 0, "gap": 0, "filtered": int(filtered.sum())}

    for n, row in status.items():
        if not filtered[n]:
            counts[row["status"]] += 1

    return counts


def document_summary(index: DocumentIndex, path: Path) -> dict:
    return {
        "name": Path(path).name,
        **document_identity(path, index),
        "propositions": len(index.propositions),
    }


def scored_pairs_rows(input_doc, reference, scored, *, deviating) -> list[dict]:
    rows = []

    for (i, j), score in scored.items():
        inc, ref = input_doc.propositions[i], reference.propositions[j]
        rows.append(
            {
                "input_seg": inc.seg_id,
                "input": flatten(inc.text),
                "reference_seg": ref.seg_id,
                "reference": flatten(ref.text),
                "relevance": score.relevance,
                "same_matter": score.same_matter,
                "differs": score.differs,
                "conflict": score.conflict,
                "deviating": (i, j) in deviating,
            }
        )

    rows.sort(key=lambda row: -row["relevance"])

    return rows


def build_report(
    *,
    reference,
    input_doc,
    reference_path,
    input_path,
    deviation_rows,
    gaps_reference,
    gaps_input,
    input_status,
    reference_status,
    input_filtered,
    reference_filtered,
    pairs,
) -> dict:
    return {
        "schema": SCHEMA,
        "reference": document_summary(reference, reference_path),
        "input": document_summary(input_doc, input_path),
        "summary": {
            "deviation_rows": len(deviation_rows),
            "deviations": sum(row["match_count"] for row in deviation_rows),
            "shortlisted_pairs": len(pairs),
            "gaps_reference": len(gaps_reference),
            "gaps_input": len(gaps_input),
            "reference_coverage": coverage_counts(reference_status, reference_filtered),
            "input_coverage": coverage_counts(input_status, input_filtered),
        },
        "summaries": None,
        "recommendations": None,
        "deviations": deviation_rows,
        "gaps_reference": gaps_reference,
        "gaps_input": gaps_input,
        "pairs": pairs,
    }


def slug(name: str) -> str:
    return re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-")[:60] or "document"


def pair_slug(reference_path, input_path) -> str:
    return f"{slug(Path(reference_path).stem)}__{slug(Path(input_path).stem)}"


def analysis_params() -> dict:
    return {
        "schema": SCHEMA,
        "pipeline_version": PIPELINE_VERSION,
        "enricher_model": ENRICHER_MODEL,
        "embedder_model": EMBEDDER_MODEL,
        "classifier_model": CLASSIFIER_MODEL,
        "shortlist_min_relevance": SHORTLIST_MIN_RELEVANCE,
        "shortlist_top_k": SHORTLIST_TOP_K,
        "filter_query": FILTER_QUERY,
        "filter_threshold": FILTER_THRESHOLD,
        "reranker_model": RERANKER_MODEL,
        "empty_query": EMPTY_QUERY,
        "care_statement": CARE_STATEMENT,
        "agree_query": AGREE_QUERY,
        "money_statement": MONEY_STATEMENT,
        "classify_query_max_chars": CLASSIFY_QUERY_MAX_CHARS,
        "match_temperature": MATCH_TEMPERATURE,
        "same_slope": SAME_SLOPE,
        "same_bias": SAME_BIAS,
        "same_threshold": SAME_THRESHOLD,
        "materiality_weights": [EMPTY_WEIGHT, CARE_WEIGHT, AGREE_WEIGHT, MONEY_WEIGHT],
        "differs_slope": DIFFERS_SLOPE,
        "differs_bias": DIFFERS_BIAS,
        "conflict_slopes": [CONFLICT_SAME_SLOPE, CONFLICT_DIFFERS_SLOPE],
        "conflict_bias": CONFLICT_BIAS,
        "conflict_threshold": CONFLICT_THRESHOLD,
        "structural_categories": list(STRUCTURAL_CATEGORIES),
    }


def comparison_dir_for(reference_path, input_path, *, runs_dir=REPORTS_DIR) -> Path:
    return Path(runs_dir) / pair_slug(reference_path, input_path)


def read_json(path):
    try:
        return json.loads(Path(path).read_text(encoding="utf-8"))

    except (OSError, ValueError):
        return None


def saved_report_if_current(comparison_dir, reference_path, input_path) -> dict | None:
    report = read_json(Path(comparison_dir) / "report.json")
    saved = read_json(Path(comparison_dir) / "params.json")

    if report is None or saved is None:
        return None

    current = analysis_params()

    if {k: saved.get(k) for k in current} != current:
        return None

    for side, path in (("reference", reference_path), ("input", input_path)):
        text = load_text(path).encode("utf-8")

        if report.get(side, {}).get("text_sha256") != hashlib.sha256(text).hexdigest():
            return None

    return report


def write_analysis_run(
    report,
    *,
    reference: DocumentIndex,
    input_doc: DocumentIndex,
    reference_path: Path,
    input_path: Path,
    run_dir: Path,
    label=None,
    documents_dir=DOCUMENTS_DIR,
) -> Path:
    dest = Path(run_dir)

    if dest.exists():
        shutil.rmtree(dest)

    dest.mkdir(parents=True)

    params = {
        **analysis_params(),
        "slug": dest.name,
        "label": label,
        "created_at": datetime.now(UTC).isoformat().replace("+00:00", "Z"),
    }

    for name, payload in (
        ("report", {k: v for k, v in report.items() if k != "pairs"}),
        ("pairs", report.get("pairs") or []),
        ("params", params),
    ):
        (dest / f"{name}.json").write_text(json.dumps(payload, indent=2), encoding="utf-8")

    write_document_artifact(reference_path, reference, documents_dir=documents_dir)
    write_document_artifact(input_path, input_doc, documents_dir=documents_dir)

    return dest

def run_comparison(client, reference_path, input_path, *, use_cache: bool = True, label: str | None = None) -> dict:
    reference_path, input_path = Path(reference_path), Path(input_path)
    dest = comparison_dir_for(reference_path, input_path)

    if use_cache and (report := saved_report_if_current(dest, reference_path, input_path)) is not None:
        return report

    print("STAGE: Reading documents")
    reference = index_document(client, reference_path, use_cache=use_cache)
    input_doc = index_document(client, input_path, use_cache=use_cache)

    print("STAGE: Filtering")
    input_filtered, reference_filtered = non_operative(client, input_doc), non_operative(client, reference)

    print("STAGE: Matching")
    shortlist = shortlist_pairs(
        mask_non_operative(pair_matrix(input_doc, reference), input_filtered, reference_filtered),
        min_relevance=SHORTLIST_MIN_RELEVANCE,
        top_k=SHORTLIST_TOP_K,
    )

    print("STAGE: Classifying")
    scored = score_pairs(client, input_doc, reference, shortlist=shortlist)

    print("STAGE: Comparing")
    deviating_pairs, input_status, reference_status = partition(input_doc, reference, scored)

    deviation_rows = roll_up_deviations(deviating_pairs, input_doc, reference)
    gaps_reference = roll_up_gaps(reference, reference_status, other=input_doc, filtered=reference_filtered)
    gaps_input = roll_up_gaps(input_doc, input_status, other=reference, filtered=input_filtered)

    print("STAGE: Assembling report")
    report = build_report(
        reference=reference,
        input_doc=input_doc,
        reference_path=reference_path,
        input_path=input_path,
        deviation_rows=deviation_rows,
        gaps_reference=gaps_reference,
        gaps_input=gaps_input,
        input_status=input_status,
        reference_status=reference_status,
        input_filtered=input_filtered,
        reference_filtered=reference_filtered,
        pairs=scored_pairs_rows(input_doc, reference, scored, deviating={(i, j) for i, j, _ in deviating_pairs}),
    )

    print("STAGE: Saving results")
    write_analysis_run(
        report,
        reference=reference,
        input_doc=input_doc,
        reference_path=reference_path,
        input_path=input_path,
        run_dir=dest,
        label=label,
    )

    return report

report = run_comparison(
    client,
    SAMPLES_DIR / "03-common-paper-cloud-service-agreement.md",
    SAMPLES_DIR / "04-mistral-ai-terms-of-use.md",
)

import sys

sys.path.insert(0, str(PROJECT_ROOT))

import generation

report = generation.generate_run(
    comparison_dir_for(
        SAMPLES_DIR / "03-common-paper-cloud-service-agreement.md",
        SAMPLES_DIR / "04-mistral-ai-terms-of-use.md",
    ),
    model="claude-cli-opus",
    api_key=None,
)