# Original Algorithm:
# By Steve Hanov, 2011. Released to the public domain.
# Please see http://stevehanov.ca/blog/index.php?id=115 for the accompanying article.
#
# Adapted for PyPy/CPython by Carl Friedrich Bolz-Tereick
#
# Based on Daciuk, Jan, et al. "Incremental construction of minimal acyclic finite-state automata."
# Computational linguistics 26.1 (2000): 3-16.
#
# Updated 2014 to use DAWG as a mapping; see
# Kowaltowski, T.; CL. Lucchesi (1993), "Applications of finite automata representing large vocabularies",
# Software-Practice and Experience 1993

from collections import defaultdict
from functools import cached_property


# This class represents a node in the directed acyclic word graph (DAWG). It
# has a list of edges to other nodes. It has functions for testing whether it
# is equivalent to another node. Nodes are equivalent if they have identical
# edges, and each identical edge leads to identical states. The __hash__ and
# __eq__ functions allow it to be used as a key in a python dictionary.


class DawgNode:

    def __init__(self, dawg):
        self.id = dawg.next_id
        dawg.next_id += 1
        self.final = False
        self.edges = {}

        self.linear_edges = None # later: list of (string, next_state)

    def __str__(self):
        if self.final:
            arr = ["1"]
        else:
            arr = ["0"]

        for (label, node) in sorted(self.edges.items()):
            arr.append(label)
            arr.append(str(node.id))

        return "_".join(arr)
    __repr__ = __str__

    def _as_tuple(self):
        edges = sorted(self.edges.items())
        edge_tuple = tuple((label, node.id) for label, node in edges)
        return (self.final, edge_tuple)

    def __hash__(self):
        return hash(self._as_tuple())

    def __eq__(self, other):
        return self._as_tuple() == other._as_tuple()

    @cached_property
    def num_reachable_linear(self):
        # returns the number of different paths to final nodes reachable from
        # this one

        count = 0
        # staying at self counts as a path if self is final
        if self.final:
            count += 1
        for label, node in self.linear_edges:
            count += node.num_reachable_linear

        return count


