"""Encode English text, commit supported graph edits, and decode a revision."""

import argparse
import json
from dataclasses import dataclass, replace
from graphlib import TopologicalSorter

import spacy
import tiktoken


MODEL_ENCODING = "o200k_base"


@dataclass(frozen=True)
class Token:
    text: str
    whitespace: str


@dataclass(frozen=True)
class Clause:
    id: int
    start: int
    stop: int


@dataclass(frozen=True)
class Edge:
    source: int
    target: int
    kind: str


@dataclass(frozen=True)
class Proposal:
    edge: Edge
    marker: int
    rule: str


@dataclass(frozen=True)
class SyntaxDAG:
    tokens: tuple[Token, ...]
    clauses: tuple[Clause, ...]
    edges: tuple[Edge, ...]
    proposals: tuple[Proposal, ...]
    history: tuple[Proposal, ...] = ()
    input_ids: tuple[int, ...] = ()

    def __post_init__(self):
        nodes = set(range(len(self.tokens))) | {clause.id for clause in self.clauses}
        sorter = TopologicalSorter()
        for node in nodes:
            sorter.add(node)
        for edge in self.edges:
            if edge.source not in nodes or edge.target not in nodes:
                raise ValueError("Edge endpoint is missing")
            sorter.add(edge.target, edge.source)
        tuple(sorter.static_order())

    def commit(self, proposal: Proposal):
        if proposal not in self.proposals:
            raise ValueError("Proposal was not mined from this graph")
        if proposal in self.history:
            return self
        return replace(
            self,
            edges=(*self.edges, proposal.edge),
            history=(*self.history, proposal),
        )


class Encoder:
    def __init__(self):
        self.tokenizer = tiktoken.get_encoding(MODEL_ENCODING)
        self.nlp = spacy.load("en_core_web_sm", disable=["ner"])

    def encode(self, text: str) -> SyntaxDAG:
        doc = self.nlp(text)
        tokens = tuple(Token(token.text, token.whitespace_) for token in doc)
        edges = [
            Edge(token.head.i, token.i, token.dep_)
            for token in doc
            if token.head != token
        ]
        clauses = []
        proposals = []
        for sentence in doc.sents:
            content = [token.i for token in sentence if not token.is_space]
            if not content:
                continue
            sentence = doc[content[0]:content[-1] + 1]
            match = self._causal_marker(sentence)
            if match is None:
                continue
            marker, end = match
            effect = Clause(len(tokens) + len(clauses), sentence.start, marker.i)
            cause = Clause(effect.id + 1, marker.i + 1, end)
            clauses.extend((effect, cause))
            for clause in (effect, cause):
                edges.extend(
                    Edge(clause.id, index, "contains")
                    for index in range(clause.start, clause.stop)
                )
            proposals.append(Proposal(Edge(cause.id, effect.id, "causes"),
                                      marker.i, "trailing-because"))
        return SyntaxDAG(
            tokens, tuple(clauses), tuple(edges), tuple(proposals),
            input_ids=tuple(self.tokenizer.encode_ordinary(text)),
        )

    @staticmethod
    def _causal_marker(sentence):
        markers = [token for token in sentence if token.lower_ == "because"]
        if len(markers) != 1 or sentence[-1].text != ".":
            return None
        marker = markers[0]
        verb = marker.head
        end = sentence.end - 1
        if (
            marker.dep_ != "mark"
            or verb.dep_ != "advcl"
            or verb.head != sentence.root
            or not sentence.start <= sentence.root.i < marker.i
            or not marker.i < verb.i < end
            or sentence.root.pos_ != "VERB"
            or verb.pos_ != "VERB"
            or sentence.root.tag_ not in {"VBD", "VBN"}
            or verb.tag_ != "VBD"
            or any(token.is_punct or token.is_space for token in sentence[:-1])
            or any(token.dep_ in {"neg", "aux", "conj"} for token in sentence)
            or not any(child.dep_ == "nsubj" for child in sentence.root.children)
            or not any(child.dep_ == "nsubj" for child in verb.children)
            or {token.i for token in verb.subtree if not token.is_punct}
            != set(range(marker.i, end))
        ):
            return None
        return marker, end


class Decoder:
    def decode_ids(self, graph: SyntaxDAG) -> tuple[int, ...]:
        tokenizer = tiktoken.get_encoding(MODEL_ENCODING)
        return tuple(tokenizer.encode_ordinary(self.decode(graph)))

    def decode(self, graph: SyntaxDAG) -> str:
        clauses = {clause.id: clause for clause in graph.clauses}
        edits = {}
        for edge in graph.edges:
            if edge.kind != "causes":
                continue
            cause, effect = clauses[edge.source], clauses[edge.target]
            cause_text = self._surface(graph, cause.start, cause.stop).rstrip()
            effect_text = self._surface(graph, effect.start, effect.stop).rstrip()
            edits[effect.start] = (cause.stop, f"Cause: {cause_text}. Effect: {effect_text}")
        output = []
        index = 0
        while index < len(graph.tokens):
            stop, surface = edits.get(index, (index + 1, self._surface(graph, index, index + 1)))
            output.append(surface)
            index = stop
        return "".join(output)

    @staticmethod
    def _surface(graph, start, stop):
        return "".join(token.text + token.whitespace for token in graph.tokens[start:stop])


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("text")
    parser.add_argument("--commit", action="store_true", help="Apply mined causal edits")
    parser.add_argument("--ids", action="store_true", help="Emit model token IDs as JSON")
    args = parser.parse_args()
    graph = Encoder().encode(args.text)
    if args.commit:
        for proposal in graph.proposals:
            graph = graph.commit(proposal)
    decoder = Decoder()
    text = decoder.decode(graph)
    if args.ids:
        print(json.dumps({
            "encoding": MODEL_ENCODING,
            "input_ids": graph.input_ids,
            "output_ids": decoder.decode_ids(graph),
            "text": text,
        }, ensure_ascii=False))
    else:
        print(text)


if __name__ == "__main__":
    main()