#!/usr/bin/env python3

from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser, FileType, Namespace
from dataclasses import dataclass
from itertools import chain
from typing import Any, Generator, Iterable, Optional

import torch
from tabulate import tabulate, tabulate_formats
from tqdm import tqdm
from transformers import (
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
    M2M100ForConditionalGeneration,
    set_seed,
)
from transformers.generation.utils import GenerateOutput, GenerationMixin

from mbrs import timer


def buffer_lines(input_stream: Iterable[str], buffer_size: int = 64):
    buf: list[str] = []
    for i, line in enumerate(tqdm(input_stream)):
        buf.append(line.strip())
        if (i + 1) % buffer_size == 0:
            yield buf
            buf = []
    if len(buf) > 0:
        yield buf


def get_argparser() -> ArgumentParser:
    parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
    # fmt: off
    parser.add_argument("input", nargs="?", default="-",
                        type=FileType("r", encoding="utf-8"),
                        help="Input file. If not specified, read from stdin.")
    parser.add_argument("--output", "-o", default="-", type=FileType("w"),
                        help="Output file.")
    parser.add_argument("--lprobs", default=None, type=FileType("w"),
                        help="Reference log-probabilities file. "
                        "This option is useful for the model-based estimation.")
    parser.add_argument("--length_normalized_lprobs", default=None, type=FileType("w"),
                        help="Length-normalized reference log-probabilities file. "
                        "This option is useful for the model-based estimation.")
    parser.add_argument("--model", "-m", type=str, default="facebook/m2m100_418M",
                        help="Model name or path.")
    parser.add_argument("--num_candidates", "-n", type=int, default=1,
                        help="Number of candidates to be returned.")
    parser.add_argument("--sampling", "-s", type=str, default="",
                        choices=["eps"],
                        help="Sampling method.")
    parser.add_argument("--beam_size", type=int, default=5,
                        help="Beam size.")
    parser.add_argument("--epsilon", "--eps", "-e", type=float, default=0.02,
                        help="Cutoff parameter for epsilon sampling.")
    parser.add_argument("--lang_pair", "-l", type=str, default="en-de",
                        help="Language name pair. Some models like M2M100 uses this information.")
    parser.add_argument("--max_length", type=int, default=1024,
                        help="Maximum length of an output sentence.")
    parser.add_argument("--min_length", type=int, default=1,
                        help="Minimum length of an output sentence.")
    parser.add_argument("--length_penalty", type=float, default=None,
                        help="Length penalty.")
    parser.add_argument("--batch_size", "-b", type=int, default=8,
                        help="Batch size.")
    parser.add_argument("--sampling_size", type=int, default=8,
                        help="Sampling size in a single inference. "
                        "The model generates this number of samples at a time "
                        "until the total number of samples reaches `--num_candidates`.")
    parser.add_argument("--fp16", action="store_true",
                        help="Use float16.")
    parser.add_argument("--bf16", action="store_true",
                        help="Use bfloat16.")
    parser.add_argument("--cpu", action="store_true",
                        help="Force to use CPU.")
    parser.add_argument("--seed", type=int, default=0,
                        help="Random number seed.")
    parser.add_argument("--quiet", "-q", action="store_true",
                        help="No report statistics.")
    parser.add_argument("--report", default="-", type=FileType("w"),
                        help="Report file.")
    parser.add_argument("--report_format", type=str, default="rounded_outline",
                        choices=tabulate_formats,
                        help="Report runtime statistics.")
    parser.add_argument("--width", "-w", type=int, default=1,
                        help="Number of digits for values of float point.")
    # fmt: on
    return parser


@dataclass
class Sample:
    """A sample generated by a model."""

    text: str
    lprob: Optional[float] = None
    length_normalized_lprob: Optional[float] = None


