"""Mateli: streaming Finnish speech recognition (MinGRU encoder + RNN-T), PyTorch reference.

    python mateli.py speech.wav                  # text
    python mateli.py speech.wav --words          # one word per line with start/end times

    from mateli import Mateli, load_audio
    asr = Mateli.from_pretrained(".")            # directory with config.json and model.safetensors
    result = asr.transcribe(load_audio("speech.wav"))
    print(result.text)

The encoder is chunk-causal: audio is processed in 320 ms chunks (32 log-mel frames -> 8 encoder
frames of 40 ms), each chunk sees only itself and the past, and the state carried between chunks
has a fixed size. transcribe() runs the same chunk loop a live stream would, so its output equals
streaming recognition of the same audio.

Model: log-mel (80 bins, 25 ms window, 10 ms hop) -> two stride-2 Conv2d -> 16 macaron blocks
(feed-forward / MinGRU / causal depthwise conv / feed-forward), with a windowed self-attention
sublayer after blocks 4, 7, 10 and 13 (keys = the current chunk + 64 previous frames) -> additive
joiner with a one-layer MinGRU prediction network over 44 characters. Greedy decoding.
"""
from __future__ import annotations

import argparse
import json
import math
from dataclasses import dataclass
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F


# --------------------------------------------------------------------------- MinGRU
class MinGRU(nn.Module):
    """minGRU (Feng et al. 2024, "Were RNNs All We Needed?") with a sigmoid highway.
    h_t = (1 - z_t) * h_{t-1} + z_t * g(x_t), computed for a whole chunk with a log-space scan."""

    def __init__(self, hidden: int):
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(hidden, 3 * hidden, bias=False)])

    @staticmethod
    def _log_g(x):
        return torch.where(x >= 0, (F.relu(x) + 0.5).log(), -F.softplus(-x))

    def forward(self, x, state):
        """x [B, T, H], state [B, H] (>= 0) -> (output [B, T, H], new state [B, H])."""
        hidden, gate, proj = self.layers[0](x).chunk(3, dim=-1)
        log_coeffs = -F.softplus(gate)
        log_values = -F.softplus(-gate) + self._log_g(hidden)
        # the carried state enters the scan as an extra first step with coefficient 1
        log_coeffs = torch.cat([torch.zeros_like(log_coeffs[:, :1]), log_coeffs], dim=1)
        log_values = torch.cat([state.clamp_min(0).log().unsqueeze(1), log_values], dim=1)
        a_star = log_coeffs.cumsum(dim=1)
        h = (a_star + (log_values - a_star).logcumsumexp(dim=1)).exp()[:, 1:]
        g = proj.sigmoid()
        return g * h + (1.0 - g) * x, h[:, -1]

    def step(self, x, state):
        """One time step without the scan: x [B, H], state [B, H] -> (output, new state)."""
        hidden, gate, proj = self.layers[0](x).chunk(3, dim=-1)
        g_hidden = torch.where(hidden >= 0, hidden + 0.5, hidden.sigmoid())
        h = torch.lerp(state, g_hidden, gate.sigmoid())
        g = proj.sigmoid()
        return g * h + (1.0 - g) * x, h


# --------------------------------------------------------------------------- encoder
class FeedForward(nn.Module):
    def __init__(self, d: int, expansion: int):
        super().__init__()
        self.net = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, expansion * d), nn.SiLU(),
                                 nn.Linear(expansion * d, d))

    def forward(self, x):
        return self.net(x)


class MinGRUMixer(nn.Module):
    def __init__(self, d: int):
        super().__init__()
        self.norm = nn.LayerNorm(d)
        self.gru = MinGRU(d)

    def forward(self, x, state):
        return self.gru(self.norm(x), state)


class ConvModule(nn.Module):
    """Conformer convolution module with a causal depthwise convolution."""

    def __init__(self, d: int, kernel: int):
        super().__init__()
        self.norm = nn.LayerNorm(d)
        self.pw1 = nn.Conv1d(d, 2 * d, 1)
        self.dw = nn.Conv1d(d, d, kernel, groups=d)
        self.act_norm = nn.LayerNorm(d)
        self.pw2 = nn.Conv1d(d, d, 1)

    def forward(self, x, cache):
        """x [B, T, D], cache [B, D, kernel-1] (the previous GLU outputs)."""
        y = F.glu(self.pw1(self.norm(x).transpose(1, 2)), dim=1)
        y = torch.cat([cache, y], dim=2)
        new_cache = y[:, :, -cache.shape[2]:]
        y = F.silu(self.act_norm(self.dw(y).transpose(1, 2)))
        return self.pw2(y.transpose(1, 2)).transpose(1, 2), new_cache


