"""
Kokoro TTS Engine for AX650 NPU.

Pipeline: AXERA Encoder U16 → F0N Mixed FP32 → Decoder S16 → HAR + ISTFT
"""
import os, json, time, numpy as np
import onnxruntime as ort

MODEL_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "models")
MAX_LEN = 96
T_PAD = 192

class KokoroEngine:
    def __init__(self, model_dir=None):
        if model_dir is None:
            model_dir = MODEL_DIR
        self.model_dir = model_dir

        # Load axengine
        from axengine import InferenceSession

        # Load vocabulary
        with open(os.path.join(model_dir, "vocab.json")) as f:
            self.vocab = json.load(f)

        # Load voice
        self.voice = np.load(os.path.join(model_dir, "voice.npy")).astype(np.float32)
        self.style_head = self.voice[:, :128]
        self.style_tail = self.voice[:, 128:]

        # Load models
        self.enc = InferenceSession(os.path.join(model_dir, "kokoro_enc_axera.axmodel"))
        self.f0n = InferenceSession(os.path.join(model_dir, "kokoro_f0n.axmodel"))
        self.dec = InferenceSession(os.path.join(model_dir, "kokoro_dec.axmodel"))
        self.har = ort.InferenceSession(os.path.join(model_dir, "kokoro_har_noup.onnx"),
                                         providers=["CPUExecutionProvider"])
        self.istft = ort.InferenceSession(os.path.join(model_dir, "kokoro_istft.onnx"),
                                           providers=["CPUExecutionProvider"])

        # Sigmoid via frompyfunc (avoids numpy float64 issue on ARM)
        self._sigmoid = np.frompyfunc(lambda x: 1.0 / (1.0 + np.exp(-float(x))), 1, 1)

    def tokenize(self, phonemes):
        """Convert phoneme string to input_ids."""
        ids = [self.vocab[p] for p in phonemes if p in self.vocab][:MAX_LEN - 2]
        input_ids = np.zeros((1, MAX_LEN), dtype=np.int32)
        input_ids[0, 0] = 0
        for i, t in enumerate(ids):
            if i + 1 < MAX_LEN - 1:
                input_ids[0, i + 1] = t
        return input_ids

    def synthesize(self, phonemes, speed=1.0):
        """Generate audio from phoneme string. Returns (audio, sample_rate)."""
        input_ids = self.tokenize(phonemes)

        # 1. Encoder
        d, t_en, dur_logits = self.enc.run(
            ["d", "t_en", "duration_logits"],
            {"input_ids": input_ids, "style_tail": self.style_tail}
        )

        # 2. Duration + Alignment (CPU)
        dt = d.transpose(0, 2, 1)  # [1, 96, 640]
        duration = self._sigmoid(dur_logits).sum(axis=-1).astype(np.float64).squeeze()
        pred_dur = np.round(duration / speed).clip(1).astype(np.int64)
        Ta_raw = int(pred_dur.sum())
        Ta = int(Ta_raw * 224 / 151)  # duration correction factor

        indices = np.repeat(np.arange(MAX_LEN), pred_dur)
        T_out = indices.shape[0]
        alignment = np.zeros((1, MAX_LEN, T_out), dtype=np.float32)
        for j, i in enumerate(indices):
            alignment[0, i, j] = 1.0

        en = (dt @ alignment).astype(np.float32)  # [1, 640, T_out]
        asr = (t_en @ alignment).astype(np.float32)  # [1, 512, T_out]

        # Pad to fixed length
        if en.shape[-1] < T_PAD:
            en = np.pad(en, ((0, 0), (0, 0), (0, T_PAD - en.shape[-1])))
            asr = np.pad(asr, ((0, 0), (0, 0), (0, T_PAD - asr.shape[-1])))
        else:
            en = en[:, :, :T_PAD]
            asr = asr[:, :, :T_PAD]

        # 3. F0N
        F0, N = self.f0n.run(["F0", "N"],
                              {"en": en.astype(np.float32), "style_tail": self.style_tail})

        # 4. HAR (CPU)
        f0_up = np.repeat(F0, 300, axis=-1).astype(np.float32)
        har = self.har.run(["har"], {"f0_up": f0_up})[0]

        # 5. Decoder
        raw_x = self.dec.run(
            ["raw_x"],
            {"asr": asr.astype(np.float32), "F0_pred": F0.astype(np.float32),
             "N_pred": N.astype(np.float32), "ref_s_head": self.style_head,
             "har": har.astype(np.float32)}
        )[0]

        # 6. ISTFT (CPU)
        waveform = self.istft.run(["waveform"], {"raw_x": raw_x.astype(np.float32)})[0]

        audio = waveform.squeeze()[:int(Ta * 300)]
        return audio, 24000