from pathlib import Path

import gradio as gr
import librosa
import numpy as np
import spaces
import torch
from loguru import logger
from transformers import AutoModelForCTC, AutoProcessor

REPO_ID_LIST = [
    "AndrewMcDowell/wav2vec2-xls-r-1b-japanese-hiragana-katakana",
    "natsuo/ja_hiragana",
    "prj-beatrice/japanese-hubert-base-phoneme-ctc",
    "prj-beatrice/japanese-hubert-base-phoneme-ctc-v2",
    "slplab/wav2vec2-xls-r-300m-japanese-hiragana",
    "snu-nia-12/wav2vec2-xls-r-300m_nia12_phone-hiragana_japanese",
    "thunninoi/wav2vec2-japanese-hiragana-vtuber",
    "vumichien/wav2vec2-large-xlsr-japanese-hiragana",
]
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
logger.info("Loading models...")

MODEL_DICT = {}
for repo_id in REPO_ID_LIST:
    model = AutoModelForCTC.from_pretrained(repo_id).to(DEVICE)
    processor = AutoProcessor.from_pretrained(repo_id)
    MODEL_DICT[repo_id] = (model, processor)
    logger.info(f"Loaded model from {repo_id}")

logger.success("All models loaded successfully")


@spaces.GPU
@torch.inference_mode()
def transcribe(audio_path: Path) -> list[str]:
    duration = librosa.get_duration(path=audio_path)
    logger.info(f"audio: {Path(audio_path).name}, duration: {duration:.2f} seconds")
    if duration > 30:
        raise gr.Error("Audio duration exceeds 30 seconds")

    y, sr = librosa.load(audio_path, sr=16000)
    y = np.concatenate([np.zeros(sr), y, np.zeros(sr // 2)])

    results = []
    for repo_id in REPO_ID_LIST:
        logger.info(f"Transcribing with model: {repo_id}")
        model, processor = MODEL_DICT[repo_id]
        inputs = processor(y, sampling_rate=sr, return_tensors="pt").to(DEVICE)
        with torch.no_grad():
            outputs = model(**inputs)
            predicted_ids = outputs.logits.argmax(-1)
            phonemes = processor.decode(
                predicted_ids[0], spaces_between_special_tokens=True
            )
        logger.info(f"Transcription for {repo_id}: {phonemes}")
        results.append(phonemes)
    return results


md = """
# Japanese Phoneme ASR CTC Comparison

日本語の音素単位のCTC ASRモデルの比較デモ。


"""

for repo_id in REPO_ID_LIST:
    md += f"- [{repo_id}](https://huggingface.co/{repo_id})\n"

with gr.Blocks() as demo:
    gr.Markdown(md)
    audio_input = gr.Audio(label="Audio Input", type="filepath")
    transcribe_button = gr.Button("Transcribe")
    outputs = []
    with gr.Column():
        for repo_id in REPO_ID_LIST:
            output_text = gr.Textbox(label=f"{repo_id}")
            outputs.append(output_text)
    transcribe_button.click(transcribe, inputs=[audio_input], outputs=outputs)

demo.launch()