import unittest
from dataclasses import replace
from graphlib import CycleError

import tiktoken

from syntax_dag import Decoder, Edge, Encoder, Proposal


class SyntaxDAGTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.encoder = Encoder()
        cls.decoder = Decoder()

    def test_round_trip(self):
        for text in ("", "  \t", "  It rained.\n\nThe road flooded.  "):
            with self.subTest(text=text):
                graph = self.encoder.encode(text)
                self.assertEqual(self.decoder.decode(graph), text)
                self.assertFalse(graph.proposals)

    def test_causal_commit(self):
        text = "The road flooded because it rained.  They waited."
        graph = self.encoder.encode(text)
        self.assertEqual(len(graph.proposals), 1, [
            (token.text, token.head.text, token.dep_, token.pos_, token.tag_)
            for token in self.encoder.nlp(text)
        ])
        proposal, = graph.proposals
        self.assertTrue(any(
            len({edge.source for edge in graph.edges if edge.target == node}) > 1
            for node in range(len(graph.tokens))
        ))
        committed = graph.commit(proposal)
        self.assertEqual(self.decoder.decode(graph), text)
        self.assertEqual(
            self.decoder.decode(committed),
            "Cause: it rained. Effect: The road flooded.  They waited.",
        )
        self.assertEqual(committed.history, (proposal,))
        self.assertIs(committed.commit(proposal), committed)
        self.assertNotIn(proposal.edge, graph.edges)

    def test_cycle_rejects_entire_commit(self):
        graph = self.encoder.encode("The road flooded because it rained.")
        original, = graph.proposals
        reverse = Proposal(Edge(original.edge.target, original.edge.source, "causes"),
                           original.marker, "test")
        committed = replace(graph.commit(original), proposals=(original, reverse))
        with self.assertRaises(CycleError):
            committed.commit(reverse)
        self.assertEqual(committed.history, (original,))
        self.assertNotIn(reverse.edge, committed.edges)

    def test_negation_and_questions_are_left_alone(self):
        for text in ("The road did not flood because it rained.",
                     "Did the road flood because it rained?"):
            with self.subTest(text=text):
                self.assertFalse(self.encoder.encode(text).proposals)


    def test_model_token_ids(self):
        tokenizer = tiktoken.get_encoding("o200k_base")
        text = "中文 🙂 <|endoftext|>"
        graph = self.encoder.encode(text)
        self.assertEqual(graph.input_ids, tuple(tokenizer.encode_ordinary(text)))
        self.assertEqual(tokenizer.decode(graph.input_ids), text)
        self.assertEqual(self.decoder.decode_ids(graph), graph.input_ids)

        graph = self.encoder.encode("The road flooded because it rained.")
        committed = graph.commit(graph.proposals[0])
        self.assertEqual(committed.input_ids, graph.input_ids)
        self.assertNotEqual(self.decoder.decode_ids(committed), graph.input_ids)
        self.assertEqual(
            tokenizer.decode(self.decoder.decode_ids(committed)),
            self.decoder.decode(committed),
        )