# -*- coding: utf-8 -*-
"""Minimal inference for the Ewe Spark-TTS model."""
import re, sys, torch, numpy as np, soundfile as sf
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL = "."                 # this repo
SPARK_SRC = "./Spark-TTS"   # git clone https://github.com/SparkAudio/Spark-TTS

sys.path.append(SPARK_SRC)
from sparktts.models.audio_tokenizer import BiCodecTokenizer
from ewe_text import normalize_ewe

dev = "cuda" if torch.cuda.is_available() else "cpu"
tok = AutoTokenizer.from_pretrained(f"{MODEL}/LLM")
llm = AutoModelForCausalLM.from_pretrained(
    f"{MODEL}/LLM", torch_dtype=torch.float32).to(dev).eval()
audio = BiCodecTokenizer(MODEL, dev)

stops = {tok.convert_tokens_to_ids(t)
         for t in ["<|im_end|>", "<|end_semantic_token|>", "<|endoftext|>"]}
stops = sorted(i for i in stops if isinstance(i, int) and i >= 0)


@torch.inference_mode()
def say(text, temperature=0.55, max_seconds=20):
    # Normalisation is required — see the model card.
    clean = normalize_ewe(text, digit_policy="verbalize",
                          strip_verse_numbers=False)
    prompt = ("<|task_tts|><|start_content|>" + clean +
              "<|end_content|><|start_global_token|>")
    ins = tok([prompt], return_tensors="pt").to(dev)
    out = llm.generate(**ins, max_new_tokens=int(max_seconds * 50) + 128,
                       do_sample=True, temperature=temperature,
                       top_k=50, top_p=0.95, eos_token_id=stops,
                       pad_token_id=tok.pad_token_id or tok.eos_token_id)
    txt = tok.batch_decode(out[:, ins.input_ids.shape[1]:],
                           skip_special_tokens=False)[0]
    sem = re.findall(r"<\|bicodec_semantic_(\d+)\|>", txt)
    glo = re.findall(r"<\|bicodec_global_(\d+)\|>", txt)
    if not sem:
        raise RuntimeError("no semantic tokens generated")
    return audio.detokenize(
        torch.tensor([int(i) for i in glo or ["0"]]).long().unsqueeze(0).to(dev),
        torch.tensor([int(i) for i in sem]).long().unsqueeze(0).to(dev))


if __name__ == "__main__":
    wav = say("Ŋdi na mi. Nye ŋkɔe nye Kofi, eye medzea Eʋegbe nyuie.")
    sf.write("out.wav", wav, audio.config.get("sample_rate", 16000))
    print("wrote out.wav")