import spaces  # must come before torch

import json
import math
import time

import gradio as gr
import librosa
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from transformers import EmbeddingGemma2Model, Gemma4AudioFeatureExtractor

# Full float32 math, as the reference (no TF32).
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

EG2_REPO = "google/embeddinggemma-2"
EG2_REV = "914f7f89142e33e77833254d9c9b90c3cef7303b"
READER_REPO = "mlboydaisuke/EmbeddingGemma-2-Transcribe-Reader"
SR = 16000
MAX_SECONDS = 15 * 60

# Window contract (reader_config.json -> window): 30 s windows (3000 frames / 750 tokens),
# 96 tokens of left overlap, 1 token dropped at the right of every window but the last.
WIN_FRAMES, LEFT_OVERLAP, DROP_LAST = 3000, 96, 1
WIN_TOKENS = WIN_FRAMES // 4
STEP_TOKENS = WIN_TOKENS - LEFT_OVERLAP - DROP_LAST  # 653

# ------------------------------------------------------------------ reader (reader/train/model.py, arch "a")
D_IN, D_OUT = 512, 768


def rope(T, dim, base, device):
    inv = 1.0 / (base ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim))
    f = torch.arange(T, device=device, dtype=torch.float32)[:, None] * inv[None, :]
    emb = torch.cat([f, f], dim=-1)
    return emb.cos(), emb.sin()


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


