"""Iterative left-to-right beam search over the masked positions."""
from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional, Tuple

import torch

from .text_utils import clean_generated_text, split_sentences, strip_sentiment_marker


@dataclass
class Beam:
    text: str
    score: float


def _terminal_token_ids(tokenizer) -> set:
    ids = set()
    for punct in (".", "!", "?"):
        ids.update(tokenizer.encode(punct, add_special_tokens=False))
    return ids


@torch.inference_mode()
def fill_masks_beam_search(
    model,
    tokenizer,
    masked_text: str,
    sentence_index: int,
    top_k: int = 40,
    beam_size: int = 2,
    max_iters: int = 30,
    alpha: float = 0.7,
    gamma: float = 0.05,
    temperature: float = 0.8,
    min_words: int = 3,
) -> Tuple[List[Beam], str]:
    """Fill masks one at a time, keeping ``beam_size`` hypotheses alive.

    ``alpha`` is the length-normalisation exponent, ``gamma`` penalises beams
    that reuse a token already proposed this step, and a beam finishes early
    once it emits terminal punctuation (after at least ``min_words`` fills).

    Returns the ranked beams over the *whole* passage plus the rewritten
    sentence extracted from the best beam.
    """
    device = next(model.parameters()).device
    mask_token_id = tokenizer.mask_token_id
    pad_token_id = tokenizer.pad_token_id
    if pad_token_id is None:
        pad_token_id = tokenizer.eos_token_id
    terminal_tokens = _terminal_token_ids(tokenizer)

    encoded = tokenizer(masked_text, return_tensors="pt").to(device)
    # (ids, cumulative score, tokens filled, finished, last token)
    beams: List[tuple] = [(encoded.input_ids[0], 0.0, 0, False, None)]

    for _ in range(max_iters):
        active = [i for i, b in enumerate(beams) if (b[0] == mask_token_id).any() and not b[3]]
        if not active:
            break

        active_beams = [beams[i] for i in active]
        batch_ids = torch.stack([b[0] for b in active_beams]).to(device)
        batch_scores = torch.tensor([b[1] for b in active_beams], device=device)
        batch_filled = torch.tensor([b[2] for b in active_beams], device=device)
        mask_positions = (batch_ids == mask_token_id).int().argmax(dim=1)

        logits = model(input_ids=batch_ids).logits
        mask_logits = logits[torch.arange(batch_ids.shape[0]), mask_positions] / temperature

        probs = torch.softmax(mask_logits, dim=-1)
        top_probs, top_tokens = torch.topk(probs, top_k)
        top_probs = top_probs / top_probs.sum(dim=-1, keepdim=True)
        log_p = torch.log(top_probs + 1e-10)

        candidates = [b for i, b in enumerate(beams) if i not in active]

        for i in range(len(active_beams)):
            filled = batch_filled[i].item() + 1
            length_norm = ((5 + filled) ** alpha) / ((5 + 1) ** alpha)
            scores = (batch_scores[i] + log_p[i]) / length_norm

            for j in range(top_k):
                token_id = top_tokens[i, j].item()
                # discourage every beam from picking the same continuation
                penalty = sum(1 for c in candidates if c[4] == token_id)
                score = scores[j].item() - gamma * penalty

                new_ids = batch_ids[i].clone()
                new_ids[mask_positions[i]] = token_id

                finished = token_id in terminal_tokens and filled >= min_words
                if finished:
                    new_ids[new_ids == mask_token_id] = pad_token_id

                candidates.append((new_ids, score, filled, finished, token_id))

        beams = sorted(candidates, key=lambda x: x[1], reverse=True)[:beam_size]

    ranked = [Beam(tokenizer.decode(b[0], skip_special_tokens=True), float(b[1])) for b in beams]
    best_text = ranked[0].text if ranked else ""

    sentences_out = split_sentences(best_text)
    if sentence_index < len(sentences_out):
        best_sentence = sentences_out[sentence_index]
    else:
        best_sentence = sentences_out[-1] if sentences_out else best_text

    return ranked, clean_generated_text(strip_sentiment_marker(best_sentence)).strip()