"""
Core models for LongBEL
"""
# Copyright (c) Facebook, Inc. and its affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import json
import logging
import os
import pickle
import re
from html import escape

import torch
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from IPython.display import HTML, display
from tqdm.auto import tqdm
from transformers import (
    AutoTokenizer,
    LlamaForCausalLM,
    PretrainedConfig,
)

logger = logging.getLogger(__name__)
logging.basicConfig(
    level=logging.INFO,  # Display INFO and above
    format="%(levelname)s - %(message)s",
)


# Define a simple config class that inherits from PretrainedConfig
class LLamaLongBELConfig(PretrainedConfig):
    model_type = "llama_longbel"

    def __init__(self, **kwargs):
        # Ensure it has llama as base
        kwargs.setdefault("model_type", "llama")
        super().__init__(**kwargs)


def clean_natural(text):
    return (
        text.replace("\xa0", " ")
        .replace("{", "(")
        .replace("}", ")")
        .replace("[", "(")
        .replace("]", ")")
        .replace("\n", " ")
    )


def parse_text(
    data,
    start_entity,
    end_entity,
    start_group,
    end_group,
) -> tuple[list[str], list[str], list[dict[str, str]]]:
    """Create simple (source, target) pairs per entity.

    For each entity in the BigBio page, returns one pair where:
      - source: the sentence text that contains the entity mention
      - target: "<entity> is <annotation>" where <annotation> is the best synonym
        if available (or the normalized id otherwise).
    """
    source_sentences: list[str] = []
    tsv_lines: list[dict[str, str]] = []
    target_texts_dict: dict[tuple[tuple[int, int], ...], str] = {}
    source_texts_dict: dict[tuple[tuple[int, int], ...], str] = {}
    tsv_lines_dict: dict[tuple[tuple[int, int], ...], dict[str, str]] = {}
    all_passages = {}
    for i, passage in enumerate(data.get("passages", [])):
        all_passages[i] = clean_natural(passage["text"][0])
    for passage_id, passage in enumerate(data.get("passages", [])):
        passage_text = passage["text"][0]
        start_offset_passage = passage["offsets"][0][0]
        end_offset_passage = passage["offsets"][0][1]

        passage_text = clean_natural(passage_text)

        # Iterate over entities and emit one pair per entity found in this passage
        for entity in data.get("entities", []):
            # min and max of all entity offsets to get the global span of the entity for filtering sentences
            global_start = min(off[0] for off in entity["offsets"])
            global_end = max(off[1] for off in entity["offsets"])
            # Keep only entities whose start falls inside this passage
            if not (start_offset_passage <= global_start < end_offset_passage):
                continue
            entity_text = " ".join(entity["text"])
            entity_text = clean_natural(entity_text)
            # Define entity group
            group_annotation = entity.get("type")
            # Get all offsets, convert to relative, and filter for this sentence
            relative_entity_spans = []
            for off in entity["offsets"]:
                global_start_off, global_end_off = off
                if not (start_offset_passage <= global_start_off < end_offset_passage):
                    continue

                rel_start_off = global_start_off - start_offset_passage
                rel_end_off = global_end_off - start_offset_passage
                relative_entity_spans.append((rel_start_off, rel_end_off))
            relative_entity_spans.sort(key=lambda x: x[0])

            marked_text = passage_text
            for start_in_sent, end_in_sent in relative_entity_spans:
                marked_text = (
                    marked_text[:start_in_sent]
                    + start_entity
                    + marked_text[start_in_sent:end_in_sent]
                    + end_entity
                    + marked_text[end_in_sent:]
                )

            for other_passage_id, other_passage_text in all_passages.items():
                if other_passage_id < passage_id:
                    marked_text = other_passage_text + "\n" + marked_text
                elif other_passage_id > passage_id:
                    marked_text = marked_text + "\n" + other_passage_text
            # Emit the pair
            doc_id = data.get("id", "")
            tsv_line = {
                "doc_id": doc_id,
                "semantic_group": group_annotation,
                "start_span": global_start,
                "end_span": global_end,
                "mention": entity_text,
            }

            # Keep original entity id if present
            if entity.get("id") is not None:
                tsv_line["mention_id"] = str(entity["id"])

            if entity.get("normalized"):
                tsv_line["gold_concept_code"] = entity["normalized"][0]["db_id"]
                tsv_line["gold_concept_name"] = entity["normalized"][0]["db_match"]

            key = (global_start, global_end, group_annotation)
            tsv_lines_dict[key] = tsv_line
            source_texts_dict[key] = marked_text
            target_entity_text = (
                start_entity
                + entity_text
                + end_entity
                + start_group
                + group_annotation
                + end_group
            )
            target_texts_dict[key] = target_entity_text
    # Sort keys to have a deterministic order
    target_texts = []
    sorted_keys = sorted(
        tsv_lines_dict.keys(),
        key=lambda x: (x[0], x[1], str(x[2]) if len(x) > 2 else ""),
    )
    for entity_idx, entity_span in enumerate(sorted_keys):
        tsv_line = tsv_lines_dict[entity_span]
        if "mention_id" not in tsv_line:
            tsv_line["mention_id"] = f"{data.get('id', '')}.{entity_idx + 1}"
        tsv_lines.append(tsv_line)
        source_sentences.append(source_texts_dict[entity_span])
        target_texts.append(target_texts_dict[entity_span])

    return source_sentences, target_texts, tsv_lines  # type: ignore