class Block(nn.Module):
    def __init__(self, d: int, kernel: int, expansion: int):
        super().__init__()
        self.ff1 = FeedForward(d, expansion)
        self.mixer = MinGRUMixer(d)
        self.conv = ConvModule(d, kernel)
        self.ff2 = FeedForward(d, expansion)
        self.norm = nn.LayerNorm(d)

    def forward(self, x, state):
        gru_h, conv_cache = state
        x = x + 0.5 * self.ff1(x)
        m, gru_h = self.mixer(x, gru_h)
        x = x + m
        c, conv_cache = self.conv(x, conv_cache)
        x = x + c
        x = x + 0.5 * self.ff2(x)
        return self.norm(x), (gru_h, conv_cache)


class _Attn(nn.Module):
    def __init__(self, d: int, heads: int):
        super().__init__()
        self.norm = nn.LayerNorm(d)
        self.mha = nn.MultiheadAttention(d, heads, batch_first=True)


class WindowedAttention(nn.Module):
    """Residual self-attention: the chunk's frames attend to the chunk and up to `left` cached
    earlier inputs. No positional encoding (the causal convolutions provide position)."""

    def __init__(self, d: int, heads: int, left: int):
        super().__init__()
        self.attn = _Attn(d, heads)
        self.left = left

    def forward(self, x, cache):
        ctx = torch.cat([cache, x], dim=1)
        q, kv = self.attn.norm(x), self.attn.norm(ctx)
        y, _ = self.attn.mha(q, kv, kv, need_weights=False)
        return x + y, ctx[:, -self.left:]


class Encoder(nn.Module):
    def __init__(self, cfg: dict):
        super().__init__()
        d = cfg["d_model"]
        self.blocks = nn.ModuleList([Block(d, cfg["conv_kernel"], cfg["ffn_expansion"])
                                     for _ in range(cfg["n_blocks"])])
        self.attn = nn.ModuleDict({str(p): WindowedAttention(d, cfg["attn_heads"], cfg["attn_left_context"])
                                   for p in cfg["attn_after_blocks"]})

    def initial_state(self, device):
        d, k = self.blocks[0].norm.normalized_shape[0], self.blocks[0].conv.dw.kernel_size[0]
        return {"blocks": [(torch.zeros(1, d, device=device), torch.zeros(1, d, k - 1, device=device))
                           for _ in self.blocks],
                "attn": {p: torch.zeros(1, 0, d, device=device) for p in self.attn}}

    def forward(self, x, state):
        new_blocks, new_attn = [], {}
        for i, blk in enumerate(self.blocks):
            x, st = blk(x, state["blocks"][i])
            new_blocks.append(st)
            p = str(i + 1)
            if p in self.attn:
                x, new_attn[p] = self.attn[p](x, state["attn"][p])
        return x, {"blocks": new_blocks, "attn": new_attn}


