"""Kokoro-7M ONNX CPU & af_msa Voice Test.

Usage:
    python test.py
    python test.py --text "Custom sentence to synthesize"
"""
import argparse
import os
import sys
import time
import numpy as np
import onnxruntime as ort
import soundfile as sf
from kokoro import KPipeline

sys.stdout.reconfigure(encoding="utf-8")
HERE = os.path.dirname(os.path.abspath(__file__))

# 114-symbol Kokoro phonetic vocabulary mapping
VOCAB = {
    ';': 1, ':': 2, ',': 3, '.': 4, '!': 5, '?': 6, '—': 9, '…': 10, '"': 11,
    '(': 12, ')': 13, '“': 14, '”': 15, ' ': 16, '̃': 17, 'ʣ': 18, 'ʥ': 19,
    'ʦ': 20, 'ʨ': 21, 'ᵝ': 22, 'ꭧ': 23, 'A': 24, 'I': 25, 'O': 31, 'Q': 33,
    'S': 35, 'T': 36, 'W': 39, 'Y': 41, 'ᵊ': 42, 'a': 43, 'b': 44, 'c': 45,
    'd': 46, 'e': 47, 'f': 48, 'h': 50, 'i': 51, 'j': 52, 'k': 53, 'l': 54,
    'm': 55, 'n': 56, 'o': 57, 'p': 58, 'q': 59, 'r': 60, 's': 61, 't': 62,
    'u': 63, 'v': 64, 'w': 65, 'x': 66, 'y': 67, 'z': 68, 'ɑ': 69, 'ɐ': 70,
    'ɒ': 71, 'æ': 72, 'β': 75, 'ɔ': 76, 'ɕ': 77, 'ç': 78, 'ɖ': 80, 'ð': 81,
    'ʤ': 82, 'ə': 83, 'ɚ': 85, 'ɛ': 86, 'ɜ': 87, 'ɟ': 90, 'ɡ': 92, 'ɥ': 99,
    'ɨ': 101, 'ɪ': 102, 'ʝ': 103, 'ɯ': 110, 'ɰ': 111, 'ŋ': 112, 'ɳ': 113,
    'ɲ': 114, 'ɴ': 115, 'ø': 116, 'ɸ': 118, 'θ': 119, 'œ': 120, 'ɹ': 123,
    'ɾ': 125, 'ɻ': 126, 'ʁ': 128, 'ɽ': 129, 'ʂ': 130, 'ʃ': 131, 'ʈ': 132,
    'ʧ': 133, 'ʊ': 135, 'ʋ': 136, 'ʌ': 138, 'ɣ': 139, 'ɤ': 140, 'χ': 142,
    'ʎ': 143, 'ʒ': 147, 'ʔ': 148, 'ˈ': 156, 'ˌ': 157, 'ː': 158, 'ʰ': 162,
    'ʲ': 164, '↓': 169, '→': 171, '↗': 172, '↘': 173, 'ᵻ': 177
}