def memory_efficient_compute_transition_scores(
    model: GenerationMixin,
    sequences: torch.Tensor,
    scores: tuple[torch.Tensor, ...],
    beam_indices: Optional[torch.Tensor] = None,
    normalize_logits: bool = False,
) -> torch.Tensor:
    """
    Computes the transition scores of sequences given the generation scores (and beam indices, if beam search was
    used). This is a convenient method to quicky obtain the scores of the selected tokens at generation time.

    This is another implementation for memory efficiency.

    Parameters:
        model (`GenerationMixin`): Generation model.
        sequences (`torch.LongTensor`):
            The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or
            shorter if all batches finished early due to the `eos_token_id`.
        scores (`tuple(torch.FloatTensor)`):
            Transition scores for each vocabulary token at each generation step. Beam transition scores consisting
            of log probabilities of tokens conditioned on log softmax of previously generated tokens in this beam.
            Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for each generated token),
            with each tensor of shape `(batch_size*num_beams, config.vocab_size)`.
        beam_indices (`torch.LongTensor`, *optional*):
            Beam indices of generated token id at each generation step. `torch.LongTensor` of shape
            `(batch_size*num_return_sequences, sequence_length)`. Only required if a `num_beams>1` at
            generate-time.
        normalize_logits (`bool`, *optional*, defaults to `False`):
            Whether to normalize the logits (which, for legacy reasons, may be unnormalized).

    Return:
        `torch.Tensor`: A `torch.Tensor` of shape `(batch_size*num_return_sequences, sequence_length)` containing
            the transition scores (logits)

    Examples:

    ```python
    >>> from transformers import GPT2Tokenizer, AutoModelForCausalLM
    >>> import numpy as np

    >>> tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
    >>> tokenizer.pad_token_id = tokenizer.eos_token_id
    >>> inputs = tokenizer(["Today is"], return_tensors="pt")

    >>> # Example 1: Print the scores for each token generated with Greedy Search
    >>> outputs = model.generate(**inputs, max_new_tokens=5, return_dict_in_generate=True, output_scores=True)
    >>> transition_scores = memory_efficient_compute_transition_scores(
    ...     model, outputs.sequences, outputs.scores, normalize_logits=True
    ... )
    >>> # input_length is the length of the input prompt for decoder-only models, like the GPT family, and 1 for
    >>> # encoder-decoder models, like BART or T5.
    >>> input_length = 1 if model.config.is_encoder_decoder else inputs.input_ids.shape[1]
    >>> generated_tokens = outputs.sequences[:, input_length:]
    >>> for tok, score in zip(generated_tokens[0], transition_scores[0]):
    ...     # | token | token string | log probability | probability
    ...     print(f"| {tok:5d} | {tokenizer.decode(tok):8s} | {score.numpy():.3f} | {np.exp(score.numpy()):.2%}")
    |   262 |  the     | -1.414 | 24.33%
    |  1110 |  day     | -2.609 | 7.36%
    |   618 |  when    | -2.010 | 13.40%
    |   356 |  we      | -1.859 | 15.58%
    |   460 |  can     | -2.508 | 8.14%

    >>> # Example 2: Reconstruct the sequence scores from Beam Search
    >>> outputs = model.generate(
    ...     **inputs,
    ...     max_new_tokens=5,
    ...     num_beams=4,
    ...     num_return_sequences=4,
    ...     return_dict_in_generate=True,
    ...     output_scores=True,
    ... )
    >>> transition_scores = memory_efficient_compute_transition_scores(
    ...     model, outputs.sequences, outputs.scores, outputs.beam_indices, normalize_logits=False
    ... )
    >>> # If you sum the generated tokens' scores and apply the length penalty, you'll get the sequence scores.
    >>> # Tip 1: recomputing the scores is only guaranteed to match with `normalize_logits=False`. Depending on the
    >>> # use case, you might want to recompute it with `normalize_logits=True`.
    >>> # Tip 2: the output length does NOT include the input length
    >>> output_length = np.sum(transition_scores.numpy() < 0, axis=1)
    >>> length_penalty = model.generation_config.length_penalty
    >>> reconstructed_scores = transition_scores.sum(axis=1) / (output_length**length_penalty)
    >>> print(np.allclose(outputs.sequences_scores, reconstructed_scores))
    True
    ```"""
    # 1. In absence of `beam_indices`, we can assume that we come from e.g. greedy search, which is equivalent
    # to a beam search approach were the first (and only) beam is always selected
    if beam_indices is None:
        beam_indices = torch.arange(scores[0].shape[0]).view(-1, 1).to(sequences.device)
        beam_indices = beam_indices.expand(-1, len(scores))

    # 2. Optionally normalize the logits (across the vocab dimension)
    if normalize_logits:
        scores = tuple(s.log_softmax(dim=1) for s in scores)

    # 3. cut beam_indices to longest beam length
    beam_indices_mask = beam_indices < 0
    max_beam_length = (1 - beam_indices_mask.long()).sum(-1).max()
    beam_indices = beam_indices.clone()[:, :max_beam_length]
    beam_indices_mask = beam_indices_mask[:, :max_beam_length]

    # 4. Set indices of beams that finished early to 0; such indices will be masked correctly afterwards
    beam_indices[beam_indices_mask] = 0

    # 5. multiply beam_indices with vocab size to gather correctly from scores
    beam_sequence_indices = beam_indices * model.config.vocab_size

    # 6. Define which indices contributed to scores
    cut_idx = sequences.shape[-1] - max_beam_length
    indices = sequences[:, cut_idx:] + beam_sequence_indices

    # 7. Compute scores
    transition_scores = torch.zeros(
        indices.size(), dtype=torch.float32, device=scores[0].device
    )
    for step, step_scores in enumerate(scores):
        transition_scores[:, step] = step_scores.view(-1).gather(0, indices[:, step])

    # 8. Mask out transition_scores of beams that stopped early
    transition_scores[beam_indices_mask] = 0

    return transition_scores


