from __future__ import annotations

from pathlib import Path
import numpy as np
import onnxruntime as ort

from .ort_opts import session_options


class MimiDecoder:
    def __init__(self, path: str | Path, threads: int | None = None) -> None:
        self.path = Path(path)
        self.opts = session_options(threads)
        self.session = ort.InferenceSession(str(self.path), self.opts, providers=["CPUExecutionProvider"])

    def decode(self, codes: np.ndarray) -> np.ndarray:
        """
        Input:
            codes: (16, F) int64
        Returns:
            wav: (F * 1920,) float32
        """
        codes = np.asarray(codes, dtype=np.int64)
        if codes.ndim == 2:
            codes_input = codes[None, :, :]  # (1, 16, F)
        elif codes.ndim == 3:
            codes_input = codes
        else:
            raise ValueError(f"Expected codes to have 2 or 3 dims, got {codes.shape}")

        outputs = self.session.run(None, {"codes": codes_input})
        wav = outputs[0]  # (1, 1, samples)
        wav = np.squeeze(wav)
        return wav.astype(np.float32)