class Dawg:
    def __init__(self):
        self.previous_word = ""
        self.next_id = 0
        self.root = DawgNode(self)

        # Here is a list of nodes that have not been checked for duplication.
        self.unchecked_nodes = []

        # To deduplicate, maintain a dictionary with
        # minimized_nodes[canonical_node] is canonical_node.
        # Based on __hash__ and __eq__, minimized_nodes[n] is the
        # canonical node equal to n.
        # In other words, self.minimized_nodes[x] == x for all nodes found in
        # the dict.
        self.minimized_nodes = {}

        # word: value mapping
        self.data = {}
        # value: word mapping
        self.inverse = {}

    def insert(self, word, value):
        if not all(0 <= ord(c) < 128 for c in word):
            raise ValueError("Use 7-bit ASCII characters only")
        if word <= self.previous_word:
            raise ValueError("Error: Words must be inserted in alphabetical order.")
        if value in self.inverse:
            raise ValueError(f"value {value} is duplicate, got it for word {self.inverse[value]} and now {word}")

        # find common prefix between word and previous word
        common_prefix = 0
        for i in range(min(len(word), len(self.previous_word))):
            if word[i] != self.previous_word[i]:
                break
            common_prefix += 1

        # Check the unchecked_nodes for redundant nodes, proceeding from last
        # one down to the common prefix size. Then truncate the list at that
        # point.
        self._minimize(common_prefix)

        self.data[word] = value
        self.inverse[value] = word

        # add the suffix, starting from the correct node mid-way through the
        # graph
        if len(self.unchecked_nodes) == 0:
            node = self.root
        else:
            node = self.unchecked_nodes[-1][2]

        for letter in word[common_prefix:]:
            next_node = DawgNode(self)
            node.edges[letter] = next_node
            self.unchecked_nodes.append((node, letter, next_node))
            node = next_node

        node.final = True
        self.previous_word = word

    def finish(self):
        if not self.data:
            raise ValueError("need at least one word in the dawg")
        # minimize all unchecked_nodes
        self._minimize(0)

        self._linearize_edges()

        topoorder, linear_data, inverse = self._topological_order()
        return self.compute_packed(topoorder), linear_data, inverse

    def _minimize(self, down_to):
        # proceed from the leaf up to a certain point
        for i in range(len(self.unchecked_nodes) - 1, down_to - 1, -1):
            (parent, letter, child) = self.unchecked_nodes[i]
            if child in self.minimized_nodes:
                # replace the child with the previously encountered one
                parent.edges[letter] = self.minimized_nodes[child]
            else:
                # add the state to the minimized nodes.
                self.minimized_nodes[child] = child
            self.unchecked_nodes.pop()

    def _lookup(self, word):
        """ Return an integer 0 <= k < number of strings in dawg
        where word is the kth successful traversal of the dawg. """
        node = self.root
        skipped = 0  # keep track of number of final nodes that we skipped
        index = 0
        while index < len(word):
            for label, child in node.linear_edges:
                if word[index] == label[0]:
                    if word[index:index + len(label)] == label:
                        if node.final:
                            skipped += 1
                        index += len(label)
                        node = child
                        break
                    else:
                        return None
                skipped += child.num_reachable_linear
            else:
                return None
        return skipped

    def enum_all_nodes(self):
        stack = [self.root]
        done = set()
        while stack:
            node = stack.pop()
            if node.id in done:
                continue
            yield node
            done.add(node.id)
            for label, child in sorted(node.edges.items()):
                stack.append(child)

    def prettyprint(self):
        for node in sorted(self.enum_all_nodes(), key=lambda e: e.id):
            s_final = " final" if node.final else ""
            print(f"{node.id}: ({node}) {s_final}")
            for label, child in sorted(node.edges.items()):
                print(f"    {label} goto {child.id}")

    def _inverse_lookup(self, number):
        assert 0, "not working in the current form, but keep it as the pure python version of compact lookup"
        result = []
        node = self.root
        while 1:
            if node.final:
                if pos == 0:
                    return "".join(result)
                pos -= 1
            for label, child in sorted(node.edges.items()):
                nextpos = pos - child.num_reachable_linear
                if nextpos < 0:
                    result.append(label)
                    node = child
                    break
                else:
                    pos = nextpos
            else:
                assert 0

    def _linearize_edges(self):
        # compute "linear" edges. the idea is that long chains of edges without
        # any of the intermediate states being final or any extra incoming or
        # outgoing edges can be represented by having removing them, and
        # instead using longer strings as edge labels (instead of single
        # characters)
        incoming = defaultdict(list)
        nodes = sorted(self.enum_all_nodes(), key=lambda e: e.id)
        for node in nodes:
            for label, child in sorted(node.edges.items()):
                incoming[child].append(node)
        for node in nodes:
            node.linear_edges = []
            for label, child in sorted(node.edges.items()):
                s = [label]
                while len(child.edges) == 1 and len(incoming[child]) == 1 and not child.final:
                    (c, child), = child.edges.items()
                    s.append(c)
                node.linear_edges.append((''.join(s), child))

    def _topological_order(self):
        # compute reachable linear nodes, and the set of incoming edges for each node
        order = []
        stack = [self.root]
        seen = set()
        while stack:
            # depth first traversal
            node = stack.pop()
            if node.id in seen:
                continue
            seen.add(node.id)
            order.append(node)
            for label, child in node.linear_edges:
                stack.append(child)

        # do a (slightly bad) topological sort
        incoming = defaultdict(set)
        for node in order:
            for label, child in node.linear_edges:
                incoming[child].add((label, node))
        no_incoming = [order[0]]
        topoorder = []
        positions = {}
        while no_incoming:
            node = no_incoming.pop()
            topoorder.append(node)
            positions[node] = len(topoorder)
            # use "reversed" to make sure that the linear_edges get reorderd
            # from their alphabetical order as little as necessary (no_incoming
            # is LIFO)
            for label, child in reversed(node.linear_edges):
                incoming[child].discard((label, node))
                if not incoming[child]:
                    no_incoming.append(child)
                    del incoming[child]
        # check result
        assert set(topoorder) == set(order)
        assert len(set(topoorder)) == len(topoorder)

        for node in order:
            node.linear_edges.sort(key=lambda element: positions[element[1]])

        for node in order:
            for label, child in node.linear_edges:
                assert positions[child] > positions[node]
        # number the nodes. afterwards every input string in the set has a
        # unique number in the 0 <= number < len(data). We then put the data in
        # self.data into a linear list using these numbers as indexes.
        topoorder[0].num_reachable_linear
        linear_data = [None] * len(self.data)
        inverse = {} # maps value back to index
        for word, value in self.data.items():
            index = self._lookup(word)
            linear_data[index] = value
            inverse[value] = index

        return topoorder, linear_data, inverse

    def compute_packed(self, order):
        def compute_chunk(node, offsets):
            """ compute the packed node/edge data for a node. result is a
            list of bytes as long as order. the jump distance calculations use
            the offsets dictionary to know where in the final big output
            bytestring the individual nodes will end up. """
            result = bytearray()
            offset = offsets[node]
            encode_varint_unsigned(number_add_bits(node.num_reachable_linear, node.final), result)
            if len(node.linear_edges) == 0:
                assert node.final
                encode_varint_unsigned(0, result) # add a 0 saying "done"
            prev_child_offset = offset + len(result)
            for edgeindex, (label, targetnode) in enumerate(node.linear_edges):
                label = label.encode('ascii')
                child_offset = offsets[targetnode]
                child_offset_difference = child_offset - prev_child_offset

                info = number_add_bits(child_offset_difference, len(label) == 1, edgeindex == len(node.linear_edges) - 1)
                if edgeindex == 0:
                    assert info != 0
                encode_varint_unsigned(info, result)
                prev_child_offset = child_offset
                if len(label) > 1:
                    encode_varint_unsigned(len(label), result)
                result.extend(label)
            return result

        def compute_new_offsets(chunks, offsets):
            """ Given a list of chunks, compute the new offsets (by adding the
            chunk lengths together). Also check if we cannot shrink the output
            further because none of the node offsets are smaller now. if that's
            the case return None. """
            new_offsets = {}
            curr_offset = 0
            should_continue = False
            for node, result in zip(order, chunks):
                if curr_offset < offsets[node]:
                    # the new offset is below the current assumption, this
                    # means we can shrink the output more
                    should_continue = True
                new_offsets[node] = curr_offset
                curr_offset += len(result)
            if not should_continue:
                return None
            return new_offsets

        # assign initial offsets to every node
        offsets = {}
        for i, node in enumerate(order):
            # we don't know position of the edge yet, just use something big as
            # the starting position. we'll have to do further iterations anyway,
            # but the size is at least a lower limit then
            offsets[node] = i * 2 ** 30


        # due to the variable integer width encoding of edge targets we need to
        # run this to fixpoint. in the process we shrink the output more and
        # more until we can't any more. at any point we can stop and use the
        # output, but we might need padding zero bytes when joining the chunks
        # to have the correct jump distances
        last_offsets = None
        while 1:
            chunks = [compute_chunk(node, offsets) for node in order]
            last_offsets = offsets
            offsets = compute_new_offsets(chunks, offsets)
            if offsets is None: # couldn't shrink
                break

        # build the final packed string
        total_result = bytearray()
        for node, result in zip(order, chunks):
            node_offset = last_offsets[node]
            if node_offset > len(total_result):
                # need to pad to get the offsets correct
                padding = b"\x00" * (node_offset - len(total_result))
                total_result.extend(padding)
            assert node_offset == len(total_result)
            total_result.extend(result)
        return bytes(total_result)


