"""The node tree of scienting-doc/1.

Non-normative; see spec/03-nodes.md.
"""

BLOCKS = frozenset(
    ["paragraph", "equation", "figure", "table", "list", "quote", "code", "algorithm", "theorem", "footnote"]
)

# Allowed child kinds for each node kind.
CHILDREN = {
    "document": frozenset(["front", "section", "back"]),
    "front": frozenset(["title", "authors", "abstract", "keywords", "notes"]),
    "title": frozenset(["sentence"]),
    "authors": frozenset(),
    "abstract": BLOCKS | {"section"},
    "keywords": frozenset(),
    "notes": frozenset(["footnote", "paragraph"]),
    "section": BLOCKS | {"heading", "section"},
    "heading": frozenset(["sentence"]),
    "paragraph": frozenset(["sentence"]),
    "sentence": frozenset(["equation"]),
    "equation": frozenset(["equation-line"]),
    "equation-line": frozenset(),
    "figure": frozenset(["panel", "caption", "figure-note"]),
    "panel": frozenset(),
    "caption": frozenset(["sentence"]),
    "figure-note": frozenset(["sentence"]),
    "table": frozenset(["caption", "table-head", "table-body", "table-note"]),
    "table-head": frozenset(["row"]),
    "table-body": frozenset(["row"]),
    "table-note": frozenset(["sentence"]),
    "row": frozenset(["cell"]),
    "cell": frozenset(),
    "list": frozenset(["item"]),
    "item": BLOCKS,
    "quote": BLOCKS,
    "code": frozenset(),
    "algorithm": frozenset(),
    "theorem": BLOCKS,
    "footnote": frozenset(["sentence"]),
    "back": frozenset(["acknowledgments", "references", "appendix", "notes"]),
    "acknowledgments": BLOCKS,
    "references": frozenset(["reference"]),
    "reference": frozenset(),
    "appendix": frozenset(["section"]),
}
KINDS = frozenset(CHILDREN)

# Kinds that must have at least one child.
NONEMPTY = frozenset(
    ["title", "heading", "paragraph", "caption", "figure-note", "table-note", "table-head", "table-body",
     "row", "list", "item", "quote", "theorem", "footnote", "appendix"]
)
# Kinds that occur at most once among their siblings.
AT_MOST_ONCE = frozenset(
    ["front", "back", "title", "authors", "abstract", "keywords", "caption", "table-head", "table-body",
     "acknowledgments", "references"]
)

RUN_KINDS = frozenset(["sentence", "cell", "reference", "authors", "keywords"])
MATH_KINDS = frozenset(["equation", "equation-line"])
VERBATIM_KINDS = frozenset(["code", "algorithm"])
LEAVES = RUN_KINDS | MATH_KINDS | VERBATIM_KINDS
CONTAINERS_WITH_CONTENT = frozenset(["section", "figure", "panel", "table", "theorem", "footnote"])
CONTAINERS_WITHOUT_CONTENT = KINDS - LEAVES - CONTAINERS_WITH_CONTENT
LABELLED = frozenset(
    ["section", "figure", "panel", "table", "theorem", "footnote", "equation", "equation-line", "reference",
     "code", "algorithm"]
)

RUNS = frozenset(["text", "math", "cite", "xref", "fnmark", "link", "equation"])
STYLES = frozenset(["em", "strong", "sup", "sub", "sc", "mono"])

LEVELS = ["F0", "F1", "F2", "F3"]


def anchor_class(kind):
    if kind in LEAVES:
        return "leaf"
    if kind in CONTAINERS_WITH_CONTENT:
        return "container-with-content"
    return "container-without-content"


class Doc:
    """A document: nodes by ID with parent and order, plus document order."""

    def __init__(self, key, nodes):
        self.key = key
        self.nodes = {}
        self.order = []
        self.children = {}
        for n in nodes:
            self.nodes[n["id"]] = n
            self.order.append(n["id"])
            self.children.setdefault(n["id"], [])
        for nid in self.order:
            p = self.nodes[nid].get("parent")
            if p is not None:
                self.children.setdefault(p, []).append(nid)
        for p, kids in self.children.items():
            kids.sort(key=lambda c: self.nodes[c]["ord"])
        roots = [nid for nid in self.order if self.nodes[nid].get("parent") is None]
        self.root = roots[0] if roots else None

    def kids(self, nid):
        return self.children.get(nid, [])

    def preorder(self, nid=None):
        nid = self.root if nid is None else nid
        out = []
        stack = [nid]
        while stack:
            x = stack.pop()
            out.append(x)
            stack.extend(reversed(self.kids(x)))
        return out

    def ancestors(self, nid):
        out = []
        p = self.nodes[nid].get("parent")
        while p is not None:
            out.append(p)
            p = self.nodes[p].get("parent")
        return list(reversed(out))

    def text_leaves(self, nid):
        """The text-bearing nodes a container's projections are made of, in document order."""
        out = []
        stack = [nid]
        while stack:
            x = stack.pop()
            n = self.nodes[x]
            k = n["kind"]
            if k in RUN_KINDS or k in VERBATIM_KINDS or (k == "equation" and self._parent_kind(x) != "sentence"):
                out.append(x)
                continue
            if k == "equation-line" or k == "equation":
                continue
            stack.extend(reversed(self.kids(x)))
        return out

    def leaves(self, nid):
        """Every leaf below or at `nid`, including equations inside sentences and equation lines."""
        return [x for x in self.preorder(nid) if self.nodes[x]["kind"] in LEAVES]

    def _parent_kind(self, nid):
        p = self.nodes[nid].get("parent")
        return None if p is None else self.nodes[p]["kind"]