# =====================================================================
#  Chép lời tiếng Việt — ZipFormer-150M fine-tuned (Tối ưu tốc độ)
#  Hugging Face Space (Gradio + ZeroGPU + FP16 + Inference Mode)
# =====================================================================
import os
import sys
import types
import subprocess
from pathlib import Path

import torch

# 1. Vô hiệu hóa lỗi torch.set_num_interop_threads khi train.py được import
_orig_set_interop = torch.set_num_interop_threads
def _safe_set_interop(n):
    try:
        _orig_set_interop(n)
    except RuntimeError:
        pass
torch.set_num_interop_threads = _safe_set_interop

import torch.nn.functional as F
import torchaudio
from torch.nn.utils.rnn import pad_sequence
import sentencepiece as spm
import gradio as gr
import spaces  # Thư viện ZeroGPU của Hugging Face

MODEL_REPO = "hynt/ZipFormer-150M-CR-CTC-RNNT-6000h"
FINETUNED_CKPT = "checkpoints/zipformer_best.pt"
FINETUNED_REPO = os.environ.get("FINETUNED_REPO", "").strip()
AUDIO_SCALE = 1.0
MAX_SEC = 90


# ---------- K2 stub ----------
def swoosh_l(x):
    return (F.softplus(x - 4.0) - 0.08 * x - 0.035).to(x.dtype)

def swoosh_r(x):
    return (F.softplus(x - 1.0) - 0.08 * x - 0.313261687).to(x.dtype)

def swoosh_l_forward_and_deriv(x):
    return swoosh_l(x), (torch.sigmoid(x - 4.0) - 0.08).to(x.dtype)

def swoosh_r_forward_and_deriv(x):
    return swoosh_r(x), (torch.sigmoid(x - 1.0) - 0.08).to(x.dtype)

def _need_real_k2(*a, **k):
    raise NotImplementedError("k2 thật chưa được cài: chỉ dùng CTC (PyTorch).")

k2_mock = types.ModuleType("k2")
k2_mock.__path__ = []
k2_mock.__file__ = "k2/__init__.py"
for _n, _f in {
    "swoosh_l": swoosh_l, "swoosh_l_forward": swoosh_l,
    "swoosh_r": swoosh_r, "swoosh_r_forward": swoosh_r,
    "swoosh_l_forward_and_deriv": swoosh_l_forward_and_deriv,
    "swoosh_r_forward_and_deriv": swoosh_r_forward_and_deriv,
    "ctc_loss": _need_real_k2, "rnnt_loss": _need_real_k2,
}.items():
    setattr(k2_mock, _n, _f)
for _n in ["RaggedTensor", "RaggedShape", "Fsa"]:
    setattr(k2_mock, _n, type(_n, (object,), {}))
k2_mock.SymbolTable = type("SymbolTable", (object,),
                            {"from_file": classmethod(lambda cls, f: cls())})
k2_version = types.ModuleType("k2.version")
k2_version.__file__ = "k2/version.py"
k2_version.__version__ = "1.0.0.stub"
k2_mock.version = k2_version
_k2 = types.ModuleType("_k2")
_k2.RaggedInt = type("RaggedInt", (object,), {})
_k2.simple_ragged_index_select = lambda *a, **k: None
sys.modules.update({
    "k2": k2_mock, "k2.version": k2_version,
    "k2.ragged": types.ModuleType("k2.ragged"),
    "k2.fsa": types.ModuleType("k2.fsa"),
    "_k2": _k2,
})


# ---------- Tải pretrained + Icefall ----------
from huggingface_hub import snapshot_download

repo_path = snapshot_download(repo_id=MODEL_REPO,
                              allow_patterns=["*.pt", "*.model", "*.txt", "*.yaml"])

if not os.path.exists("/tmp/icefall"):
    subprocess.run(["git", "clone", "-q", "--depth", "1",
                    "https://github.com/k2-fsa/icefall.git", "/tmp/icefall"], check=True)
for p in ["/tmp/icefall", "/tmp/icefall/egs/librispeech/ASR",
          "/tmp/icefall/egs/librispeech/ASR/zipformer"]:
    if p not in sys.path:
        sys.path.append(p)

import scaling  # noqa: F401
from train import get_model, get_parser, get_params


# ---------- Dựng model & nạp trọng số ----------
import train as _train
_train.get_env_info = lambda: {"k2-version": "stub"}

bpe_files = sorted(Path(repo_path).rglob("bpe.model"))
assert bpe_files, "Không tìm thấy bpe.model trong repo pretrained"
sp = spm.SentencePieceProcessor()
sp.load(str(bpe_files[0]))
VOCAB = sp.get_piece_size()
assert sp.id_to_piece(0) == "<blk>", "blank phải là id 0"

pts = sorted(Path(repo_path).rglob("*.pt"))
ckpt_path = next((p for p in pts if p.name == "pretrained.pt"), None) \
         or next((p for p in pts if p.name.startswith("epoch-")), None) or pts[0]
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
state_dict = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt

ARCH = ("--num-encoder-layers 2,2,4,5,4,2 --feedforward-dim 512,768,1536,2048,1536,768 "
        "--encoder-dim 192,256,512,768,512,256 --encoder-unmasked-dim 192,192,256,320,256,192 "
        "--cnn-module-kernel 31,31,15,15,15,31 --num-heads 4,4,4,8,4,4 "
        "--downsampling-factor 1,2,4,8,4,2 --decoder-dim 512 --joiner-dim 512 --context-size 2 "
        "--query-head-dim 32 --value-head-dim 12 --pos-head-dim 4 --pos-dim 48 "
        "--causal 0 --use-transducer 1 --use-ctc 1").split()

