"""Run the f0+intensity -> mel conformer.

The real module, loaded from the trained checkpoint. This is the path the
published numbers come from: `generate()` sizes its output from the real
input length and takes a padding mask, so a phrase is predicted at its own
duration.

The alternative was the ONNX export, whose input width is baked in at 2000
frames -- every phrase padded to ten seconds, with the decoder attending
over the padding. Measured against the training pipeline's own predictions
that is 0.9357 correlation against 0.9985 here, which is two orders of
magnitude more than the f0 extraction differs by, so it is not a reasonable
substitute when the point is reproducing the paper.
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np

# Keys FeaturesToMelConv's constructor accepts. model_config.json carries
# training settings alongside them, and the model calls it mel_dim rather
# than mel_bins and takes no max_len -- so guessing the set does not work.
CONSTRUCTOR_KEYS = {
    "input_dim", "d_model", "nhead", "num_encoder_layers",
    "num_decoder_layers", "encoder_kernel_size", "conv_kernel_size",
    "dim_feedforward", "dropout", "loss_type", "length_ratio",
    "gradient_checkpointing", "decoder_only", "attn_window",
    "mel_dim",
}


class Conformer:
    """Predicts a mel spectrogram from 200 Hz prosody features."""

    def __init__(self, checkpoint: Path | str, cpu: bool = False) -> None:
        """`checkpoint` is a directory holding model_config.json and weights.

        The weights may sit beside the config or in a `checkpoint-*`
        subdirectory, which is how the training runs leave them.
        """
        import torch
        from safetensors.torch import load_file

        from .conformer_model import FeaturesToMelConv

        root = Path(checkpoint)
        config_path = root / "model_config.json"
        if not config_path.exists():
            raise FileNotFoundError(
                f"{config_path} not found. The conformer needs the directory "
                f"holding model_config.json and its weights."
            )

        config = json.loads(config_path.read_text())
        kwargs = {k: v for k, v in config.items() if k in CONSTRUCTOR_KEYS}

        weights = sorted(root.glob("model.safetensors"))
        weights += sorted(root.glob("checkpoint-*/model.safetensors"))
        if not weights:
            raise FileNotFoundError(f"no model.safetensors under {root}")

        self._model = FeaturesToMelConv(**kwargs)
        # strict=False because the checkpoint carries training-only buffers
        # the inference classes do not declare. Checked against the ONNX
        # export's output, which is what catches a genuinely wrong load.
        self._model.load_state_dict(load_file(str(weights[-1])), strict=False)
        self._model.eval()

        self._device = torch.device(
            "cpu" if cpu or not torch.cuda.is_available() else "cuda"
        )
        self._model.to(self._device)

        self.mel_bins: int = int(config.get("mel_dim", config.get("mel_bins", 12)))
        self.length_ratio: float = float(config.get("length_ratio", 0.5))
        self._source = str(weights[-1])

    @property
    def provider(self) -> str:
        return f"torch ({self._device.type})"

    @property
    def max_duration_s(self) -> float:
        """No fixed window: generate() sizes its output from the input."""
        return float("inf")

    def predict(self, features: np.ndarray) -> np.ndarray:
        """[T, 2] normalised features -> [T*length_ratio, mel_bins] mel.

        The padding mask is all-false: a single phrase has nothing padded,
        and passing the mask is what lets generate() size the output from
        the real length rather than from a fixed window.
        """
        import torch

        features = np.asarray(features, dtype=np.float32)
        if features.ndim != 2 or features.shape[1] != 2:
            raise ValueError(f"expected [T, 2] features, got {features.shape}")
        if features.size == 0:
            return np.zeros((0, self.mel_bins), dtype=np.float32)

        with torch.no_grad():
            batch = torch.from_numpy(features).unsqueeze(0).to(self._device)
            mask = torch.zeros(
                (1, batch.shape[1]), dtype=torch.bool, device=self._device
            )
            out = self._model.generate(batch, input_padding_mask=mask)

        mel = (out[0] if isinstance(out, tuple) else out)[0]
        return mel.detach().cpu().numpy().astype(np.float32)