"""Haas-ja-v1.0 ONNX inference without NeMo: NeMo-compatible log-mel + TDT greedy decoding."""
import librosa
import numpy as np
import onnxruntime as ort
import sentencepiece as spm
import soundfile as sf

SR, N_FFT, WIN, HOP, N_MELS, PAD_TO = 16000, 512, 400, 160, 128, 16
DURATIONS = [0, 1, 2, 3, 4]
MAX_SYMBOLS = 10
BLANK = 3072


def _bf16(x):
    """Round float32 values to bfloat16 precision (round to nearest even), kept as float32."""
    b = np.asarray(x, np.float32).view(np.uint32)
    b = (b + 0x7FFF + ((b >> 16) & 1)) & 0xFFFF0000
    return b.astype(np.uint32).view(np.float32)


# The Haas checkpoint stores the preprocessor's window and filterbank in bfloat16 and NeMo computes with those,
# so round ours the same way to match the official features.
_FB = _bf16(librosa.filters.mel(sr=SR, n_fft=N_FFT, n_mels=N_MELS, fmin=0, fmax=SR / 2, norm="slaney"))
_WINDOW = _bf16(np.hanning(WIN).astype(np.float32))  # symmetric Hann, as torch.hann_window(periodic=False)
_WINDOW_PAD = np.pad(_WINDOW, ((N_FFT - WIN) // 2, (N_FFT - WIN + 1) // 2))


def logmel(audio):
    """NeMo AudioToMelSpectrogramPreprocessor (eval: no dither), batch 1. Returns ([1,128,T_pad], valid_len)."""
    x = np.asarray(audio, np.float32)
    seq_len = len(x) // HOP  # NeMo get_seq_len: (len + 2*(n_fft//2) - n_fft) // hop
    x = np.concatenate([x[:1], x[1:] - np.float32(0.97) * x[:-1]])
    x = np.pad(x, N_FFT // 2)  # center=True, pad_mode="constant"
    n = 1 + (len(x) - N_FFT) // HOP
    frames = np.lib.stride_tricks.as_strided(x, (n, N_FFT), (x.strides[0] * HOP, x.strides[0]))
    spec = np.fft.rfft(frames * _WINDOW_PAD, axis=1)
    power = (spec.real ** 2 + spec.imag ** 2).astype(np.float32).T  # [257, T], power 2.0
    mel = np.log(_FB @ power + np.float32(2 ** -24))
    m = mel[:, :seq_len]
    mean = m.mean(-1, keepdims=True)
    std = np.sqrt(((m - mean) ** 2).sum(-1, keepdims=True) / (seq_len - 1)) + np.float32(1e-5)
    mel = (mel - mean) / std
    mel[:, seq_len:] = 0
    T = mel.shape[-1]
    if T % PAD_TO:
        mel = np.pad(mel, ((0, 0), (0, PAD_TO - T % PAD_TO)))
    return mel[None].astype(np.float32), seq_len


def load_wav(path):
    audio, sr = sf.read(path, dtype="float32")
    if audio.ndim > 1:
        audio = audio.mean(1)
    if sr != SR:
        audio = librosa.resample(audio, orig_sr=sr, target_sr=SR)
    return audio


def sessions(enc_path, dec_path, threads=4):
    so = ort.SessionOptions()
    so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    so.intra_op_num_threads = threads
    p = ["CPUExecutionProvider"]
    return ort.InferenceSession(enc_path, so, providers=p), ort.InferenceSession(dec_path, so, providers=p)


def _ftype(sess, name):
    t = next(i.type for i in sess.get_inputs() if i.name == name)
    return np.float16 if "float16" in t else np.float32


def encode(enc_sess, mel, mel_len):
    ft = _ftype(enc_sess, "audio_signal")
    enc, enc_len = enc_sess.run(None, {"audio_signal": mel.astype(ft), "length": np.array([mel_len], np.int64)})
    return enc, int(enc_len[0])


def greedy_tdt(dec_sess, enc, T, num_layers=2, width=20480):
    """Same rule as real-weights/scripts/verify_onnx.py greedy_tdt_onnx (NeMo label-looping)."""
    ft = _ftype(dec_sess, "encoder_outputs")
    st = _ftype(dec_sess, "input_states_1")
    enc = enc.astype(ft)
    s1 = np.zeros((num_layers, 1, width), st)
    s2 = np.zeros_like(s1)
    last = BLANK
    tokens, t, last_ts, count = [], 0, -1, 0

    def step(frame):
        out = dec_sess.run(None, {
            "encoder_outputs": enc[:, :, frame:frame + 1],
            "targets": np.array([[last]], np.int32), "target_length": np.array([1], np.int32),
            "input_states_1": s1, "input_states_2": s2,
        })
        logits = out[0][0, 0, 0]
        k = int(logits[:-len(DURATIONS)].argmax())
        d = DURATIONS[int(logits[-len(DURATIONS):].argmax())]
        if k == BLANK and d == 0:
            d = 1
        return k, d, out[2], out[3]

    while t < T:
        while True:
            k, d, n1, n2 = step(t)
            t_label, t = t, t + d
            if k != BLANK or t >= T:
                break
        if k != BLANK:
            tokens.append(k)
            count = count + 1 if t_label == last_ts else 1
            last_ts = t_label
            s1, s2, last = n1, n2, k
            if count >= MAX_SYMBOLS and last_ts == t:
                t += 1
    return tokens


class Tokenizer:
    def __init__(self, path):
        self.sp = spm.SentencePieceProcessor(model_file=path)

    def decode(self, ids):
        return self.sp.decode(ids)