"""HF processor: raw audio -> mel input_features, and token ids -> text (SentencePiece)."""
import os
import torch
import sentencepiece as spm
from transformers.feature_extraction_utils import BatchFeature


def _normalize_per_feature(x, seq_len, constant):
    B, _, max_time = x.shape
    steps = torch.arange(max_time, device=x.device).unsqueeze(0).expand(B, max_time)
    valid = steps < seq_len.unsqueeze(1)
    denom = valid.sum(dim=1)
    mean = torch.where(valid.unsqueeze(1), x, torch.zeros_like(x)).sum(dim=2) / denom.unsqueeze(1)
    var = torch.sum(torch.where(valid.unsqueeze(1), x - mean.unsqueeze(2), torch.zeros_like(x)) ** 2,
                    dim=2) / (denom.unsqueeze(1) - 1.0)
    std = torch.sqrt(var).masked_fill(torch.sqrt(var).isnan(), 0.0) + constant
    return (x - mean.unsqueeze(2)) / std.unsqueeze(2)


class FastConformerProcessor:
    def __init__(self, sp, window, fb, params, blank_id=3000, sample_rate=16000):
        self.sp = sp
        self.window = window
        self.fb = fb
        self.p = params
        self.blank_id = blank_id
        self.sample_rate = sample_rate

    @classmethod
    def register_for_auto_class(cls, auto_class="AutoProcessor"):
        return None

    @classmethod
    def from_pretrained(cls, path, **kwargs):
        sp = spm.SentencePieceProcessor(model_file=os.path.join(path, "tokenizer.model"))
        pp = torch.load(os.path.join(path, "preproc.pt"), map_location="cpu")
        blank = pp.get("blank_id", 3000)
        return cls(sp, pp["window"], pp["fb"], pp["params"], blank_id=blank)

    def save_pretrained(self, path, **kwargs):
        os.makedirs(path, exist_ok=True)

    @staticmethod
    def _to_2d(audio):
        if isinstance(audio, torch.Tensor):
            a = audio
        else:
            import numpy as np
            a = torch.as_tensor(np.asarray(audio), dtype=torch.float32)
        if a.dim() == 1:
            a = a.unsqueeze(0)
        return a.float()

    @torch.no_grad()
    def __call__(self, audio, sampling_rate=16000, return_tensors="pt"):
        wav = self._to_2d(audio)
        if sampling_rate != self.sample_rate:
            import torchaudio
            wav = torchaudio.functional.resample(wav, sampling_rate, self.sample_rate)
        p = self.p
        n_fft, hop, win = p["n_fft"], p["hop_length"], p["win_length"]
        wav_len = torch.tensor([wav.shape[1]] * wav.shape[0])
        seq_len = torch.div(wav_len + n_fft - n_fft, hop, rounding_mode="floor")
        tmask = torch.arange(wav.shape[1]).unsqueeze(0) < wav_len.unsqueeze(1)
        x = torch.cat((wav[:, :1], wav[:, 1:] - p["preemph"] * wav[:, :-1]), dim=1).masked_fill(~tmask, 0.0)
        spec = torch.stft(x, n_fft=n_fft, hop_length=hop, win_length=win, window=self.window,
                          center=True, pad_mode="constant", return_complex=True)
        x = torch.view_as_real(spec)
        x = torch.sqrt(x.pow(2).sum(-1)).pow(p["mag_power"])
        x = torch.matmul(self.fb.to(x.dtype), x)
        x = torch.log(x + p["log_zero_guard_value"])
        x = _normalize_per_feature(x, seq_len, p["CONSTANT"])
        max_len = x.size(-1)
        m = torch.arange(max_len).repeat(x.size(0), 1) >= seq_len.unsqueeze(1)
        x = x.masked_fill(m.unsqueeze(1), p["pad_value"])
        return BatchFeature({"input_features": x, "feature_lengths": seq_len}, tensor_type=return_tensors)

    def _clean(self, ids):
        return [int(i) for i in ids if int(i) != self.blank_id]

    def batch_decode(self, sequences, **kwargs):
        token_lists = getattr(sequences, "token_lists", None)
        if token_lists is None:
            seqs = getattr(sequences, "sequences", sequences)
            if isinstance(seqs, torch.Tensor):
                token_lists = [self._clean(row.tolist()) for row in seqs]
            else:
                token_lists = [self._clean(row) for row in seqs]
        return [self.sp.decode(t) for t in token_lists]

    def decode(self, sequence, **kwargs):
        return self.batch_decode([sequence], **kwargs)[0]