import logging

import kaldi_native_fbank as knf
import numpy as np
import numpy.typing as npt

logger = logging.getLogger(__name__)

SAMPLE_RATE = 16000
NUM_MEL_BINS = 80
FRAME_SHIFT_MS = 10

FBankReturnT = npt.NDArray[np.float32]


def _create_fbank_opts(sample_rate: int) -> knf.FbankOptions:
    fbank_opts = knf.FbankOptions()
    fbank_opts.mel_opts.num_bins = NUM_MEL_BINS
    fbank_opts.frame_opts.frame_length_ms = 25
    fbank_opts.frame_opts.frame_shift_ms = FRAME_SHIFT_MS
    fbank_opts.frame_opts.samp_freq = sample_rate
    fbank_opts.frame_opts.window_type = "hamming"
    fbank_opts.frame_opts.dither = 0.0
    fbank_opts.frame_opts.snip_edges = True
    fbank_opts.frame_opts.round_to_power_of_two = True
    return fbank_opts


class FbankExtractor:
    def __init__(
        self,
        sample_rate: int = SAMPLE_RATE,
    ) -> None:
        self.sample_rate = sample_rate
        self.fbank_opts = _create_fbank_opts(sample_rate)

    def compute_fbank(self, audio_data: npt.NDArray[np.float32]) -> FBankReturnT:
        extractor = knf.OnlineFbank(self.fbank_opts)

        scaled_audio = (audio_data * (1 << 15)).astype(np.float32)

        extractor.accept_waveform(self.sample_rate, scaled_audio)
        extractor.input_finished()

        features = np.array([extractor.get_frame(j) for j in range(extractor.num_frames_ready)])
        features = features - np.mean(features, axis=0, keepdims=True)
        return features
