"""Run the fine-tuned Whisper over a predicted mel.

Uses transformers' own `generate()` rather than a hand-rolled decoder. An
earlier ONNX export with custom beam search produced single punctuation marks
and empty strings where `generate()` produced fluent sentences from the
identical mel -- the model was being silently crippled by the decoding rather
than by the input.

The mel contract is the subtle part. The checkpoint keeps Whisper's standard
**80-channel** input even though it was fine-tuned on 12-bin prosody-predicted
mels: `num_mel_bins` in the training config truncates incoming mels, it never
resizes conv1. The training collator reconciles the two by zero-padding, and
this module reproduces that exactly:

    conformer -> [T, 12] -> transpose -> [12, T] -> zero-pad -> [80, 3000]

Bins 12-79 are zeros. Getting this wrong fails at the first conv layer.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np

# Whisper's encoder window: 3000 mel frames at 100 Hz = 30 s.
ENCODER_FRAMES = 3000
ENCODER_MEL_BINS = 80

# The decoder prefix the checkpoint was fine-tuned behind.
#
# Its own forced_decoder_ids is [[1, None], [2, <|transcribe|>]] -- position 1,
# the language slot, is left unset. Decoding without pinning it does not make
# the model choose a language: it skips the slot and starts from
# <|startoftranscript|> alone, so the prefix is two tokens short of what
# training used.
#
# The published inference runs set language: en, and their recorded
# snapshot_meta.json confirms the full prefix
# [50258, 50259, 50359, 50363] = <|sot|> <|en|> <|transcribe|> <|notimestamps|>.
# Matching it is the point: a different prefix is a different model condition,
# so any comparison against those runs is confounded without it.
DEFAULT_LANGUAGE = "en"

# Generation limits the published inference runs used.
DEFAULT_MAX_NEW_TOKENS = 225
DEFAULT_NO_REPEAT_NGRAM = 3

# 10, matching the published inference config for this checkpoint.
#
# Wider is cheap and was tried: 385 ms at 10 beams against 634 ms at 40 on a
# healthy phrase, four times the hypotheses for a quarter more latency. It
# made the transcriptions worse -- a wider beam searches harder for the
# highest-likelihood sequence, and on a signal as weak as a drawn contour
# that is the YouTube-caption boilerplate the checkpoint overlearned. More
# beams find the attractor more reliably rather than finding better words.
DEFAULT_BEAMS = 10

# Phrases Whisper falls back on when the input carries no usable speech.
#
# Trained on YouTube captions, it learned what tends to accompany
# unintelligible audio: "Thank you." for silence, and shouted boilerplate
# from auto-generated or burned-in captions. Feeding it pure noise returns
# "I'M NOT GOING TO SAY ANYTHING" in three of the top four beams; feeding it
# a real drawing returns it in none. Naming it in the UI keeps it from
# reading as a transcription of the drawing.
HALLUCINATION_MARKERS = (
    "i'm not going to say anything",
    "thank you for watching",
    "subscribe",
    "thanks for watching",
)


def looks_like_fallback(text: str) -> bool:
    """Whether a hypothesis is the model's no-speech boilerplate."""
    lowered = text.strip().lower()
    if not lowered:
        return False
    if len(lowered) > 3 and text.strip().isupper():
        return True
    return any(marker in lowered for marker in HALLUCINATION_MARKERS)


# Measured across the templates: real speech runs about 4.3 words/sec, p90
# 5.6. The drawing is the only evidence about how long the utterance is, so
# the token budget is derived from its duration rather than left at a flat
# 48 -- otherwise a half-second scribble can come back as a paragraph, which
# is the language model talking rather than a transcription.
WORDS_PER_SECOND = 5.6  # p90, so an unusually fast phrase is not truncated
TOKENS_PER_WORD = 2.5  # Whisper splits a word into roughly this many tokens
MIN_TOKENS = 4


def pad_mel_for_whisper(
    mel: np.ndarray,
    mel_bins: int = ENCODER_MEL_BINS,
    frames: int = ENCODER_FRAMES,
) -> np.ndarray:
    """[T, bins] from the conformer -> [bins_padded, frames] for Whisper.

    Transposes to Whisper's [bins, time] layout, then zero-pads both axes.
    This mirrors the training collator's `pad_mel_to_expected` branch, so the
    model sees the same shape it was fine-tuned on.
    """
    mel = np.asarray(mel, dtype=np.float32)
    if mel.ndim != 2:
        raise ValueError(f"expected a 2-D mel, got {mel.shape}")

    # The conformer emits [T, bins]; Whisper wants [bins, T].
    mel = mel.T

    out = np.zeros((mel_bins, frames), dtype=np.float32)
    bins = min(mel.shape[0], mel_bins)
    time = min(mel.shape[1], frames)
    out[:bins, :time] = mel[:bins, :time]
    return out