def get_prefix_allowed_tokens_fn(
    model,
    sources: list[str],
    sem_groups: list[str],
    multiple_answers: bool = False,
):
    candidates_trie = model.candidate_trie  # type: ignore
    sep_token_id = model.tokenizer.sep_token_id
    eos_token_id = model.tokenizer.eos_token_id
    pad_token_id = model.tokenizer.pad_token_id
    plus_token_id = model.tokenizer.convert_tokens_to_ids("<+>")  # type: ignore
    end_group_token_id = model.tokenizer.convert_tokens_to_ids("}")  # type: ignore

    def prefix_allowed_tokens_fn(batch_id, sent):
        sent = sent.tolist()
        if len(sent) > 1 and sent[-1] in [eos_token_id, pad_token_id, sep_token_id]:
            if sep_token_id:
                return [sep_token_id, pad_token_id, eos_token_id]
            else:
                return [pad_token_id, eos_token_id]

        # Remove the prefix from the sent
        index_sep = len(sent) - 1 - sent[::-1].index(end_group_token_id)
        sent = sent[index_sep:]

        sem_group = sem_groups[batch_id]
        # Remove everything up to last sep_token_id and add prefix
        if multiple_answers and plus_token_id in sent:
            index_plus = len(sent) - 1 - sent[::-1].index(plus_token_id)
            # Start fresh with decoder start
            if index_plus == len(sent) - 1:
                sent = [end_group_token_id]
            # If there are tokens after the last plus_token_id, keep them
            else:
                sent = [end_group_token_id] + sent[index_plus + 1 :]
        trie_out = candidates_trie[
            sem_group  # type: ignore
        ].get(sent)
        if eos_token_id in trie_out:
            if sep_token_id:
                trie_out += [sep_token_id]
            if multiple_answers:
                trie_out += [plus_token_id]
        elif not trie_out:
            if sep_token_id:
                return [sep_token_id, pad_token_id, eos_token_id]
            else:
                return [pad_token_id, eos_token_id]
        return trie_out

    return prefix_allowed_tokens_fn


def add_headers_to_prompt(source: str, target: str, previous_targets: str):
    if not previous_targets:
        previous_targets = "None"
    input_sentence = f"### Context\n{source.rstrip()}\n\n### Previous Normalizations\n{previous_targets.rstrip()}\n\n### Prediction\n{target.rstrip()}"
    return input_sentence


def parse_prediction(
    outputs: list[str],
    sem_groups: list[str],
    text_to_code: dict[str, dict[str, str]] | None = None,
    multiple_answers: bool = False,
) -> tuple[list[str], list[str]]:
    codes = []
    predictions = []
    for output, group in zip(outputs, sem_groups, strict=False):
        splits = output.split("} ")  # type: ignore
        if len(splits) > 1 and splits[-1].strip():
            prediction = splits[-1].strip().replace("<SEP>", "")
            if text_to_code:
                if multiple_answers:
                    prediction_list = prediction.split("<+>")  # type: ignore
                    code_list = set()
                    for pred in prediction_list:
                        code_list.add(text_to_code[group].get(pred.strip(), "NO_CODE"))
                    if len(code_list) > 1 and "NO_CODE" in code_list:
                        code_list.remove("NO_CODE")
                    code = "+".join(code_list)
                else:
                    code = text_to_code[group].get(prediction, "NO_CODE")
            else:
                code = "NO_CODE"
        else:
            print(
                "IndexError: splitting failed or empty prediction, adding empty string as prediction."
            )
            prediction = "NO_PREDICTION"
            code = "NO_CODE"
        codes.append(code)
        predictions.append(prediction)
    return codes, predictions