def main():
    parser = argparse.ArgumentParser(description="Test Kokoro-7M ONNX model with af_msa voice.")
    parser.add_argument(
        "--model",
        default=None,
        help="Path to Kokoro ONNX model (defaults to kokoro_7m.onnx, or kokoro_7m_int8.onnx if --int8 is passed)."
    )
    parser.add_argument(
        "--int8",
        action="store_true",
        help="Use INT8 dynamic quantized model (kokoro_7m_int8.onnx)."
    )
    parser.add_argument(
        "--voice",
        default=os.path.join(HERE, "af_msa.onnx"),
        help="Path to af_msa voice ONNX model."
    )
    parser.add_argument(
        "--text",
        default="Hello! This is Kokoro-7M speaking with the af_msa voice on ONNX CPU.",
        help="Text to synthesize."
    )
    parser.add_argument(
        "--output",
        default=os.path.join(HERE, "out.wav"),
        help="Output WAV audio file."
    )
    parser.add_argument(
        "--speed",
        type=float,
        default=1.0,
        help="Speech speed factor."
    )
    args = parser.parse_args()

    # Select model
    if args.model is None:
        model_file = "kokoro_7m_int8.onnx" if args.int8 else "kokoro_7m.onnx"
        model_path = os.path.join(HERE, model_file)
    else:
        model_path = args.model

    print("=" * 60)
    print("Kokoro-7M ONNX Model & Voice Verification")
    print("=" * 60)
    print(f"Model:  {model_path} ({os.path.getsize(model_path) / (1024*1024):.2f} MB)")
    print(f"Voice:  {args.voice}")
    print(f"Text:   \"{args.text}\"")

    # 1. Initialize ONNX Runtime sessions for both Model and Voice
    print("\n[1/3] Initializing ONNX Runtime CPU sessions...")
    sess_opts = ort.SessionOptions()
    sess_opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

    model_sess = ort.InferenceSession(model_path, sess_opts, providers=["CPUExecutionProvider"])
    voice_sess = ort.InferenceSession(args.voice, sess_opts, providers=["CPUExecutionProvider"])

    # 2. Phonemize
    print("[2/3] Phonemizing input text...")
    pipeline = KPipeline(lang_code="a", repo_id="oddadmix/Kokoro-7M-Distill", model=False)
    _, tokens = pipeline.g2p(args.text)

    # 3. Synthesize
    print("[3/3] Generating audio...")
    t0 = time.perf_counter()
    chunks = []
    silence = np.zeros(int(24000 * 0.15), dtype=np.float32)

    for _, ps, _ in pipeline.en_tokenize(tokens):
        if not ps:
            continue
        if len(ps) > 510:
            ps = ps[:510]

        # Map phonemes to token IDs
        input_ids = [VOCAB[p] for p in ps if p in VOCAB]
        input_ids_arr = np.array([[0] + input_ids + [0]], dtype=np.int64)

        # Style lookup via voice ONNX model
        seq_len = input_ids_arr.shape[1]
        style_idx = np.array([min(seq_len - 1, 509)], dtype=np.int64)
        style_vec = voice_sess.run(None, {"index": style_idx})[0]

        # Audio generation via model ONNX
        audio_chunk = model_sess.run(None, {
            "input_ids": input_ids_arr,
            "style": style_vec.astype(np.float32),
            "speed": np.array([args.speed], dtype=np.float32)
        })[0].squeeze()

        chunks.append(audio_chunk)
        chunks.append(silence)

    t_elapsed = time.perf_counter() - t0

    if not chunks:
        print("Error: No audio was generated.")
        return

    full_audio = np.concatenate(chunks[:-1]) if len(chunks) > 1 else chunks[0]
    dur = len(full_audio) / 24000
    rtf = t_elapsed / dur

    # Save audio
    sf.write(args.output, full_audio, 24000)

    # Quality and integrity checks
    has_nan = bool(np.isnan(full_audio).any())
    has_inf = bool(np.isinf(full_audio).any())
    peak = float(np.max(np.abs(full_audio)))
    rms = float(np.sqrt(np.mean(full_audio ** 2)))

    print("\n" + "=" * 60)
    print("Verification Summary:")
    print("=" * 60)
    print(f"Output File:     {args.output}")
    print(f"Duration:        {dur:.2f} s ({len(full_audio):,} samples @ 24 kHz)")
    print(f"Inference Time:  {t_elapsed * 1000:.1f} ms")
    print(f"RTF:             {rtf:.4f} ({1/rtf:.1f}x real-time speed)")
    print(f"Peak Amplitude:  {peak:.3f}")
    print(f"RMS Energy:      {rms:.4f}")
    print(f"Status:          PASSED [OK] (NaN={has_nan}, Inf={has_inf})")
    print("=" * 60)


if __name__ == "__main__":
    main()