# ______________________________________________________________________
# the following functions operate on the packed representation

def number_add_bits(x, *bits):
    for bit in bits:
        assert bit == 0 or bit == 1
        x = (x << 1) | bit
    return x

def encode_varint_unsigned(i, res):
    # https://en.wikipedia.org/wiki/LEB128 unsigned variant
    more = True
    startlen = len(res)
    if i < 0:
        raise ValueError("only positive numbers supported", i)
    while more:
        lowest7bits = i & 0b1111111
        i >>= 7
        if i == 0:
            more = False
        else:
            lowest7bits |= 0b10000000
        res.append(lowest7bits)
    return len(res) - startlen

def number_split_bits(x, n, acc=()):
    if n == 1:
        return x >> 1, x & 1
    if n == 2:
        return x >> 2, (x >> 1) & 1, x & 1
    assert 0, "implement me!"

def decode_varint_unsigned(b, index=0):
    res = 0
    shift = 0
    while True:
        byte = b[index]
        res = res | ((byte & 0b1111111) << shift)
        index += 1
        shift += 7
        if not (byte & 0b10000000):
            return res, index

def decode_node(packed, node):
    x, node = decode_varint_unsigned(packed, node)
    node_count, final = number_split_bits(x, 1)
    return node_count, final, node

def decode_edge(packed, edgeindex, prev_child_offset, offset):
    x, offset = decode_varint_unsigned(packed, offset)
    if x == 0 and edgeindex == 0:
        raise KeyError # trying to decode past a final node
    child_offset_difference, len1, last_edge = number_split_bits(x, 2)
    child_offset = prev_child_offset + child_offset_difference
    if len1:
        size = 1
    else:
        size, offset = decode_varint_unsigned(packed, offset)
    return child_offset, last_edge, size, offset