# --------------------------------------------------------------------------- front end
class LogMel(nn.Module):
    def __init__(self, n_mels=80, n_fft=512, win=400, hop=160):
        super().__init__()
        self.n_fft, self.win, self.hop = n_fft, win, hop
        self.register_buffer("fb", torch.zeros(n_mels, n_fft // 2 + 1))     # loaded from the checkpoint
        self.register_buffer("window", torch.hann_window(win))

    def forward(self, wav):
        """[B, samples] -> [B, frames, n_mels], frames = 1 + samples // hop."""
        spec = torch.stft(wav, self.n_fft, self.hop, win_length=self.win, window=self.window,
                          center=True, return_complex=True)
        mel = torch.einsum("mf,bft->bmt", self.fb, spec.real ** 2 + spec.imag ** 2)
        return torch.log(mel + 1e-6).transpose(1, 2)


class Subsampler(nn.Module):
    """Two Conv2d (3x3, stride 2) over (time, mel) -> 40 ms frames."""

    def __init__(self, n_mels: int, d: int):
        super().__init__()
        self.conv = nn.Sequential(nn.Conv2d(1, d, 3, stride=2, padding=1), nn.ReLU(),
                                  nn.Conv2d(d, d, 3, stride=2, padding=1), nn.ReLU())
        f_out = ((n_mels - 1) // 2 + 1 - 1) // 2 + 1
        self.out = nn.Linear(d * f_out, d)

    def forward(self, mel):
        x = self.conv(mel.unsqueeze(1))                        # [B, d, T/4, F/4]
        b, c, t, f = x.shape
        return self.out(x.transpose(1, 2).reshape(b, t, c * f))


class FrontEnd(nn.Module):
    def __init__(self, n_mels: int, d: int):
        super().__init__()
        self.logmel = LogMel(n_mels)
        self.subsample = Subsampler(n_mels, d)


# --------------------------------------------------------------------------- transducer
class PredictionNetwork(nn.Module):
    def __init__(self, vocab_size: int, hidden: int):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, hidden)
        self.rnn = MinGRU(hidden)

    def step(self, token: int, state):
        return self.rnn.step(self.embed(torch.tensor([token], device=state.device)), state)


class Joiner(nn.Module):
    def __init__(self, enc_dim: int, pred_dim: int, joint_dim: int, vocab_size: int):
        super().__init__()
        self.enc_proj = nn.Linear(enc_dim, joint_dim)
        self.pred_proj = nn.Linear(pred_dim, joint_dim)
        self.out = nn.Linear(joint_dim, vocab_size)


@dataclass(frozen=True)
class Word:
    text: str
    start: float      # seconds, time of the first character's encoder frame
    end: float        # seconds, time of the last character's encoder frame


@dataclass
class Transcript:
    words: list[Word]

    @property
    def text(self) -> str:
        return " ".join(w.text for w in self.words)


class Mateli(nn.Module):
    def __init__(self, cfg: dict):
        super().__init__()
        self.cfg = cfg
        d = cfg["d_model"]
        self.frontend = FrontEnd(cfg["n_mels"], d)
        self.encoder = Encoder(cfg)
        self.prednet = PredictionNetwork(len(cfg["vocab"]), cfg["pred_hidden"])
        self.joiner = Joiner(d, cfg["pred_hidden"], cfg["joint_dim"], len(cfg["vocab"]))

    @classmethod
    def from_pretrained(cls, model_dir=".", device: str | None = None) -> "Mateli":
        from safetensors.torch import load_file
        d = Path(model_dir)
        model = cls(json.loads((d / "config.json").read_text(encoding="utf-8")))
        model.load_state_dict(load_file(str(d / "model.safetensors")))
        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        return model.to(device).eval()

    @property
    def device(self):
        return self.joiner.out.weight.device

    @torch.inference_mode()
    def transcribe(self, samples) -> Transcript:
        """float32 mono 16 kHz samples in [-1, 1] -> words with times."""
        cfg = self.cfg
        wav = torch.as_tensor(samples, dtype=torch.float32, device=self.device)
        if wav.ndim != 1:
            raise ValueError(f"expected mono samples of shape [n], got {tuple(wav.shape)}; mix channels down first")
        wav = wav[None]
        if wav.shape[1] <= self.frontend.logmel.n_fft // 2:
            return Transcript([])
        mel = self.frontend.logmel(wav)                         # [1, M, 80]
        n_chunk, ctx = cfg["chunk_mel_frames"], 4               # 4 mel frames = one 40 ms frame (total stride 4)
        state = self.encoder.initial_state(self.device)
        blank, space = cfg["blank_id"], cfg["space_id"]
        pred, h = self.prednet.step(blank, torch.zeros(1, cfg["pred_hidden"], device=self.device))
        pred_proj = self.joiner.pred_proj(pred)
        words, chars, t0, t1, frame = [], "", 0, 0, 0
        fs = cfg["frame_seconds"]
        for start in range(0, mel.shape[1], n_chunk):
            # the subsampler sees the chunk plus 4 earlier mel frames; the first output frame,
            # which only exists because of that context, is dropped
            lo = max(0, start - ctx)
            x = self.frontend.subsample(mel[:, lo:start + n_chunk])
            x = x[:, 1:] if lo < start else x
            enc, state = self.encoder(x, state)
            enc_proj = self.joiner.enc_proj(enc)[0]
            for e in enc_proj:
                for _ in range(cfg["max_symbols_per_frame"]):
                    k = int(self.joiner.out(torch.tanh(e + pred_proj[0])).argmax())
                    if k == blank:
                        break
                    if k == space:
                        if chars:
                            words.append(Word(chars, round((t0 + 1) * fs, 2), round((t1 + 1) * fs, 2)))
                        chars = ""
                    else:
                        if not chars:
                            t0 = frame
                        chars, t1 = chars + cfg["vocab"][k], frame
                    pred, h = self.prednet.step(k, h)
                    pred_proj = self.joiner.pred_proj(pred)
                frame += 1
        if chars:
            words.append(Word(chars, round((t0 + 1) * fs, 2), round((t1 + 1) * fs, 2)))
        return Transcript(words)


def load_audio(path, sample_rate: int = 16000):
    import numpy as np
    import soundfile as sf
    x, sr = sf.read(str(path), dtype="float32", always_2d=True)
    x = x.mean(axis=1)
    if sr != sample_rate:
        from scipy.signal import resample_poly
        g = math.gcd(sr, sample_rate)
        x = resample_poly(x, sample_rate // g, sr // g).astype(np.float32)
    return x


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("audio")
    ap.add_argument("--model-dir", default=str(Path(__file__).resolve().parent))
    ap.add_argument("--device", default=None)
    ap.add_argument("--words", action="store_true", help="print words with start and end times")
    a = ap.parse_args()
    asr = Mateli.from_pretrained(a.model_dir, a.device)
    t = asr.transcribe(load_audio(a.audio, asr.cfg["sample_rate"]))
    if a.words:
        for w in t.words:
            print(f"{w.start:7.2f} {w.end:7.2f}  {w.text}")
    else:
        print(t.text)


if __name__ == "__main__":
    main()