"""Prosody to text: f0 and intensity in, words out.

Two signals -- interpolated f0 and intensity at 200 Hz -- through a conformer
to a mel spectrogram, and the mel through a fine-tuned Whisper. No phonetic
content at any point: the model sees only how the pitch moved and how loud
it was.

Loaded the way any custom architecture on the Hub is:

    from transformers import AutoModel
    model = AutoModel.from_pretrained(
        "Daporte/prosody2text-weights", trust_remote_code=True
    )
    print(model.transcribe("utterance.wav").text)

A note on the shape of this class. It subclasses `PreTrainedModel` so that
`from_pretrained`, the `AutoModel` API and the Hub's download metrics all
work, but it is a pipeline rather than a single network: a conformer and a
`WhisperForConditionalGeneration` held alongside each other, loaded from
their own checkpoints rather than from one state dict. `forward` runs them
in sequence under `no_grad`, so this is not trainable as it stands -- it
exists to run the published pipeline, which is what the paper reports.
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import torch

from transformers import PreTrainedModel

from .configuration_prosody import ProsodyToTextConfig
from .conformer import Conformer
from .features import ProsodyFeatures
from .speech import extract_features
from .whisper_decoder import Whisper, looks_like_fallback


class ProsodyResult:
    """What one utterance produces.

    Carries the intermediate stages as well as the words, because the whole
    point is being able to see them: a transcription that looks wrong is
    usually a contour that looks wrong first.
    """

    def __init__(
        self,
        features: ProsodyFeatures,
        mel: np.ndarray,
        ranked: list[tuple[str, float]],
    ) -> None:
        self.features = features
        self.mel = mel
        self.ranked = ranked

    @property
    def text(self) -> str:
        """The chosen hypothesis, empty if nothing was transcribed."""
        return self.ranked[0][0] if self.ranked else ""

    @property
    def candidates(self) -> list[str]:
        """The whole beam. The top hypothesis is often a narrow winner on an
        input this uncertain, so the runners-up say more about what the model
        recovered than the single best string does."""
        return [text for text, _ in self.ranked]

    @property
    def fallback(self) -> bool:
        """Whether the result is Whisper's no-speech boilerplate rather than
        a transcription. Trained on captions, it answers unintelligible
        input with "Thank you." and shouted auto-caption filler, which reads
        as a result but is not one."""
        return bool(self.text) and looks_like_fallback(self.text)

    def __repr__(self) -> str:
        return (
            f"ProsodyResult(text={self.text!r}, "
            f"frames={self.features.num_frames}, mel={self.mel.shape})"
        )


class ProsodyToTextModel(PreTrainedModel):
    """The published pipeline, as a Hub model."""

    config_class = ProsodyToTextConfig
    # Nothing here is trained by this class, and the weights are loaded by
    # the two components rather than from a single state dict.
    base_model_prefix = "prosody"
    supports_gradient_checkpointing = False
    _supports_sdpa = False

    def __init__(self, config: ProsodyToTextConfig) -> None:
        super().__init__(config)
        self.config = config
        self._conformer: Conformer | None = None
        self._whisper: Whisper | None = None
        self._root: Path | None = None

    # ---- loading ---------------------------------------------------------

    @classmethod
    def from_pretrained(cls, model_id, *args, **kwargs):  # noqa: ANN001
        """Build the pipeline from a repo or a local directory.

        Overridden rather than relying on the base implementation: there is
        no single state dict to load. Each half is its own checkpoint in its
        own subfolder, so both are resolved from the snapshot and loaded by
        their own code.
        """
        from transformers import AutoConfig

        half = kwargs.pop("half", False)
        cpu = kwargs.pop("cpu", False)
        load_whisper = kwargs.pop("whisper", True)
        config = kwargs.pop("config", None)
        # Consumed so a caller can pass the usual from_pretrained arguments
        # without them reaching code that has no use for them.
        kwargs.pop("torch_dtype", None)
        kwargs.pop("device_map", None)

        root = Path(model_id)
        if not root.is_dir():
            from huggingface_hub import snapshot_download

            patterns = ["*.py", "*.json", config_file(model_id)]
            patterns += ["conformer/*", "tokenizer/*"]
            if load_whisper:
                patterns.append("whisper/*")
            root = Path(snapshot_download(
                repo_id=str(model_id),
                allow_patterns=[p for p in patterns if p],
                revision=kwargs.pop("revision", None),
                token=kwargs.pop("token", None),
            ))

        if config is None:
            config = AutoConfig.from_pretrained(
                str(root), trust_remote_code=True
            )

        model = cls(config)
        model._root = root
        model._load(root, half=half, cpu=cpu, whisper=load_whisper)
        return model

    def _load(
        self, root: Path, half: bool, cpu: bool, whisper: bool
    ) -> None:
        warn_if_driver_hidden(whisper)

        self._conformer = Conformer(
            root / self.config.conformer_subfolder, cpu=cpu
        )
        if not whisper:
            return

        tokenizer = root / self.config.tokenizer_subfolder
        self._whisper = Whisper(
            root / self.config.whisper_subfolder,
            half=half,
            num_beams=self.config.num_beams,
            tokenizer_dir=(
                tokenizer if (tokenizer / "tokenizer.json").exists() else None
            ),
        )

    @property
    def device_summary(self) -> str:
        """Where each half is running, which is worth printing: the whole
        difference between a minute and an hour is whether Whisper found the
        GPU."""
        parts = []
        if self._conformer is not None:
            parts.append(f"conformer: {self._conformer.provider}")
        if self._whisper is not None:
            parts.append(f"whisper: {self._whisper.provider}")
        return ", ".join(parts) or "nothing loaded"

    # ---- running ---------------------------------------------------------

    def forward(
        self,
        f0_hz: np.ndarray | torch.Tensor | None = None,
        intensity_db: np.ndarray | torch.Tensor | None = None,
        transcribe: bool = True,
        **kwargs,
    ) -> ProsodyResult:
        """The two contours in, the mel and the ranked beam out.

        Physical units -- Hz and dB on the 200 Hz grid -- not normalised. The
        normalisation to the model's own scale happens here so a caller
        cannot apply it twice, which produces a mel that looks plausible and
        is wrong.
        """
        if f0_hz is None or intensity_db is None:
            raise ValueError("f0_hz and intensity_db are both required")

        features = ProsodyFeatures(
            f0_hz=as_array(f0_hz), intensity_db=as_array(intensity_db)
        )
        return self.run(features, transcribe=transcribe)

    def run(
        self, features: ProsodyFeatures, transcribe: bool = True
    ) -> ProsodyResult:
        """Features to words, with the mel kept on the way through."""
        if self._conformer is None:
            raise RuntimeError(
                "the conformer is not loaded; build this with "
                "from_pretrained()"
            )
        if features.num_frames == 0:
            return ProsodyResult(
                features, np.zeros((0, self.config.mel_bins), np.float32), []
            )

        mel = self._conformer.predict(features.normalised())
        ranked: list[tuple[str, float]] = []
        if transcribe and self._whisper is not None and len(mel):
            ranked = self._whisper.hypotheses(
                mel,
                language=self.config.language,
                # None under `constrain`, so the budget comes from the
                # utterance's duration rather than the flat ceiling.
                max_tokens=(
                    None if self.config.constrain
                    else self.config.max_new_tokens
                ),
                no_repeat_ngram_size=self.config.no_repeat_ngram_size,
                constrain=self.config.constrain,
            )
        return ProsodyResult(features, mel, ranked)

    def transcribe(
        self,
        audio,
        sampling_rate: int | None = None,
        trim: bool = True,
    ) -> ProsodyResult:
        """Audio to words, for the common case.

        `audio` is a path, or samples with `sampling_rate` given. `trim` cuts
        the quiet lead-in and tail a live recording carries -- yaapt finds a
        pitch in room noise, so without it several hundred milliseconds of
        nothing reach the model. Dataset clips usually arrive trimmed
        already; pass False to match the published runs exactly.
        """
        if isinstance(audio, (str, Path)):
            # In a try block on purpose. trust_remote_code scans a modeling
            # file's imports statically and refuses to load unless every one
            # is installed -- and it skips Try nodes precisely so a module
            # can mark a dependency optional. Reading a path is the only
            # thing here that needs soundfile; a caller passing arrays, as
            # the drawing demo does, should not be made to install it.
            try:
                import soundfile
            except ImportError as error:  # pragma: no cover
                raise ImportError(
                    "reading audio from a path needs soundfile "
                    "(pip install soundfile); pass samples and "
                    "sampling_rate instead to avoid it"
                ) from error

            samples, sampling_rate = soundfile.read(
                str(audio), dtype="float64", always_2d=False
            )
            if getattr(samples, "ndim", 1) > 1:
                samples = samples.mean(axis=1)
        else:
            if sampling_rate is None:
                raise ValueError(
                    "sampling_rate is required when passing samples"
                )
            samples = as_array(audio)

        return self.run(extract_features(samples, sampling_rate, trim=trim))


# ---- helpers -------------------------------------------------------------


def as_array(value) -> np.ndarray:  # noqa: ANN001
    if isinstance(value, torch.Tensor):
        value = value.detach().cpu().numpy()
    array = np.asarray(value, dtype=np.float64)
    return array.squeeze() if array.ndim > 1 else array


def config_file(model_id) -> str:  # noqa: ANN001
    """`config.json`, named separately so the download pattern reads clearly."""
    return "config.json"


# Directories where a GPU driver may sit without being on the loader path.
#
# A pip-installed torch finds libcuda.so.1 through the dynamic loader, and on
# most systems the driver package puts it where the loader already looks. Not
# everywhere: NixOS and some container images keep it outside that path, so
# torch reports "Found no NVIDIA driver on your system" on a machine whose
# nvidia-smi works perfectly.
DRIVER_DIRS = (
    "/run/opengl-driver/lib",          # NixOS
    "/usr/lib/x86_64-linux-gnu",       # Debian, Ubuntu
    "/usr/lib64",                      # Fedora, RHEL
    "/usr/local/nvidia/lib64",         # some container images
)


def warn_if_driver_hidden(whisper: bool) -> None:
    """Say when a GPU is present but unreachable.

    Falling back to CPU silently is the wrong default for something twenty
    times slower: across a few thousand utterances that is the difference
    between an hour and a day, and the cause is one environment variable.
    """
    if not whisper:
        return
    if torch.cuda.is_available():
        return

    import ctypes

    try:
        ctypes.CDLL("libcuda.so.1")
        return          # loadable, so CUDA is unavailable for another reason
    except OSError:
        pass

    for directory in DRIVER_DIRS:
        if (Path(directory) / "libcuda.so.1").exists():
            print(
                f"warning: a GPU driver is installed but not on the library "
                f"path, so this will run on CPU. Re-run with:\n"
                f"    LD_LIBRARY_PATH={directory}:$LD_LIBRARY_PATH",
                flush=True,
            )
            return