def _match_edge(packed, s, size, node_offset, stringpos):
    if size > 1 and stringpos + size > len(s):
        # past the end of the string, can't match
        return False
    for i in range(size):
        if packed[node_offset + i] != s[stringpos + i]:
            # if a subsequent char of an edge doesn't match, the word isn't in
            # the dawg
            if i > 0:
                raise KeyError
            return False
    return True

def lookup(packed, data, s):
    return data[_lookup(packed, s)]

def _lookup(packed, s):
    stringpos = 0
    node_offset = 0
    skipped = 0  # keep track of number of final nodes that we skipped
    false = False
    while stringpos < len(s):
        #print(f"{node_offset=} {stringpos=}")
        _, final, edge_offset = decode_node(packed, node_offset)
        prev_child_offset = edge_offset
        edgeindex = 0
        while 1:
            child_offset, last_edge, size, edgelabel_chars_offset = decode_edge(packed, edgeindex, prev_child_offset, edge_offset)
            #print(f"    {edge_offset=} {child_offset=} {last_edge=} {size=} {edgelabel_chars_offset=}")
            edgeindex += 1
            prev_child_offset = child_offset
            if _match_edge(packed, s, size, edgelabel_chars_offset, stringpos):
                # match
                if final:
                    skipped += 1
                stringpos += size
                node_offset = child_offset
                break
            if last_edge:
                raise KeyError
            descendant_count, _, _ = decode_node(packed, child_offset)
            skipped += descendant_count
            edge_offset = edgelabel_chars_offset + size
    _, final, _ = decode_node(packed, node_offset)
    if final:
        return skipped
    raise KeyError

def inverse_lookup(packed, inverse, x):
    pos = inverse[x]
    return _inverse_lookup(packed, pos)

def _inverse_lookup(packed, pos):
    result = bytearray()
    node_offset = 0
    while 1:
        node_count, final, edge_offset = decode_node(packed, node_offset)
        if final:
            if pos == 0:
                return bytes(result)
            pos -= 1
        prev_child_offset = edge_offset
        edgeindex = 0
        while 1:
            child_offset, last_edge, size, edgelabel_chars_offset = decode_edge(packed, edgeindex, prev_child_offset, edge_offset)
            edgeindex += 1
            prev_child_offset = child_offset
            descendant_count, _, _ = decode_node(packed, child_offset)
            nextpos = pos - descendant_count
            if nextpos < 0:
                assert edgelabel_chars_offset >= 0
                result.extend(packed[edgelabel_chars_offset: edgelabel_chars_offset + size])
                node_offset = child_offset
                break
            elif not last_edge:
                pos = nextpos
                edge_offset = edgelabel_chars_offset + size
            else:
                raise KeyError
        else:
            raise KeyError


def build_compression_dawg(ucdata):
    d = Dawg()
    ucdata.sort()
    for name, value in ucdata:
        d.insert(name, value)
    packed, pos_to_code, reversedict = d.finish()
    print("size of dawg [KiB]", round(len(packed) / 1024, 2))
    # check that lookup and inverse_lookup work correctly on the input data
    for name, value in ucdata:
        assert lookup(packed, pos_to_code, name.encode('ascii')) == value
        assert inverse_lookup(packed, reversedict, value) == name.encode('ascii')
    return packed, pos_to_code