"""CPU streaming greedy inference for the Sorani Run B TorchScript model."""
from pathlib import Path
import argparse
import math

import kaldi_native_fbank as knf
import numpy as np
import sentencepiece as spm
import torch


class Recognizer:
    def __init__(self, model_dir=None):
        root = Path(model_dir or Path(__file__).parent)
        torch.set_num_threads(1)
        torch._C._jit_set_profiling_executor(False)
        torch._C._jit_set_profiling_mode(False)
        torch._C._set_graph_executor_optimize(False)
        self.model = torch.jit.load(str(root / 'model.pt'), map_location='cpu').eval()
        self.tokenizer = spm.SentencePieceProcessor(model_file=str(root / 'bpe.model'))

    def create_stream(self):
        return Stream(self.model, self.tokenizer)


class Stream:
    def __init__(self, model, tokenizer):
        self.model, self.tokenizer = model, tokenizer
        opts = knf.FbankOptions()
        opts.frame_opts.samp_freq = 16000
        opts.frame_opts.frame_length_ms = 25
        opts.frame_opts.frame_shift_ms = 10
        opts.frame_opts.dither = 0
        opts.frame_opts.snip_edges = False
        opts.frame_opts.preemph_coeff = 0.97
        opts.frame_opts.window_type = 'povey'
        opts.frame_opts.remove_dc_offset = True
        opts.frame_opts.round_to_power_of_two = True
        opts.mel_opts.num_bins = 80
        opts.mel_opts.low_freq = 20
        opts.mel_opts.high_freq = 7600
        opts.use_energy = False
        opts.use_log_fbank = True
        opts.use_power = True
        self.fbank = knf.OnlineFbank(opts)
        self.step = model.encoder.chunk_size * 2
        self.window = self.step + model.encoder.pad_length
        self.states = model.encoder.get_init_states(device=torch.device('cpu'))
        self.processed = 0
        self.hyp = [model.decoder.blank_id] * model.decoder.context_size
        self.decoder_out = None
        self.finished = False

    def update_decoder(self):
        context = torch.tensor([self.hyp[-self.model.decoder.context_size:]], dtype=torch.int32)
        self.decoder_out = self.model.decoder(context, torch.tensor([False])).squeeze(1)

    @torch.inference_mode()
    def accept(self, audio, final=False):
        """Accept mono float PCM at 16 kHz (normally in [-1, 1]); return partial text.

        Call with final=True once to flush. Create a fresh stream for the next
        utterance. No spelling correction or transcript normalization is applied.
        """
        if self.finished:
            raise ValueError('Stream already finished; create a new stream.')
        audio = np.asarray(audio, dtype=np.float32)
        if audio.ndim != 1 or not np.isfinite(audio).all():
            raise ValueError('Expected finite, one-dimensional mono float PCM.')
        self.fbank.accept_waveform(16000, audio.tolist())
        if final:
            self.fbank.accept_waveform(16000, [0.0] * 8000)
            self.fbank.input_finished()
            self.finished = True
        while self.fbank.num_frames_ready - self.processed >= self.window:
            features = torch.from_numpy(np.stack([
                self.fbank.get_frame(self.processed + i) for i in range(self.window)
            ])).unsqueeze(0)
            encoded, lengths, self.states = self.model.encoder(
                features, torch.tensor([self.window], dtype=torch.int32), self.states)
            self.processed += self.step
            if self.decoder_out is None:
                self.update_decoder()
            for frame in encoded[0, :int(lengths[0])]:
                token = int(self.model.joiner(frame.unsqueeze(0), self.decoder_out).argmax(-1).item())
                if token != self.model.decoder.blank_id:
                    self.hyp.append(token)
                    self.update_decoder()
        return self.tokenizer.decode(self.hyp[self.model.decoder.context_size:])


def main():
    import soundfile as sf
    from scipy.signal import resample_poly
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('audio', help='WAV/FLAC or another format supported by libsndfile')
    parser.add_argument('--model-dir', type=Path, default=Path(__file__).parent)
    parser.add_argument('--partials', action='store_true', help='Print changed partial hypotheses')
    args = parser.parse_args()
    audio, sr = sf.read(args.audio, dtype='float32', always_2d=True)
    audio = audio.mean(axis=1)
    if not len(audio):
        parser.error('Audio is empty')
    if sr != 16000:
        factor = math.gcd(sr, 16000)
        audio = resample_poly(audio, 16000 // factor, sr // factor).astype(np.float32)
    stream = Recognizer(args.model_dir).create_stream()
    previous = ''
    for start in range(0, len(audio), 5120):
        text = stream.accept(audio[start:start + 5120])
        if args.partials and text != previous:
            print(text, flush=True)
        previous = text
    print(stream.accept(np.empty(0, dtype=np.float32), final=True))


if __name__ == '__main__':
    main()