class Block(nn.Module):
    def __init__(self, d, heads, ffn):
        super().__init__()
        self.heads = heads
        self.ln1 = nn.LayerNorm(d)
        self.qkv = nn.Linear(d, 3 * d)
        self.proj = nn.Linear(d, d)
        self.ln2 = nn.LayerNorm(d)
        self.ff1 = nn.Linear(d, ffn)
        self.ff2 = nn.Linear(ffn, d)

    def forward(self, x, mask, cos, sin):
        B, T, D = x.shape
        q, k, v = self.qkv(self.ln1(x)).view(B, T, 3, self.heads, D // self.heads).permute(2, 0, 3, 1, 4)
        q = q * cos + rotate_half(q) * sin
        k = k * cos + rotate_half(k) * sin
        a = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        x = x + self.proj(a.transpose(1, 2).reshape(B, T, D))
        return x + self.ff2(F.gelu(self.ff1(self.ln2(x))))


class ReaderA(nn.Module):
    def __init__(self, d=512, layers=6, heads=8, ffn=2048, rope_base=10000.0):
        super().__init__()
        self.heads, self.rope_base = heads, rope_base
        self.ln_in = nn.LayerNorm(D_IN)
        self.inp = nn.Linear(D_IN, d)
        self.blocks = nn.ModuleList([Block(d, heads, ffn) for _ in range(layers)])
        self.ln_out = nn.LayerNorm(d)
        self.out = nn.Linear(d, D_OUT)

    def forward(self, x, lengths):
        T = x.shape[1]
        mask = (torch.arange(T, device=x.device)[None, :] < lengths[:, None])[:, None, None, :]
        cos, sin = rope(T, self.inp.out_features // self.heads, self.rope_base, x.device)
        h = self.inp(self.ln_in(x))
        for blk in self.blocks:
            h = blk(h, mask, cos, sin)
        return self.out(self.ln_out(h))


class Reader(nn.Module):
    def __init__(self, layers=6):
        super().__init__()
        self.body = ReaderA(layers=layers)
        self.blank = nn.Linear(D_OUT, 1)
        self.log_scale = nn.Parameter(torch.tensor(math.log(20.0)))

    @property
    def scale(self):
        return self.log_scale.exp()

    def forward(self, x, lengths):
        v = self.body(x, lengths)
        return F.normalize(v, dim=-1), self.blank(v).squeeze(-1)


# ------------------------------------------------------------------ load everything at module scope
cfg = json.load(open(hf_hub_download(READER_REPO, "reader_config.json")))
with open(hf_hub_download(READER_REPO, "vocab.txt"), encoding="utf-8") as f:
    VOCAB = [line.rstrip("\n") for line in f if line.strip()]
BLANK_ID = len(VOCAB)

reader = Reader(layers=cfg["architecture"]["layers"])
reader.load_state_dict(torch.load(hf_hub_download(READER_REPO, "reader.pt"), map_location="cpu", weights_only=True))
reader = reader.to("cuda").eval()

table = torch.from_numpy(np.load(hf_hub_download(READER_REPO, "word_table.npy")).astype(np.float32))
table = F.normalize(table, dim=-1).to("cuda")

feature_extractor = Gemma4AudioFeatureExtractor.from_pretrained(EG2_REPO, revision=EG2_REV)
encoder = EmbeddingGemma2Model.from_pretrained(EG2_REPO, revision=EG2_REV, dtype=torch.float32, vision_config=None)
encoder = encoder.to("cuda").eval()
print("audio attention:", getattr(encoder.config.audio_config, "_attn_implementation", None))


def n_tokens_for(n_frames):
    return ((n_frames + 1) // 2 + 1) // 2


def plan_windows(n_frames):
    """Windowing.plan of the eg2-transcribe app: (token0, frame0, frames, keep_start, keep_end)."""
    total = n_tokens_for(n_frames)
    plan, i = [], 0
    while total > 0:
        a = i * STEP_TOKENS
        nf = min(WIN_FRAMES, n_frames - 4 * a)
        m = n_tokens_for(nf)
        final = a + m >= total
        plan.append((a, 4 * a, nf, 0 if i == 0 else LEFT_OVERLAP, m if final else WIN_TOKENS - DROP_LAST))
        if final:
            break
        i += 1
    return plan


def load_audio(path):
    wav, _ = librosa.load(path, sr=SR, mono=True)
    return wav.astype(np.float32)


def get_duration(audio_path, *args, **kwargs):
    try:
        secs = librosa.get_duration(path=audio_path)
    except Exception:
        secs = 60
    return int(min(60, 8 + secs * 0.01))


@spaces.GPU(duration=get_duration)
@torch.no_grad()
def transcribe(audio_path: str) -> tuple[str, str]:
    """Transcribe English speech with the frozen EmbeddingGemma 2 audio encoder + the trained word reader.

    Args:
        audio_path: path to an audio file (any format librosa can read; resampled to 16 kHz mono).

    Returns:
        The lower-case transcript (no punctuation) and a short stats line.
    """
    if not audio_path:
        raise gr.Error("Please upload or record some audio first.")
    t0 = time.time()
    wav = load_audio(audio_path)
    secs = len(wav) / SR
    if secs > MAX_SECONDS:
        raise gr.Error(f"Audio is {secs / 60:.1f} min long; this demo accepts up to {MAX_SECONDS // 60} min.")
    n_frames = max(0, (len(wav) - 161) // 160 + 1)
    if n_frames < 4:
        raise gr.Error("Audio is too short.")

    # log-mel of the whole recording (no 30 s cut); each window is a slice of its frames
    out = feature_extractor([wav], return_tensors="np", max_length=None, truncation=False)
    feats = out["input_features"][0][:n_frames].astype(np.float32)

    words_all = []
    plan = plan_windows(n_frames)
    for token0, frame0, nf, k0, k1 in plan:
        x = torch.from_numpy(np.ascontiguousarray(feats[frame0:frame0 + nf]))[None].to("cuda")
        m = torch.ones(1, nf, dtype=torch.bool, device="cuda")
        res = encoder.get_audio_features(x, m, return_dict=True)
        soft = res.pooler_output[0][res.attention_mask[0]].float()  # [tokens, 512]
        lengths = torch.tensor([soft.shape[0]], device="cuda")
        h, blank = reader(soft[None], lengths)
        best, ids = (h[0] @ table.T).max(dim=1)
        labels = torch.where(blank[0] > best * reader.scale, torch.full_like(ids, BLANK_ID), ids)
        words_all.extend(labels[k0:k1].tolist())

    # greedy CTC read-out over the kept tokens of all windows
    text, prev = [], None
    for c in words_all:
        if c != prev and c != BLANK_ID:
            text.append(VOCAB[c])
        prev = c
    dt = time.time() - t0
    stats = f"{secs:.1f} s of audio · {len(plan)} window(s) of 30 s · {len(words_all)} tokens · {dt:.2f} s"
    return " ".join(text), stats


CSS = """
#col-container { max-width: 1000px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# EmbeddingGemma 2 Transcribe\n"
            "English speech-to-text from the **frozen audio encoder of "
            "[EmbeddingGemma 2](https://huggingface.co/google/embeddinggemma-2)** plus a small trained reader "
            "that maps every 40 ms token to one of 200,000 English words (or blank) — greedy CTC, no language model. "
            "Runs the PyTorch reference path of "
            "[EmbeddingGemma-2-Transcribe-LiteRT](https://huggingface.co/mlboydaisuke/EmbeddingGemma-2-Transcribe-LiteRT) "
            "/ [Reader v2](https://huggingface.co/mlboydaisuke/EmbeddingGemma-2-Transcribe-Reader), "
            "with the same 30 s sliding-window contract as the on-device app."
        )
        with gr.Row():
            with gr.Column():
                audio = gr.Audio(label="English speech", type="filepath", sources=["upload", "microphone"])
                run = gr.Button("Transcribe", variant="primary")
            with gr.Column():
                transcript = gr.Textbox(label="Transcript", lines=8, buttons=["copy"])
                stats = gr.Markdown()
        gr.Markdown(
            "Output is lower case without punctuation. Trained on read audiobook speech (LibriSpeech), "
            f"so clean read English works best. Up to {MAX_SECONDS // 60} min of audio."
        )
        gr.Examples(
            examples=[
                ["examples/1284-1180-0000.flac"],
                ["examples/1089-134686-0000.flac"],
                ["examples/2300-131720-0000.flac"],
                ["examples/4507-16021-0047.flac"],
            ],
            inputs=[audio],
            outputs=[transcript, stats],
            fn=transcribe,
            cache_examples=True,
            cache_mode="lazy",
        )
        gr.Markdown(
            "Examples: LibriSpeech test-clean (Panayotov et al., ICASSP 2015, CC BY 4.0). "
            "Use must follow the [Gemma Prohibited Use Policy](https://ai.google.dev/gemma/prohibited_use_policy)."
        )
    run.click(transcribe, inputs=[audio], outputs=[transcript, stats])

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)