def main(args: Namespace) -> None:
    set_seed(args.seed)

    src_lang, tgt_lang = tuple(args.lang_pair.split("-"))
    tokenizer = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForSeq2SeqLM.from_pretrained(args.model)

    model.eval()
    for param in model.parameters():
        param.requires_grad = False
    if torch.cuda.is_available() and not args.cpu:
        if args.fp16:
            model.half()
        elif args.bf16:
            model.bfloat16()
        model.cuda()

    generation_kwargs = {
        "max_length": args.max_length,
        "min_length": args.min_length,
        "return_dict_in_generate": True,
    }
    length_penalty = getattr(model.generation_config, "length_penalty", 1.0)
    if args.length_penalty is not None:
        generation_kwargs["length_penalty"] = args.length_penalty
        length_penalty = args.length_penalty

    if args.lprobs is not None or args.length_normalized_lprobs is not None:
        generation_kwargs["output_scores"] = True

    if isinstance(model, M2M100ForConditionalGeneration):
        tokenizer.src_lang = src_lang
        generation_kwargs["forced_bos_token_id"] = tokenizer.get_lang_id(tgt_lang)

    if args.sampling == "eps":
        generation_kwargs["do_sample"] = True
        generation_kwargs["epsilon_cutoff"] = args.epsilon
        generation_kwargs["num_beams"] = 1
        generation_kwargs["early_stopping"] = False
    else:
        generation_kwargs["num_beams"] = max(args.beam_size, args.num_candidates)

    def decode(
        inputs: list[str], num_candidates: int, generation_kwargs: dict[str, Any]
    ) -> list[Sample]:
        model_inputs = tokenizer(inputs, return_tensors="pt", padding=True).to(
            device=model.device
        )
        with timer.measure("generate"):
            model_outputs: GenerateOutput = model.generate(
                **model_inputs, **generation_kwargs, num_return_sequences=num_candidates
            )

        if model_outputs.scores is None:
            return [
                Sample(s)
                for s in tokenizer.batch_decode(
                    model_outputs.sequences, skip_special_tokens=True
                )
            ]

        sequences = model_outputs.sequences
        scores = model_outputs.scores
        max_length = len(scores)
        if hasattr(tokenizer, "pad_token_id"):
            pad_token_id = tokenizer.pad_token_id
            for s in scores:
                s[:, pad_token_id] = float("-inf")
            scores = tuple(s.log_softmax(dim=1) for s in scores)
            for s in scores:
                s[:, pad_token_id] = 0.0
            sequence_lengths = max_length - (sequences.eq(pad_token_id)).sum(dim=-1)
        else:
            scores = tuple(s.log_softmax(dim=1) for s in scores)
            sequence_lengths = torch.full(
                (sequences.size(0),), fill_value=max_length, device=sequences.device
            )

        transition_scores = memory_efficient_compute_transition_scores(
            model,
            sequences,
            scores,
            beam_indices=model_outputs.get("beam_indices", None),
        )
        lprobs = transition_scores.sum(dim=-1)
        length_normalized_lprobs = lprobs / (sequence_lengths**length_penalty)
        return [
            Sample(text, lprob=lprob, length_normalized_lprob=length_normalized_lprob)
            for text, lprob, length_normalized_lprob in zip(
                tokenizer.batch_decode(sequences, skip_special_tokens=True),
                lprobs.cpu().tolist(),
                length_normalized_lprobs.cpu().tolist(),
            )
        ]

    def generate(inputs: list[str]) -> Generator[Sample, None, None]:
        if (
            not generation_kwargs.get("do_sample", False)
            or generation_kwargs.get("num_beams", 1) != 1
        ):
            yield from decode(inputs, args.num_candidates, generation_kwargs)
        else:
            samples: list[list[Sample]] = [[] for _ in range(args.batch_size)]
            for n in range(0, args.num_candidates, args.sampling_size):
                sampling_size = min(args.sampling_size, args.num_candidates - n)
                shards = decode(inputs, sampling_size, generation_kwargs)
                for i in range(args.batch_size):
                    samples[i] += shards[i * sampling_size : (i + 1) * sampling_size]
            yield from chain.from_iterable(samples)

    num_sentences = 0
    for lines in buffer_lines(args.input, buffer_size=args.batch_size):
        for sample in generate(lines):
            print(sample.text.strip(), file=args.output)
            if sample.lprob is not None and args.lprobs is not None:
                print(str(sample.lprob), file=args.lprobs)
            if (
                sample.length_normalized_lprob is not None
                and args.length_normalized_lprobs is not None
            ):
                print(
                    str(sample.length_normalized_lprob),
                    file=args.length_normalized_lprobs,
                )
            num_sentences += 1

    if not args.quiet:
        statistics = timer.aggregate().result(num_sentences)
        table = tabulate(
            statistics,
            headers="keys",
            tablefmt=args.report_format,
            floatfmt=f".{args.width}f",
        )
        print(table, file=args.report)


def cli_main():
    args = get_argparser().parse_args()
    main(args)


if __name__ == "__main__":
    cli_main()
