"""Phrase boosting for Klang Pianissimo. A numpy port of NeMo's GPU-PB boosting tree
(ContextGraph and GPUBoostingTreeModel, NeMo 2.7, Apache 2.0)."""
import math
from collections import deque

import numpy as np
import sentencepiece as spm


class _Node:
    __slots__ = ("id", "children", "token_score", "node_score", "is_end", "fail")

    def __init__(self, node_id, token_score, node_score, is_end):
        self.id = node_id
        self.children = {}
        self.token_score = token_score
        self.node_score = node_score
        self.is_end = is_end
        self.fail = None


class PhraseBoost:
    def __init__(self, phrases, tokenizer_path, vocab_size, context_score=1.0, depth_scaling=2.0):
        sp = spm.SentencePieceProcessor(model_file=str(tokenizer_path))
        self.vocab_size = vocab_size
        self.root = _Node(0, 0.0, 0.0, False)
        self.root.fail = self.root
        self.nodes = [self.root]
        for phrase in dict.fromkeys(phrases):
            self._add(sp.encode(phrase), context_score, depth_scaling)
        self._fill_fail()
        self._cache = {}

    def _add(self, tokens, context_score, depth_scaling):
        node = self.root
        for i, token in enumerate(tokens):
            last = i == len(tokens) - 1
            child = node.children.get(token)
            if child is None:
                token_score = context_score * depth_scaling + math.log(i + 1) if i > 0 else context_score
                child = _Node(len(self.nodes), token_score, node.node_score + token_score, last)
                node.children[token] = child
                self.nodes.append(child)
            else:
                child.token_score = max(context_score, child.token_score)
                child.node_score = node.node_score + child.token_score
                child.is_end = last or child.is_end
            node = child

    def _fill_fail(self):
        queue = deque()
        for child in self.root.children.values():
            child.fail = self.root
            queue.append(child)
        while queue:
            node = queue.popleft()
            for token, child in node.children.items():
                fail = node.fail
                while token not in fail.children and fail is not self.root:
                    fail = fail.fail
                child.fail = fail.children.get(token, self.root)
                queue.append(child)

    def scores(self, state):
        """Bonus and next state for every token from `state`."""
        if state not in self._cache:
            bonus = np.zeros(self.vocab_size, dtype=np.float32)
            nxt = np.full(self.vocab_size, -1, dtype=np.int64)
            node, taken_back = self.nodes[state], 0.0
            while True:
                for token, child in node.children.items():
                    if nxt[token] < 0:
                        bonus[token] = taken_back + child.token_score
                        nxt[token] = child.id
                if node is self.root:
                    unset = nxt < 0
                    bonus[unset] = taken_back
                    nxt[unset] = 0
                    break
                if not node.is_end:
                    taken_back += node.fail.node_score - node.node_score
                node = node.fail
            self._cache[state] = (bonus, nxt)
        return self._cache[state]