"""Greedy and beam-search decoding (no KV cache: simple and correct; fast enough for small models)."""
import math

import torch


def _encode_source(model, tok, text, device):
    ids = tok.encode(text)[: model.max_seq_len - 2]
    src = torch.tensor([[tok.BOS_ID] + ids + [tok.EOS_ID]], dtype=torch.long, device=device)
    src_mask = model.make_src_mask(src)
    return src, src_mask, model.encode(src, src_mask)


def _max_steps(model, src):
    return max(1, min(model.max_seq_len - 1, int(src.size(1) * 2) + 10))


@torch.no_grad()
def greedy_decode(model, tok, text, device, max_len=None):
    model.eval()
    if not text.strip():
        return ""
    src, src_mask, memory = _encode_source(model, tok, text, device)
    steps = min(max_len or 10 ** 9, _max_steps(model, src))
    ys = torch.tensor([[tok.BOS_ID]], dtype=torch.long, device=device)
    for _ in range(steps):
        logits = model.decode(ys, memory, src_mask)[:, -1, :].float()
        logits[:, tok.PAD_ID] = float("-inf")
        logits[:, tok.BOS_ID] = float("-inf")
        nxt = int(logits.argmax(-1))
        if nxt == tok.EOS_ID:
            break
        ys = torch.cat([ys, torch.tensor([[nxt]], device=device)], dim=1)
    return tok.decode(ys[0, 1:].tolist())


def _length_penalty(length, alpha):
    return ((5.0 + length) / 6.0) ** alpha


@torch.no_grad()
def beam_search_decode(model, tok, text, device, beam_size=5, max_len=None, alpha=0.6):
    model.eval()
    if not text.strip():
        return ""
    src, src_mask, memory = _encode_source(model, tok, text, device)
    steps = min(max_len or 10 ** 9, _max_steps(model, src))

    seqs = torch.tensor([[tok.BOS_ID]], dtype=torch.long, device=device)   # [k, t]
    scores = torch.zeros(1, device=device)                                  # [k]
    finished = []                                                           # (normalised_score, ids)

    for step in range(steps):
        k = seqs.size(0)
        logits = model.decode(seqs, memory.expand(k, -1, -1), src_mask.expand(k, -1, -1, -1))[:, -1, :]
        logp = torch.log_softmax(logits.float(), dim=-1)
        logp[:, tok.PAD_ID] = float("-inf")
        logp[:, tok.BOS_ID] = float("-inf")
        V = logp.size(-1)
        flat = (scores.unsqueeze(1) + logp).view(-1)
        top_scores, top_idx = flat.topk(min(2 * beam_size, flat.numel()))

        new_seqs, new_scores = [], []
        for s, idx in zip(top_scores.tolist(), top_idx.tolist()):
            if not math.isfinite(s):
                continue
            b, t = divmod(idx, V)
            if t == tok.EOS_ID:
                finished.append((s / _length_penalty(step + 1, alpha), seqs[b, 1:].tolist()))
            else:
                new_seqs.append(torch.cat([seqs[b], torch.tensor([t], device=device)]))
                new_scores.append(s)
            if len(new_seqs) == beam_size:
                break
        if not new_seqs:
            break
        seqs = torch.stack(new_seqs)
        scores = torch.tensor(new_scores, device=device)
        if len(finished) >= beam_size:
            # Stop only when even the WORST of the best `beam_size` finished hypotheses already beats
            # the best hypothesis still being extended. (Stopping as soon as `beam_size` hypotheses
            # have finished would let early, low-quality endings crowd out the real best one.)
            finished.sort(key=lambda x: x[0], reverse=True)
            finished = finished[:beam_size]
            best_alive = scores.max().item() / _length_penalty(step + 1, alpha)
            if finished[-1][0] >= best_alive:
                break

    if not finished:  # ran out of steps: fall back to the best unfinished hypotheses
        for s, seq in zip(scores.tolist(), seqs):
            finished.append((s / _length_penalty(seqs.size(1) - 1, alpha), seq[1:].tolist()))
    best = max(finished, key=lambda x: x[0])[1]
    return tok.decode(best)


def translate(model, tok, text, device, beam_size=5, max_len=None):
    """beam_size <= 1 -> greedy."""
    if beam_size and beam_size > 1:
        return beam_search_decode(model, tok, text, device, beam_size=beam_size, max_len=max_len)
    return greedy_decode(model, tok, text, device, max_len=max_len)