"""Sampling from the model one token at a time."""
import torch
from torch.nn import functional as F

from .bard_core import KVCache, softcap

NEG_INF = float("-inf")


class Reader:

    def __init__(self, lm, ids, vals, context):
        self.lm = lm
        self.ids = ids
        self.vals = vals
        self.context = context
        self.cache = KVCache(lm.body)
        self.read = 0

    def last_hidden(self):
        if self.cache.pos + len(self.ids) - self.read > self.context:
            self.cache.reset()
            self.read = max(0, len(self.ids) - (3 * self.context) // 4)
        device = self.lm.head.weight.device
        ids = torch.tensor([self.ids[self.read:]], device=device)
        vals = torch.tensor([self.vals[self.read:]], device=device)
        hidden = self.lm.hidden(ids, vals, cache=self.cache)[0, -1]
        self.read = len(self.ids)
        return hidden

    def add(self, token, value):
        self.ids.append(token)
        self.vals.append(value)


def never_sampled(lm, tok):
    banned = torch.zeros(lm.head.weight.shape[0], dtype=torch.bool, device=lm.head.weight.device)
    first_special = min(tok.specials.values())
    banned[tok.vocab_size:] = True
    banned[first_special:tok.vocab_size] = True
    allowed = {tok.eos_id, tok.num_id, tok.snum_id}
    allowed |= {i for name, i in tok.specials.items() if not name.startswith("<|start|>")}
    banned[list(allowed)] = False
    return banned


def penalize_recent(logits, recent_ids, penalty):
    recent = torch.tensor(sorted(set(recent_ids)), device=logits.device)
    picked = logits[recent]
    logits[recent] = torch.where(picked > 0, picked / penalty, picked * penalty)


def sample(logits, temperature, top_k, top_p, rng):
    logits = logits / temperature
    if top_k:
        threshold = torch.topk(logits, min(top_k, logits.numel())).values[-1]
        logits[logits < threshold] = NEG_INF
    probs = F.softmax(logits, -1)
    if top_p < 1.0:
        sorted_probs, order = probs.sort(descending=True)
        outside = sorted_probs.cumsum(0) - sorted_probs > top_p
        probs[order[outside]] = 0.0
        probs /= probs.sum()
    return int(torch.multinomial(probs, 1, generator=rng))


@torch.no_grad()
def generate(lm, tok, seed, max_new, context, rng, temperature=0.8, top_k=50, top_p=0.95, rep_penalty=1.1,
             rep_window=128, adjust_logits=None):
    seed_ids, seed_vals = tok.encode(seed)
    reader = Reader(lm, [tok.eos_id] + seed_ids.tolist(), [0] + seed_vals.tolist(), context)
    banned = never_sampled(lm, tok)
    generated = []

    for _ in range(max_new):
        hidden = reader.last_hidden()
        logits = softcap(F.linear(hidden, lm.head.weight)).float()
        logits[banned] = NEG_INF
        if adjust_logits is not None:
            adjusted = adjust_logits(generated, logits)
            logits = logits if torch.isinf(adjusted).all() else adjusted
        if rep_penalty != 1.0 and not getattr(adjust_logits, "rep_free", False):
            window = reader.ids[max(0, len(reader.ids) - rep_window):]
            if window:
                penalize_recent(logits, window, rep_penalty)

        token = sample(logits, temperature, top_k, top_p, rng)
        if token == tok.eos_id:
            return
        number = 0
        if token in (tok.num_id, tok.snum_id):
            number = int(lm.num_dec.generate(hidden[None], temperature)[0])
        reader.add(token, number)
        generated.append(token)
        yield token, number