"""Steering a turn while it is sampled: how long the narration runs, and what the choices look like."""
import math
import re

import torch

from .bard_rules import is_story_choice

LENGTHS = {"short": 60, "medium": 150, "long": None}
SENTENCE_END = re.compile(r"[.!?…][\"”’')\]]*\s*$|\n\s*$")
SENTENCE_DONE = re.compile(r"[.!?…][\"”’')\]]*(\s|$)")
NEG_INF = float("-inf")


class LengthControl:

    MIN_CHOICES = 2
    MAX_CHOICES = 5
    MAX_CHOICE_TOKENS = 28
    COPY_PENALTY = 1.5
    COPY_PENALTY_MAX = 8.0
    ECHO_PENALTY = 4.0
    MIN_WORDS = 4
    GIVE_UP_AFTER = 90

    def __init__(self, tok, mode="medium", vars_weight=1.0):
        self.tok = tok
        self.vars_id = tok.specials["<|vars|>"]
        self.choices_id = tok.specials["<|choices|>"]
        self.sep_id = tok.specials["<|sep|>"]
        self.go_on = [self.vars_id, self.choices_id]
        self.stop = self.go_on + [tok.eos_id]
        self.rep_free = False

        first_special = min(tok.specials.values())
        decoded = tok.vocab.id_to_token_decoded
        self.silent = {i for i in range(first_special) if decoded(i) == ""}
        self.pipes = torch.tensor([i for i in range(first_special) if "|" in decoded(i)], dtype=torch.long)
        text_like = {self.sep_id, tok.eos_id, tok.specials["<|num|>"], tok.specials["<|snum|>"]}
        self.no_specials = torch.tensor(
            [i for i in range(first_special, 1 << 14) if i not in text_like], dtype=torch.long)

        self.vars_bias = math.log(vars_weight) if vars_weight > 1 else 0.0
        self.set_length(mode)

    def set_length(self, mode):
        self.mode = mode
        self.budget = LENGTHS[mode]

    def start_part(self, raw):
        self.written = raw
        self.done = "<|" in raw
        self.offset = 0 if self.done else len(self.tok.encode(raw)[0])
        self.prior = None
        if "<|choices|>" in raw:
            texts = raw.split("<|choices|>", 1)[1].split("<|sep|>")
            self.prior = [[int(i) for i in self.tok.encode(text)[0]] if text else [] for text in texts]
        self.rep_free = False

    def __call__(self, generated, logits):
        logits = logits.clone()
        logits[self.pipes.to(logits.device)] = NEG_INF
        if self.prior is not None or self.choices_id in generated:
            self.rep_free = True
            return self.steer_choices(generated, logits)
        if self.done or any(token in self.stop for token in generated):
            self.done = True
            return logits
        return self.steer_narration(generated, logits)

    def steer_narration(self, generated, logits):
        if not self.has_own_sentence(generated):
            logits[self.stop] = NEG_INF
        if self.vars_bias:
            logits[self.vars_id] -= self.vars_bias
        if self.budget is None or not generated:
            return logits

        length = self.offset + len(generated)
        ends_a_sentence = SENTENCE_END.search(self.tok.decode(generated[-4:], [0] * 4))
        if length < self.budget or not ends_a_sentence:
            return logits
        if length >= 2 * self.budget:
            return self.end_narration(logits)
        logits[self.go_on] += 2.0 + 6.0 * (length - self.budget) / self.budget
        return logits

    def has_own_sentence(self, generated):
        if self.offset + len(generated) >= self.GIVE_UP_AFTER:
            return True
        text = (self.written + self.tok.decode(generated)).strip()
        return bool(SENTENCE_DONE.search(text)) and len(text.split()) >= self.MIN_WORDS

    def end_narration(self, logits):
        best = max(self.go_on, key=lambda token: float(logits[token]))
        kept = logits[best].clone()
        logits.fill_(NEG_INF)
        logits[best] = kept
        return logits

    def steer_choices(self, generated, logits):
        finished, current = self.split_choices(generated)
        real_choices = sum(1 for choice in finished if is_story_choice(self.tok.vocab.decode(choice)))

        if not current or real_choices + 1 < self.MIN_CHOICES:
            logits[self.tok.eos_id] = NEG_INF
        if not current:
            logits[self.sep_id] = NEG_INF
        stray = self.no_specials.to(logits.device)
        logits[stray[stray < logits.numel()]] = NEG_INF
        if len(finished) + 1 >= self.MAX_CHOICES:
            logits[self.sep_id] = NEG_INF

        if len(current) >= self.MAX_CHOICE_TOKENS:
            return self.close_choice(logits)
        self.discourage_copies(logits, finished, current)
        return logits

    def split_choices(self, generated):
        if self.prior is not None:
            *finished, head = self.prior
            written = list(generated)
        else:
            finished, head = [], []
            after_marker = len(generated) - generated[::-1].index(self.choices_id)
            written = list(generated[after_marker:])

        pieces = [[]]
        for token in written:
            if token == self.sep_id:
                pieces.append([])
            else:
                pieces[-1].append(token)
        pieces[0] = head + pieces[0]
        return [choice for choice in finished + pieces[:-1] if choice], pieces[-1]

    def close_choice(self, logits):
        closers = [self.sep_id, self.tok.eos_id]
        closing = torch.full_like(logits, NEG_INF)
        closing[closers] = logits[closers]
        return logits if torch.isinf(closing).all() else closing

    def discourage_copies(self, logits, finished, current):
        penalties = {}
        for choice in finished:
            if choice == current:
                logits[[self.sep_id, self.tok.eos_id]] = NEG_INF
            self.penalize_retracing(penalties, choice, current)
            self.penalize_echo(penalties, choice, current)
        for token, amount in penalties.items():
            logits[token] -= amount

    def penalize_retracing(self, penalties, choice, current):
        done = len(current)
        if done < len(choice) and choice[:done] == current and choice[done] not in self.silent:
            amount = min(self.COPY_PENALTY_MAX, self.COPY_PENALTY * (done + 1))
            self.raise_penalty(penalties, choice[done], amount)

    def penalize_echo(self, penalties, choice, current):
        if len(current) < 2:
            return
        pair = current[-2:]
        for i in range(len(choice) - 2):
            if choice[i:i + 2] == pair and choice[i + 2] not in self.silent:
                self.raise_penalty(penalties, choice[i + 2], self.ECHO_PENALTY)

    @staticmethod
    def raise_penalty(penalties, token, amount):
        penalties[token] = max(penalties.get(token, 0.0), amount)