viewer.py

#!/usr/bin/env python3
"""Build a standalone HTML viewer for one extraction record.

    python3 viewer.py xevP8UDRAVh9

Reads the gold extraction and the source text it was extracted from, flattens the
record into entities/fields/evidence spans, and inlines the result into
`viewer/template.html`. The output has no external requests, so it can be mailed
or dropped on a static host as a single file.

See docs/evidence-viewer.md for what the two views are for.
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent
TEMPLATE = ROOT / "viewer" / "template.html"

# A colour family is one hue; every entity inside it gets its own step of that hue,
# so "which kind of thing" and "which one of them" are both readable at a glance.
# Hues and their order come from the data-viz reference palette (slots 1-8).
FAMILIES = [
    ("analysis", "Analyses", "#2a78d6", "#3987e5"),
    ("study", "Study & design", "#eb6834", "#d95926"),
    ("sample", "Sample", "#1baf7a", "#199e70"),
    ("table", "Tables", "#eda100", "#c98500"),
    ("acquisition", "Acquisition", "#e87ba4", "#d55181"),
    ("preprocessing", "Preprocessing", "#008300", "#008300"),
    ("model", "Models", "#4a3aa7", "#9085e9"),
    ("measure", "Measures, regions, inference", "#e34948", "#e66767"),
    # The palette's eight slots are spoken for; a container the schema grows next folds
    # into this neutral rather than being cycled back onto a hue that means something.
    ("other", "Other", "#6d6c66", "#9a988e"),
]

# Container key -> (singular label, colour family). Order fixes both the legend and
# the order instances take colour steps in.
TYPES = [
    ("study", "Record", "study"),
    ("design", "Design", "study"),
    ("arms", "Arm", "study"),
    ("timepoints", "Timepoint", "study"),
    ("groups", "Group", "sample"),
    ("assessments", "Assessment", "sample"),
    ("analyses", "Analysis", "analysis"),
    ("acquisitions", "Acquisition", "acquisition"),
    ("devices", "Device", "acquisition"),
    ("preprocessings", "Preprocessing", "preprocessing"),
    ("model_estimations", "Model estimation", "model"),
    ("terms", "Term", "model"),
    ("regions", "Region", "measure"),
    ("measures", "Measure", "measure"),
    ("inference_settings", "Inference settings", "measure"),
    ("tables", "Table", "table"),
]
TYPE_LABEL = {k: label for k, label, _ in TYPES}
TYPE_FAMILY = {k: fam for k, _, fam in TYPES}
TYPE_ORDER = {k: i for i, (k, _, _) in enumerate(TYPES)}
UNKNOWN_ORDER = len(TYPES)

# Keys whose nested items carry their own local_id, and the type they map to.
NESTED_TYPES = {"arms": "arms", "timepoints": "timepoints", "terms": "terms"}


def is_field(obj) -> bool:
    """A leaf value slot: the extraction status is what distinguishes it from a container."""
    return isinstance(obj, dict) and "extraction_status" in obj


def humanize(key: str) -> str:
    return key.replace("_", " ").strip()


def short(text: str, limit: int = 90) -> str:
    text = re.sub(r"\s+", " ", text).strip()
    return text if len(text) <= limit else text[: limit - 1].rstrip() + "…"


class Build:
    def __init__(self, record: dict, text: str):
        self.record = record
        self.text = text
        self.ids = self._collect_ids(record)
        self.entities: dict[str, dict] = {}
        self.edges: list[dict] = []
        self.spans: dict[tuple[int, int], int] = {}
        self.span_list: list[dict] = []
        self.unknown_types: set[str] = set()
        self.misplaced_spans = 0

    @staticmethod
    def _collect_ids(obj, out=None) -> set[str]:
        out = set() if out is None else out
        if isinstance(obj, dict):
            if isinstance(obj.get("local_id"), str):
                out.add(obj["local_id"])
            for value in obj.values():
                Build._collect_ids(value, out)
        elif isinstance(obj, list):
            for value in obj:
                Build._collect_ids(value, out)
        return out

    # -- entities ---------------------------------------------------------

    def entity(self, local_id: str, type_key: str) -> dict:
        ent = self.entities.get(local_id)
        if ent is None:
            if type_key not in TYPE_LABEL:
                self.unknown_types.add(type_key)
            ent = {
                "id": local_id,
                "type": type_key,
                "typeLabel": TYPE_LABEL.get(type_key, humanize(type_key).rstrip("s").capitalize()),
                "family": TYPE_FAMILY.get(type_key, "other"),
                "label": humanize(local_id),
                "fields": [],
                "spans": [],
            }
            self.entities[local_id] = ent
        return ent

    def span_index(self, span: dict) -> int:
        key = (span["start_char"], span["end_char"])
        # The offsets are the contract: if one no longer addresses its own quote, the
        # record and the text file have come apart and every highlight after it is suspect.
        if self.text[key[0] : key[1]] != span.get("text"):
            self.misplaced_spans += 1
        idx = self.spans.get(key)
        if idx is None:
            idx = len(self.span_list)
            self.spans[key] = idx
            self.span_list.append({"start": key[0], "end": key[1]})
        return idx

    def add_edge(self, source: str, target: str, label: str, kind: str) -> None:
        if source == target or source not in self.ids or target not in self.ids:
            return
        for edge in self.edges:
            if edge["source"] == source and edge["target"] == target and edge["label"] == label:
                return
        self.edges.append({"source": source, "target": target, "label": label, "kind": kind})

    def add_field(self, ent: dict, path: list[str], obj: dict) -> None:
        status = obj.get("extraction_status")
        evidence = obj.get("evidence") or {}
        sets = []
        for group in evidence.get("sets") or []:
            indices = [self.span_index(span) for span in group.get("spans") or []]
            if indices:
                sets.append(indices)
        field = {
            "id": f"{ent['id']}::{'.'.join(path)}",
            "path": path,
            "label": " › ".join(humanize(part) for part in path),
            "status": status,
            "value": obj.get("value"),
            "valueSource": obj.get("value_source"),
            "sets": sets,
        }
        ent["fields"].append(field)
        for idx in {i for group in sets for i in group}:
            if idx not in ent["spans"]:
                ent["spans"].append(idx)

    def add_literal(self, ent: dict, path: list[str], value) -> None:
        """A plain scalar in the record (`details_type`, `acquisition_type`): no evidence to cite."""
        ent["fields"].append(
            {
                "id": f"{ent['id']}::{'.'.join(path)}",
                "path": path,
                "label": " › ".join(humanize(part) for part in path),
                "status": "literal",
                "value": value,
                "valueSource": None,
                "sets": [],
            }
        )

    # -- traversal --------------------------------------------------------

    def walk_entity(self, obj: dict, ent: dict, path: list[str]) -> None:
        for key, value in obj.items():
            if key == "local_id":
                continue
            self.walk_value(value, ent, path + [key])

    def walk_value(self, value, ent: dict, path: list[str]) -> None:
        key = path[-1]
        if is_field(value):
            self.add_field(ent, path, value)
        elif isinstance(value, dict):
            if isinstance(value.get("local_id"), str):
                self.child_entity(value, ent, key)
            else:
                self.walk_entity(value, ent, path)
        elif isinstance(value, list):
            for i, item in enumerate(value):
                if isinstance(item, str) and item in self.ids:
                    self.add_edge(ent["id"], item, humanize(key), "ref")
                elif isinstance(item, dict) and isinstance(item.get("local_id"), str):
                    self.child_entity(item, ent, key)
                elif isinstance(item, dict):
                    self.walk_entity(item, ent, path[:-1] + [self.item_label(item, key, i)])
        elif isinstance(value, str) and value in self.ids:
            self.add_edge(ent["id"], value, humanize(key), "ref")
        elif value is not None and not isinstance(value, (dict, list)):
            self.add_literal(ent, path, value)

    def item_label(self, item: dict, key: str, index: int) -> str:
        """Label a list entry by what it points at (`effect › term perfusion condition`)
        when it points at all: the position in the list means nothing to a reader."""
        for sub_key, sub_value in item.items():
            if isinstance(sub_value, str) and sub_value in self.ids:
                return humanize(sub_value)
        return f"{humanize(key)} {index + 1}"

    def child_entity(self, obj: dict, parent: dict, key: str) -> None:
        type_key = NESTED_TYPES.get(key, key)
        child = self.entity(obj["local_id"], type_key)
        self.add_edge(parent["id"], child["id"], humanize(key), "contains")
        self.walk_entity(obj, child, [])

    # -- assembly ---------------------------------------------------------

    def run(self) -> dict:
        record = self.record
        study = self.entity(record["local_id"], "study")
        for key, value in record.items():
            if key in ("local_id", "extraction_metadata"):
                continue
            if isinstance(value, list) and value and isinstance(value[0], dict) and "local_id" in value[0]:
                for item in value:
                    if isinstance(item, dict) and isinstance(item.get("local_id"), str):
                        child = self.entity(item["local_id"], key)
                        self.walk_entity(item, child, [])
            elif key == "design":
                design = self.entity("design", "design")
                self.add_edge(study["id"], "design", "design", "contains")
                self.ids.add("design")
                self.walk_entity(value, design, [])
            else:
                self.walk_value(value, study, [key])

        # An entity nothing points at still belongs to the record; hang it off the
        # study node so the graph stays connected and nothing is unreachable.
        linked = {e["source"] for e in self.edges} | {e["target"] for e in self.edges}
        for ent in self.entities.values():
            if ent["id"] != study["id"] and ent["id"] not in linked:
                self.add_edge(study["id"], ent["id"], ent["typeLabel"].lower(), "contains")

        self.name_entities()
        self.assign_colors()
        for span in self.span_list:
            span["text"] = self.text[span["start"] : span["end"]]

        order = sorted(
            self.entities.values(),
            key=lambda e: (TYPE_ORDER.get(e["type"], UNKNOWN_ORDER), e["colorIndex"]),
        )
        return {
            "record": {
                "id": record["local_id"],
                "title": self.text.split("\n", 1)[0].strip(),
                "sections": (record.get("extraction_metadata") or {}).get("paper_sections", []),
                "metadata": record.get("extraction_metadata") or {},
            },
            "text": self.text,
            "spans": self.span_list,
            "families": [
                {"key": k, "label": label, "light": light, "dark": dark, "count": self.family_counts.get(k, 0)}
                for k, label, light, dark in FAMILIES
            ],
            "entities": order,
            "edges": self.edges,
        }

    def name_entities(self) -> None:
        for ent in self.entities.values():
            name = self.field_value(ent, "name")
            description = self.field_value(ent, "description")
            if ent["type"] == "study":
                name, description = None, self.field_value(ent, "description")
            if ent["type"] == "tables":
                name = self.field_value(ent, "table_number")
                description = self.field_value(ent, "caption")
            if ent["type"] == "devices":
                name = " ".join(
                    v for v in (self.field_value(ent, "manufacturer"), self.field_value(ent, "model")) if v
                )
            if ent["type"] == "measures":
                name = self.field_value(ent, "source_label") or self.field_value(ent, "type")
                description = self.field_value(ent, "specific_metric")
            if ent["type"] == "acquisitions":
                name = self.field_value(ent, "modality") or self.field_value(ent, "acquisition_type")
            if ent["type"] == "model_estimations":
                name = self.field_value(ent, "model_type")
            if ent["type"] == "inference_settings":
                name = self.field_value(ent, "multiple_comparison_method")
            if ent["type"] == "analyses":
                description = self.field_value(ent, "definition") or description
            ent["label"] = short(name or humanize(ent["id"]), 60)
            ent["sublabel"] = short(description or "", 140)
            ent["fieldCount"] = len(ent["fields"])
            ent["evidenceCount"] = len(ent["spans"])
            ent["reportedCount"] = sum(1 for f in ent["fields"] if f["status"] == "extracted")
        self.disambiguate()

    def disambiguate(self) -> None:
        """Two analyses both called "Positive correlation" are useless as graph labels;
        tell them apart by whatever their local_ids do not share."""
        by_label: dict[str, list[dict]] = {}
        for ent in self.entities.values():
            by_label.setdefault(ent["label"].lower(), []).append(ent)
        for group in by_label.values():
            if len(group) < 2:
                continue
            token_sets = [set(e["id"].split("_")) for e in group]
            shared = set.intersection(*token_sets)
            for ent, tokens in zip(group, token_sets):
                unique = [t for t in ent["id"].split("_") if t not in shared]
                if unique:
                    ent["label"] += " · " + " ".join(unique)

    def field_value(self, ent: dict, key: str):
        for field in ent["fields"]:
            if field["path"] == [key] and field["value"] is not None:
                value = field["value"]
                return ", ".join(str(v) for v in value) if isinstance(value, list) else str(value)
        return None

    def assign_colors(self) -> None:
        self.family_counts: dict[str, int] = {}
        for ent in sorted(self.entities.values(), key=lambda e: (TYPE_ORDER.get(e["type"], UNKNOWN_ORDER), e["id"])):
            family = ent["family"]
            ent["colorIndex"] = self.family_counts.get(family, 0)
            self.family_counts[family] = ent["colorIndex"] + 1


def build(record_path: Path, text_path: Path) -> tuple[dict, Build]:
    record = json.loads(record_path.read_text())
    text = text_path.read_text()
    builder = Build(record, text)
    return builder.run(), builder


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("local_id", help="record id, e.g. xevP8UDRAVh9")
    parser.add_argument("--record", type=Path, help="extraction JSON (default: data/gold/<id>.extraction.json)")
    parser.add_argument("--text", type=Path, help="source text (default: the record's processed/local text)")
    parser.add_argument("--out", type=Path, help="output HTML (default: viewer/<id>.html)")
    args = parser.parse_args()

    record_path = args.record or ROOT / "data" / "gold" / f"{args.local_id}.extraction.json"
    text_path = args.text or ROOT / "data" / "texts" / args.local_id / "processed" / "local" / "text.tables.txt"
    out_path = args.out or ROOT / "viewer" / f"{args.local_id}.html"

    for path in (record_path, text_path, TEMPLATE):
        if not path.exists():
            print(f"missing: {path}", file=sys.stderr)
            return 1

    data, builder = build(record_path, text_path)
    payload = json.dumps(data, ensure_ascii=False).replace("</", "<\\/")
    html = TEMPLATE.read_text().replace("__DATA__", payload)
    html = html.replace("__TITLE__", data["record"]["title"].replace("<", "&lt;"))
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(html)

    print(
        f"{out_path}  ({len(html) // 1024} KB)  "
        f"{len(data['entities'])} entities, {len(data['edges'])} links, "
        f"{sum(len(e['fields']) for e in data['entities'])} fields, {len(data['spans'])} spans"
    )
    if builder.unknown_types:
        print(f"note: no colour family for {sorted(builder.unknown_types)} -- shown under Other",
              file=sys.stderr)
    if builder.misplaced_spans:
        print(f"WARNING: {builder.misplaced_spans} spans do not match the text at their offsets; "
              f"check that {text_path.name} is the file {record_path.name} was extracted from",
              file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())