def compute_score(outputs, tokenizer, prefix_len=0):
    sequences = outputs.sequences  # (N, seq_len)
    scores = outputs.scores  # list length T = # generated tokens

    N, total_len = sequences.shape
    T = len(scores)

    # keep only the generated part (completion)
    sequences = sequences[:, prefix_len : prefix_len + T]

    # Make sure score is not longer than sequences
    if len(scores) > sequences.size(1):
        scores = scores[: sequences.size(1)]

    # Compute as usual but now only for completion tokens
    mask = (
        (sequences != tokenizer.pad_token_id)
        & (sequences != tokenizer.eos_token_id)
        & (sequences != tokenizer.bos_token_id)
    )

    # log-prob for each generated token
    logprob_steps = []
    for t, logits in enumerate(scores):
        log_probs_t = F.log_softmax(logits, dim=-1)
        token_t = sequences[:, t]
        idx = torch.arange(N)
        logprob_steps.append(log_probs_t[idx, token_t])

    logprobs = torch.stack(logprob_steps, dim=1)
    logprobs.masked_fill_(~mask, 0)

    lengths = mask.sum(dim=1).clamp(min=1)
    confidence = torch.exp(logprobs.sum(dim=1) / lengths)

    return confidence.tolist()


def skip_undesired_tokens(outputs, tokenizer):
    sep_token = "<SEP>"
    plus_token = "<+>"
    # Build the list of special tokens to remove
    tokens_to_remove = tokenizer.all_special_tokens[:2]

    cleaned_outputs = []
    for sequence in outputs:
        # Remove undesired special tokens
        for token in tokens_to_remove:
            sequence = sequence.replace(token, "")

        # Remove spaces *immediately* after the sep_token adn plus_token (e.g. "<sep>  text" → "<sep>text")
        sequence = re.sub(rf"({re.escape(plus_token)})\s+", r"\1", sequence)
        sequence = re.sub(rf"({re.escape(sep_token)})\s+", r"\1", sequence)

        cleaned_outputs.append(sequence.strip())

    return cleaned_outputs


def _score_to_rgb(score: float) -> tuple[int, int, int]:
    clipped_score = max(0.0, min(1.0, score))
    red = 255
    channel = int(255 * (1.0 - clipped_score))
    return red, channel, channel


def _build_ansi_saliency_text(
    token_texts: list[str], saliency_scores: list[float]
) -> str:
    chunks = []
    for token_text, score in zip(token_texts, saliency_scores, strict=False):
        red, green, blue = _score_to_rgb(score)
        chunks.append(f"\x1b[48;2;{red};{green};{blue}m{token_text}\x1b[0m")
    return "".join(chunks)


def _build_html_saliency_text(
    token_texts: list[str], saliency_scores: list[float]
) -> str:
    chunks = []
    for token_text, score in zip(token_texts, saliency_scores, strict=False):
        red, green, blue = _score_to_rgb(score)
        chunks.append(
            f'<span style="background-color: rgb({red}, {green}, {blue});">{escape(token_text)}</span>'
        )
    return "".join(chunks)