class Whisper:
    """Transcribes a predicted mel with the fine-tuned checkpoint."""

    def __init__(
        self,
        model_dir: Path | str,
        device: str | None = None,
        num_beams: int = DEFAULT_BEAMS,
        half: bool = False,
        tokenizer_dir: Path | str | None = None,
    ) -> None:
        import torch
        from transformers import (WhisperForConditionalGeneration,
                                  WhisperTokenizerFast)

        model_dir = Path(model_dir)
        if not (model_dir / "config.json").exists():
            raise FileNotFoundError(
                f"{model_dir} is not a Whisper checkpoint directory "
                "(no config.json). The weights come from the model repo, so "
                "load through it rather than pointing here by hand:\n"
                "    AutoModel.from_pretrained("
                "'Daporte/prosody2text-weights', trust_remote_code=True)"
            )

        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        self.device = torch.device(device)
        self.num_beams = num_beams

        # fp16 is 1.8x faster (493 ms against 793 ms per phrase at beam 10)
        # and halves the VRAM, and it transcribed all twelve templates
        # identically -- the acoustic evidence in a drawn contour is far
        # coarser than the gap between the precisions. fp32 is the default
        # because that is what the published runs used.
        #
        # CPU stays fp32 regardless: torch has no fast half kernels there, so
        # it would be slower, not faster.
        self.dtype = (
            torch.float16 if half and self.device.type == "cuda" else torch.float32
        )
        self.model = WhisperForConditionalGeneration.from_pretrained(
            model_dir, dtype=self.dtype
        )
        self.model.eval().to(self.device)
        self.mel_bins: int = self.model.config.num_mel_bins

        # The checkpoint carries no tokenizer of its own; it is the stock
        # whisper-medium one. Loading it from a local directory keeps first
        # run offline -- otherwise this is a HuggingFace Hub fetch, which on
        # a fresh Space or an air-gapped machine is a startup failure rather
        # than a slow import.
        source = str(tokenizer_dir) if tokenizer_dir else "openai/whisper-medium"
        self.tokenizer = WhisperTokenizerFast.from_pretrained(source)

    @property
    def provider(self) -> str:
        return str(self.device)

    def _digit_token_ids(self) -> list[int]:
        """Whisper tokens containing a digit.

        Banned under `constrain`: a prosody contour carries no evidence for
        "15" over "fifteen", and numerals are a large, acoustically
        meaningless slice of the vocabulary. Spelled-out numbers remain
        available.
        """
        if getattr(self, "_digits", None) is None:
            # Decode each id rather than inspecting the raw token: Whisper's
            # timestamp markers are spelled "<|24.08|>" and would otherwise
            # match on the digits in their name. They decode to nothing and
            # are suppressed already.
            self._digits = [
                index
                for index in range(len(self.tokenizer))
                if any(c.isdigit() for c in self.tokenizer.decode([index]))
            ]
        return self._digits

    @staticmethod
    def tokens_for(mel: np.ndarray) -> int:
        """A token budget proportional to the drawing's length.

        The mel runs at 100 Hz, so its frame count is the duration. Letting
        generation run to a flat limit regardless lets a short scribble come
        back as a long sentence, which is the decoder's prior rather than the
        drawing's content.
        """
        seconds = len(mel) / 100.0
        words = seconds * WORDS_PER_SECOND
        return max(MIN_TOKENS, int(round(words * TOKENS_PER_WORD)))

    def token_budget(self, requested: int, language: str | None) -> int:
        """`requested`, clipped to what the decoder can actually emit.

        Whisper's decoder has max_target_positions=448 and counts the prefix
        against it, so a budget derived from the input length has to leave
        room for the prefix or generate() raises -- which it did on a 4.8 s
        clip, where the duration-derived budget came to 478 and the
        four-token prefix pushed the total to 482.
        """
        limit = int(getattr(self.model.config, "max_target_positions", 448))
        prefix = len(self.decoder_prefix(language))
        # One spare for EOS, which also occupies a position.
        return max(1, min(int(requested), limit - prefix - 1))

    def decoder_prefix(self, language: str | None = DEFAULT_LANGUAGE) -> list[int]:
        """The tokens the decoder is primed with before any content.

        Built from the tokenizer with the language and task set, which is how
        the published runs got theirs: their snapshot_meta.json records
        [50258, 50259, 50359, 50363], and this reproduces exactly that.
        """
        if language:
            # prefix_tokens is a property derived from these two, so set them
            # on the tokenizer rather than hand-assembling the sequence.
            self.tokenizer.set_prefix_tokens(language=language, task="transcribe")
            return [int(t) for t in self.tokenizer.prefix_tokens]

        prefix = [self.model.config.decoder_start_token_id]
        forced = getattr(self.model.generation_config, "forced_decoder_ids", None)
        for _position, token in forced or []:
            if token is not None:
                prefix.append(int(token))
        return prefix

    def hypotheses(
        self,
        mel: np.ndarray,
        max_tokens: int | None = None,
        num_beams: int | None = None,
        top_k: int | None = None,
        language: str | None = DEFAULT_LANGUAGE,
        no_repeat_ngram_size: int | None = None,
        constrain: bool = True,
    ) -> list[tuple[str, float]]:
        """The beam's ranked candidates as (text, score), best first.

        `language` pins the decoder prefix; see DEFAULT_LANGUAGE. Pass None
        to leave the checkpoint's own forced_decoder_ids alone, which gives a
        two-token-shorter prefix than training used.

        On input this uncertain the top hypothesis is often a narrow winner,
        so the runners-up say more about what the model heard than the single
        best string does.

        Scores are HuggingFace's length-normalised sequence log-probabilities
        where available (higher is better), and NaN where the capture below
        cannot recover them -- the ordering is always meaningful even when the
        numbers are not.
        """
        import torch
        from transformers import GenerationMixin

        beams = num_beams or self.num_beams
        limits = {}
        if language:
            limits["language"] = language

        # The published runs set no_repeat_ngram_size: 3 and nothing else.
        ngram = (DEFAULT_NO_REPEAT_NGRAM if no_repeat_ngram_size is None
                 else no_repeat_ngram_size)
        if ngram:
            limits["no_repeat_ngram_size"] = ngram

        # `constrain` is what the published dataset runs used: ban numerals,
        # and bound the length by the utterance rather than by a flat limit.
        # Letting generation run to 225 regardless lets a half-second clip
        # come back as a paragraph, which is the decoder's prior talking
        # rather than the input.
        if constrain:
            limits["bad_words_ids"] = [[i] for i in self._digit_token_ids()]
            max_tokens = max_tokens or self.tokens_for(mel)
        else:
            max_tokens = max_tokens or DEFAULT_MAX_NEW_TOKENS
        max_tokens = self.token_budget(max_tokens, language)
        features = pad_mel_for_whisper(mel, self.mel_bins, ENCODER_FRAMES)
        batch = torch.from_numpy(features).unsqueeze(0).to(self.device, self.dtype)

        if beams <= 1:
            with torch.no_grad():
                ids = self.model.generate(
                    batch, num_beams=1, max_new_tokens=max_tokens, **limits
                )
            return [(self.tokenizer.batch_decode(
                ids, skip_special_tokens=True)[0].strip(), float("nan"))]

        # Whisper's generate() wrapper keeps only the top beam and does not
        # honour num_return_sequences; asking it for all of them inflates the
        # decoder cache until a 24 GB GPU runs out. The approach here is the
        # one the reference inference script uses: patch
        # GenerationMixin.generate, set num_return_sequences on the inner
        # call, capture every beam, and hand Whisper back the single sequence
        # it expects.
        captured: list = []
        original = GenerationMixin.generate

        def capturing(model_self, input_features, **kwargs):
            config = kwargs.get("generation_config")
            width = config.num_beams if config else 1
            if width <= 1:
                return original(model_self, input_features, **kwargs)

            config.num_return_sequences = width
            # sequences_scores only comes back when the inner call is asked
            # for a dict; Whisper's wrapper does not always pass this down.
            kwargs["return_dict_in_generate"] = True
            # ...and, since transformers 5, only when scores are explicitly
            # requested. Without this the attribute is present but None, so
            # every hypothesis came back scored NaN -- the ranking survived,
            # being the beam's own order, but the numbers beside it were
            # gone and nothing said so.
            kwargs["output_scores"] = True
            out = original(model_self, input_features, **kwargs)
            config.num_return_sequences = 1

            seqs = out.sequences if hasattr(out, "sequences") else out
            captured.append((seqs, getattr(out, "sequences_scores", None)))

            top1 = seqs[::width]
            if hasattr(out, "sequences"):
                out.sequences = top1
            else:
                out = top1
            return out

        GenerationMixin.generate = capturing
        try:
            with torch.no_grad():
                self.model.generate(
                    batch,
                    num_beams=beams,
                    max_new_tokens=max_tokens,
                    return_dict_in_generate=True,
                    **limits,
                )
        finally:
            GenerationMixin.generate = original

        if not captured:
            return []
        seqs, scores = captured[0]
        texts = self.tokenizer.batch_decode(seqs, skip_special_tokens=True)
        values = (
            scores.tolist() if scores is not None
            else [float("nan")] * len(texts)
        )
        ranked = [(t.strip(), s) for t, s in zip(texts, values)]
        return ranked[:top_k] if top_k else ranked