parser = get_parser()
args, _ = parser.parse_known_args(ARCH)
params = get_params()
params.update(vars(args))
params.blank_id = 0
params.vocab_size = VOCAB
params.feature_dim = 80
params.use_transducer = True
params.use_ctc = True
params.use_attention_decoder = False

model = get_model(params)
res = model.load_state_dict(state_dict, strict=False)
assert len(res.missing_keys) == 0, f"Thiếu trọng số: {res.missing_keys[:5]}"

ft_path = None
if FINETUNED_REPO:
    try:
        ft_dir = snapshot_download(repo_id=FINETUNED_REPO, allow_patterns=["*.pt"])
        cand = sorted(Path(ft_dir).rglob("*.pt"))
        ft_path = str(cand[0]) if cand else None
    except Exception as e:
        print(f"Không tải được FINETUNED_REPO={FINETUNED_REPO}: {e}")
if ft_path is None and os.path.exists(FINETUNED_CKPT):
    ft_path = FINETUNED_CKPT

if ft_path:
    st = torch.load(ft_path, map_location="cpu", weights_only=False)
    st = st["model"] if isinstance(st, dict) and "model" in st else st
    r = model.load_state_dict(st, strict=False)
    print(f"Đã nạp checkpoint fine-tune | missing={len(r.missing_keys)} "
          f"unexpected={len(r.unexpected_keys)}")
    USING_FINETUNED = True
else:
    print("CHƯA có checkpoint fine-tune — đang dùng model gốc.")
    USING_FINETUNED = False

# Giữ mô hình trên CPU lúc khởi tạo
model.to("cpu").eval()
N_PARAMS = sum(p.numel() for p in model.parameters()) / 1e6


# ---------- Fbank + giải mã CTC (Tối ưu FP16 & Inference Mode) ----------
def load_fbank(path):
    wav, sr = torchaudio.load(path)
    if wav.size(0) > 1:
        wav = wav.mean(0, keepdim=True)
    if sr != 16000:
        wav = torchaudio.functional.resample(wav, sr, 16000)
    return torchaudio.compliance.kaldi.fbank(
        wav * AUDIO_SCALE, num_mel_bins=80, sample_frequency=16000,
        snip_edges=False, dither=0.0, low_freq=20.0, high_freq=-400.0, energy_floor=1e-10)


def encode(model, x, x_lens):
    if hasattr(model, "forward_encoder"):
        return model.forward_encoder(x, x_lens)
    from icefall.utils import make_pad_mask
    x, x_lens = model.encoder_embed(x, x_lens)
    out, out_lens = model.encoder(x.permute(1, 0, 2), x_lens, make_pad_mask(x_lens))
    return out.permute(1, 0, 2), out_lens


@torch.inference_mode()
def ctc_greedy_decode(model, paths, batch_size=4):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Đẩy model sang GPU và ép kiểu dữ liệu FP16 để tăng tốc tính toán
    model.to(device)
    use_fp16 = (device == "cuda")
    if use_fp16:
        model.half()
    model.eval()

    hyps = []
    for i in range(0, len(paths), batch_size):
        feats = [load_fbank(p) for p in paths[i:i + batch_size]]
        lens = torch.tensor([f.size(0) for f in feats], dtype=torch.long)
        
        x = pad_sequence(feats, batch_first=True).to(device)
        if use_fp16:
            x = x.half()

        enc, enc_lens = encode(model, x, lens.to(device))
        best = model.ctc_output(enc).argmax(-1).cpu()
        for b in range(best.size(0)):
            ids = torch.unique_consecutive(best[b, : int(enc_lens[b])]).tolist()
            hyps.append(sp.decode([t for t in ids if t != 0]))
            
    return hyps


def sentence_case(t: str) -> str:
    t = (t or "").strip().lower()
    return t[0].upper() + t[1:] if t else t


# ---------- Giao diện web (Gradio + ZeroGPU) ----------
@spaces.GPU
def transcribe(audio_path):
    if not audio_path:
        return "", ""
    info = torchaudio.info(audio_path)
    duration = info.num_frames / info.sample_rate
    if duration < 0.3:
        raise gr.Error("Âm thanh quá ngắn (dưới 0,3 giây).")
    if duration > MAX_SEC:
        raise gr.Error(f"Âm thanh dài {duration:.0f} giây, vượt giới hạn {MAX_SEC} giây.")
    
    text = ctc_greedy_decode(model, [audio_path])[0]
    status = f"{duration:.1f} giây · " + (
        "mô hình đã fine-tune" if USING_FINETUNED else
        "⚠️ chưa tìm thấy checkpoint fine-tune, đang dùng model gốc")
    return sentence_case(text), status


CSS = """
.gradio-container { max-width: 720px !important; font-family: 'Be Vietnam Pro', system-ui, sans-serif; }
#out_text textarea { font-family: Georgia, 'Source Serif 4', serif !important; font-size: 1.35rem !important; line-height: 1.5 !important; }
"""

with gr.Blocks(css=CSS, title="Chép lời tiếng Việt") as demo:
    gr.Markdown(
        f"# Chép lời tiếng Việt\n"
        f"Tải lên hoặc ghi âm, mô hình ZipFormer-150M ({N_PARAMS:.1f}M tham số) "
        f"sẽ chuyển giọng nói thành văn bản. Tối đa {MAX_SEC} giây mỗi lần."
    )
    audio_in = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Âm thanh")
    btn = gr.Button("Chép lời", variant="primary")
    out_text = gr.Textbox(label="Văn bản", elem_id="out_text", lines=4)
    out_status = gr.Markdown()
    btn.click(transcribe, inputs=audio_in, outputs=[out_text, out_status])

demo.queue(max_size=20).launch(show_api=False)