class LLamaLongBEL(LlamaForCausalLM):
    config_class = LLamaLongBELConfig

    def __init__(self, config, *args, **kwargs):
        # Initialize the parent LlamaForCausalLM
        super().__init__(config, *args, **kwargs)

        self.text_to_code = None
        self.candidate_trie = None
        self.tokenizer = None

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path,
        *args,
        text_to_code_path=None,
        candidate_trie_path=None,
        **kwargs,
    ):
        # Remove custom kwargs before passing to parent
        custom_kwargs = {
            "text_to_code_path": text_to_code_path,
            "candidate_trie_path": candidate_trie_path,
        }

        # Call parent's from_pretrained
        model = super().from_pretrained(
            pretrained_model_name_or_path,
            *args,
            **{k: v for k, v in kwargs.items() if k not in custom_kwargs},
        )

        # Set up tokenizer
        model.tokenizer = AutoTokenizer.from_pretrained(
            pretrained_model_name_or_path, use_fast=True
        )
        model.tokenizer.padding_side = "left"

        # Load text_to_code
        text_to_code_file_local = (
            text_to_code_path
            if text_to_code_path is not None
            else os.path.join(pretrained_model_name_or_path, "text_to_code.json")
        )
        try:
            if os.path.exists(text_to_code_file_local):
                with open(text_to_code_file_local, encoding="utf-8") as f:
                    model.text_to_code = json.load(f)
                logger.info(
                    f"Loaded text_to_code.json from local path: {text_to_code_file_local}"
                )
            else:
                text_to_code_path_hf = hf_hub_download(
                    repo_id=pretrained_model_name_or_path,
                    filename="text_to_code.json",
                )
                with open(text_to_code_path_hf, encoding="utf-8") as f:
                    model.text_to_code = json.load(f)
                logger.info(
                    f"Loaded text_to_code.json from HF Hub: {text_to_code_path_hf}"
                )
        except Exception:
            logger.warning("text_to_code.json not found (local or HF hub)")
            model.text_to_code = None

        # Load candidate_trie
        candidate_trie_file_local = (
            candidate_trie_path
            if candidate_trie_path is not None
            else os.path.join(pretrained_model_name_or_path, "candidate_trie.pkl")
        )
        try:
            if os.path.exists(candidate_trie_file_local):
                with open(candidate_trie_file_local, "rb") as f:
                    model.candidate_trie = pickle.load(f)
                logger.info(
                    f"Loaded candidate_trie.pkl from local path: {candidate_trie_file_local}"
                )
            else:
                candidate_trie_path_hf = hf_hub_download(
                    repo_id=pretrained_model_name_or_path,
                    filename="candidate_trie.pkl",
                )
                with open(candidate_trie_path_hf, "rb") as f:
                    model.candidate_trie = pickle.load(f)
                logger.info(
                    f"Loaded candidate_trie.pkl from HF Hub: {candidate_trie_path_hf}"
                )
        except Exception:
            logger.warning("candidate_trie.pkl not found (local or HF hub)")
            model.candidate_trie = None

        return model

    def _compute_gradient_saliency(
        self,
        input_sentences: list[str],
        generated_sequences: torch.Tensor,
        num_beams: int,
        prefix_len: int,
    ) -> list[dict[str, object]]:
        if not input_sentences:
            return []

        top_sequence_indices = (
            torch.arange(
                len(input_sentences),
                device=generated_sequences.device,
            )
            * num_beams
        )
        top_sequences = generated_sequences.index_select(0, top_sequence_indices)

        attention_mask = (top_sequences != self.tokenizer.pad_token_id).long()  # type: ignore
        input_embeddings = self.get_input_embeddings()(top_sequences).detach()  # type: ignore

        next_tokens = top_sequences[:, 1:]
        output_token_mask = torch.zeros_like(next_tokens, dtype=torch.bool)
        if prefix_len > 0:
            output_token_mask[:, prefix_len - 1 :] = True

        valid_token_mask = output_token_mask & (
            (next_tokens != self.tokenizer.pad_token_id)  # type: ignore
            & (next_tokens != self.tokenizer.eos_token_id)  # type: ignore
            & (next_tokens != self.tokenizer.bos_token_id)  # type: ignore
        )

        def _objective_from_embeddings(embeddings: torch.Tensor) -> torch.Tensor:
            forward_outputs = self(  # type: ignore
                inputs_embeds=embeddings,
                attention_mask=attention_mask,
                use_cache=False,
                return_dict=True,
            )
            logits = forward_outputs.logits[:, :-1, :]
            log_probs = F.log_softmax(logits, dim=-1)
            token_log_probs = log_probs.gather(
                dim=-1,
                index=next_tokens.unsqueeze(-1),
            ).squeeze(-1)
            return token_log_probs.masked_select(valid_token_mask).sum()

        simple_embeddings = input_embeddings.detach()
        simple_embeddings.requires_grad_(True)
        self.zero_grad(set_to_none=True)  # type: ignore
        with torch.enable_grad():
            objective = _objective_from_embeddings(simple_embeddings)
        gradients = torch.autograd.grad(
            outputs=objective,
            inputs=simple_embeddings,
            retain_graph=False,
            create_graph=False,
        )[0]
        token_importance = gradients.norm(p=2, dim=-1)
        saliency_maps = []
        sequence_len = top_sequences.size(1)
        prompt_positions = torch.arange(sequence_len, device=top_sequences.device)
        prompt_mask = (prompt_positions.unsqueeze(0) < prefix_len) & (
            top_sequences != self.tokenizer.pad_token_id  # type: ignore
        )

        for sequence_ids, importance_scores, sentence, mask in zip(
            top_sequences, token_importance, input_sentences, prompt_mask, strict=False
        ):
            selected_ids = sequence_ids[mask]
            selected_scores = importance_scores[mask]

            if selected_scores.numel() == 0:
                saliency_maps.append({
                    "input_sentence": sentence,
                    "token_ids": [],
                    "token_strings": [],
                    "saliency_scores": [],
                    "saliency_ansi": "",
                    "saliency_html": "",
                })
                continue

            max_score = selected_scores.max().clamp(min=1e-12)
            normalized_scores = (selected_scores / max_score).tolist()
            selected_ids_list = selected_ids.tolist()
            token_strings = [
                self.tokenizer.decode(  # type: ignore
                    [token_id],
                    skip_special_tokens=False,
                    clean_up_tokenization_spaces=False,
                )
                for token_id in selected_ids_list
            ]

            saliency_maps.append({
                "input_sentence": sentence,
                "token_ids": selected_ids_list,
                "token_strings": token_strings,
                "saliency_scores": normalized_scores,
                "saliency_ansi": _build_ansi_saliency_text(
                    token_strings,
                    normalized_scores,
                ),
                "saliency_html": _build_html_saliency_text(
                    token_strings,
                    normalized_scores,
                ),
            })

        return saliency_maps

    def display_saliency_map(self, saliency_map):
        saliency_html = re.sub(
            r"<span[^>]*>\s*&lt;\|begin_of_text\|&gt;\s*</span>",
            "",
            saliency_map["saliency_html"],
            count=1,
        )
        pred_name = escape(str(saliency_map.get("pred_concept_name", "")))
        pred_code = escape(str(saliency_map.get("pred_concept_code", "")))
        full_html = f"""
        <div style="
            font-family: Times New Roman, Times, serif, monospace;
            font-size: 18px;
            line-height: 1.6;
            white-space: pre-wrap;
            border: 1px solid #ddd;
            border-radius: 8px;
            padding: 12px;
            background: #fafafa;
        ">{saliency_html} → {pred_name} ({pred_code})</div>
        """
        display(HTML(full_html))

    def predict_batch(
        self,
        all_outputs,
        batch_size,
        input_sentences,
        sem_groups,
        mentions,
        mentions_id,
        doc_ids,
        start_spans,
        end_spans,
        gold_concept_codes,
        gold_concept_names,
        constrained,
        multiple_answers,
        num_beams,
        with_saliency_maps: bool = False,
        **kwargs,
    ):
        input_args = {
            k: v.to(self.device)  # type: ignore
            for k, v in self.tokenizer.batch_encode_plus(  # type: ignore
                input_sentences, padding="longest", return_tensors="pt"
            ).items()
        }

        # Constrained decoding
        prefix_allowed_tokens_fn = None
        if constrained:
            if self.candidate_trie is None:  # type: ignore
                raise ValueError(
                    "candidate_trie is not loaded in the model. Use constrained=False."
                )
            prefix_allowed_tokens_fn = get_prefix_allowed_tokens_fn(
                model=self,
                sources=input_sentences,
                sem_groups=sem_groups,
                multiple_answers=multiple_answers,
            )
        if self.tokenizer.sep_token_id:  # type: ignore
            eos_token_id = self.tokenizer.sep_token_id  # type: ignore
        else:
            eos_token_id = self.tokenizer.eos_token_id  # type: ignore
        outputs = self.generate(  # type: ignore
            **input_args,
            max_new_tokens=128,
            num_beams=num_beams,
            num_return_sequences=num_beams,
            output_scores=True,
            return_dict_in_generate=True,
            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
            eos_token_id=eos_token_id,  # type: ignore
            **kwargs,
        )
        decoded_sequences = self.tokenizer.batch_decode(  # type: ignore
            outputs.sequences,  # type: ignore
            skip_special_tokens=False,
            clean_up_tokenization_spaces=True,
        )
        cleaned_output_sequences = skip_undesired_tokens(
            decoded_sequences,
            self.tokenizer,  # type: ignore
        )

        prefix_len = input_args["input_ids"].size(1)

        base_sem_groups = sem_groups.copy()
        base_mentions = mentions.copy()
        base_mentions_id = mentions_id.copy()
        base_doc_ids = doc_ids.copy()
        base_start_spans = start_spans.copy()
        base_end_spans = end_spans.copy()
        base_gold_concept_codes = gold_concept_codes.copy()
        base_gold_concept_names = gold_concept_names.copy()

        # Duplicate sem_groups and mentions for each beam
        sem_groups = [x for x in sem_groups for _ in range(num_beams)]
        mentions = [x for x in mentions for _ in range(num_beams)]
        mentions_id = [x for x in mentions_id for _ in range(num_beams)]
        gold_concept_codes = [x for x in gold_concept_codes for _ in range(num_beams)]  # type: ignore
        gold_concept_names = [x for x in gold_concept_names for _ in range(num_beams)]  # type: ignore
        start_spans = [x for x in start_spans for _ in range(num_beams)]
        end_spans = [x for x in end_spans for _ in range(num_beams)]
        doc_ids = [x for x in doc_ids for _ in range(num_beams)]
        # Parse predictions
        pred_concept_codes, pred_concept_names = parse_prediction(
            cleaned_output_sequences,
            sem_groups,
            self.text_to_code,  # type: ignore
            multiple_answers=multiple_answers,
        )
        scores = compute_score(
            outputs,
            self.tokenizer,  # type: ignore
            prefix_len=prefix_len,
        )
        beam_scores = [
            float(torch.exp(s)) if num_beams > 1 else float("nan")
            for s in (
                outputs.sequences_scores  # type: ignore
                if num_beams > 1
                else [torch.tensor(float("nan"))] * len(scores)
            )
        ]
        all_outputs.extend([
            {
                "mention": mention,
                "doc_id": doc_id,
                "mention_id": mention_id,
                "start_span": start_span,
                "end_span": end_span,
                "semantic_group": group,
                "gold_concept_code": gold_concept_code,
                "gold_concept_name": gold_concept_name,
                "pred_concept_name": pred_concept_name,
                "pred_concept_code": pred_concept_code,
                "score": score,
                "beam_score": beam_score,
                "rank": rank + 1,
            }
            for score, beam_score, pred_concept_code, pred_concept_name, mention, doc_id, mention_id, start_span, end_span, group, gold_concept_code, gold_concept_name, rank in zip(
                scores,
                beam_scores,
                pred_concept_codes,
                pred_concept_names,
                mentions,
                doc_ids,
                mentions_id,
                start_spans,
                end_spans,
                sem_groups,
                gold_concept_codes,
                gold_concept_names,
                list(range(num_beams)) * batch_size,
                strict=False,
            )
        ])

        saliency_maps = []
        if with_saliency_maps:
            saliency_maps = self._compute_gradient_saliency(
                input_sentences=input_sentences,
                generated_sequences=outputs.sequences,  # type: ignore
                num_beams=num_beams,
                prefix_len=prefix_len,
            )
            for idx, saliency_map in enumerate(saliency_maps):
                top_prediction_index = idx * num_beams
                saliency_map.update({
                    "mention": base_mentions[idx],
                    "doc_id": base_doc_ids[idx],
                    "mention_id": base_mentions_id[idx],
                    "start_span": base_start_spans[idx],
                    "end_span": base_end_spans[idx],
                    "semantic_group": base_sem_groups[idx],
                    "gold_concept_code": base_gold_concept_codes[idx],
                    "gold_concept_name": base_gold_concept_names[idx],
                    "pred_concept_name": pred_concept_names[top_prediction_index],
                    "pred_concept_code": pred_concept_codes[top_prediction_index],
                    "score": scores[top_prediction_index],
                    "rank": 1,
                })

        print(f"Sampling completed. Generated {len(all_outputs)} predictions.")
        return all_outputs, cleaned_output_sequences, saliency_maps

    def sample(
        self,
        bigbio_pages: list[dict],  # type: ignore
        num_beams: int = 5,
        constrained: bool = True,
        with_saliency_maps: bool = False,
        multiple_answers: bool = False,
        batch_size: int = 8,
        start_entity: str = "[",
        end_entity: str = "]",
        start_group: str = "{",
        end_group: str = "}",
        show_progress: bool = True,
        **kwargs,
    ) -> (
        list[dict[str, object]]
        | tuple[list[dict[str, object]], list[dict[str, object]]]
    ):

        print(
            f"Starting sampling on {len(bigbio_pages)} pages, constrained={constrained}, beams={num_beams}, batch_size={batch_size})"
        )

        def _progress(iterable, desc: str, total: int | None = None, show: bool = True):
            if show:
                return tqdm(iterable, desc=desc, total=total)
            return iterable

        all_outputs = []
        all_sources = []
        all_targets = []
        all_entities_info = []
        for data in bigbio_pages:
            sources, targets, entities_info = parse_text(
                data=data,
                start_entity=start_entity,
                end_entity=end_entity,
                start_group=start_group,
                end_group=end_group,
            )
            all_sources.append(sources)
            all_targets.append(targets)
            all_entities_info.append(entities_info)

        def _build_sequential_batches():
            # Keep per-page order while still processing multiple pages per batch.
            page_positions = [0] * len(all_sources)
            next_page_idx = 0
            active_pages = []
            batches = []

            while active_pages or next_page_idx < len(all_sources):
                while len(active_pages) < batch_size and next_page_idx < len(
                    all_sources
                ):
                    if len(all_sources[next_page_idx]) > 0:
                        active_pages.append(next_page_idx)
                    next_page_idx += 1

                if not active_pages:
                    break

                batch = []
                next_active_pages = []
                for page_idx in active_pages:
                    item_idx = page_positions[page_idx]
                    batch.append((
                        all_sources[page_idx][item_idx],
                        all_targets[page_idx][item_idx],
                        all_entities_info[page_idx][item_idx],
                    ))
                    page_positions[page_idx] += 1
                    if page_positions[page_idx] < len(all_sources[page_idx]):
                        next_active_pages.append(page_idx)

                batches.append(batch)
                active_pages = next_active_pages

            return batches

        all_batches = _build_sequential_batches()

        print(
            f"Input preparation completed. Running generation on {len(all_batches)} batches."
        )

        all_outputs = []
        all_saliency_maps = []
        batch_previous_targets = {}
        for batch in _progress(
            all_batches,
            desc="Processing batches",
            total=len(all_batches),
            show=show_progress,
        ):
            input_sentences = []
            sem_groups = []
            mentions = []
            doc_ids = []
            mentions_id = []
            gold_concept_codes = []
            gold_concept_names = []
            start_spans = []
            end_spans = []
            for source, target, entity in batch:
                doc_id = entity["doc_id"]
                if doc_id not in batch_previous_targets:
                    batch_previous_targets[doc_id] = ""
                previous_targets = batch_previous_targets.get(doc_id)

                input_sentences.append(
                    add_headers_to_prompt(
                        source,
                        target,
                        previous_targets,  # type: ignore
                    )
                )
                sem_groups.append(entity["semantic_group"])
                mentions.append(entity["mention"])
                doc_ids.append(doc_id)
                mentions_id.append(entity["mention_id"])
                start_spans.append(entity["start_span"])
                end_spans.append(entity["end_span"])
                gold_concept_codes.append(entity.get("gold_concept_code", None))  # type: ignore
                gold_concept_names.append(entity.get("gold_concept_name", None))  # type: ignore
            all_outputs, cleaned_output_sequences, batch_saliency_maps = (
                self.predict_batch(
                    all_outputs=all_outputs,
                    batch_size=batch_size,
                    input_sentences=input_sentences,
                    sem_groups=sem_groups,
                    mentions=mentions,
                    mentions_id=mentions_id,
                    doc_ids=doc_ids,
                    start_spans=start_spans,
                    end_spans=end_spans,
                    gold_concept_codes=gold_concept_codes,
                    gold_concept_names=gold_concept_names,
                    constrained=constrained,
                    multiple_answers=multiple_answers,
                    num_beams=num_beams,
                    with_saliency_maps=with_saliency_maps,
                    **kwargs,
                )
            )
            if with_saliency_maps:
                all_saliency_maps.extend(batch_saliency_maps)
            for i, doc_id in enumerate(doc_ids):
                clean_sentence = cleaned_output_sequences[num_beams * i]
                clean_sentence = start_entity + clean_sentence.split(start_entity)[-1]
                clean_sentence = clean_sentence.rstrip() + "\n"
                batch_previous_targets[doc_id] += clean_sentence

        if with_saliency_maps:
            return all_outputs, all_saliency_maps  # type: ignore
        return all_outputs  # type: ignore

    def encode(self, sentence):
        return self.tokenizer.encode(sentence, return_tensors="pt")[0]  # type: ignore