"""Word-by-word Piper: fixed 64-ID encoder window, commit up to a word boundary,
feed back where it got to and re-encode with left context.

Words arrive one at a time (simulating an LLM). After each word we try to run a
window. Audio is committed only up to a word boundary that has LOOKAHEAD_WORDS
whole words after it, so the newest words are never spoken before their right
context is known. When the text ends, everything left is flushed.

Usage: python stream_text_windows.py [noise|nonoise] [ctx_words] [lookahead_words]
"""
import json
import sys
import time
import wave
from pathlib import Path

import numpy as np
import onnxruntime as ort

ROOT = Path(__file__).resolve().parent
N_IDS, N_FRAMES, HOP, DEC_CTX = 64, 160, 256, 12
XFADE = 0  # frames crossfaded between windows; 0 preferred by ear (exact cuts)
CFG = json.loads((ROOT / "data.onnx.json").read_text(encoding="utf-8"))
MAP = CFG["phoneme_id_map"]
RATE = CFG["audio"]["sample_rate"]
PAD, SPACE = MAP["_"][0], MAP[" "][0]

TEXT_IPA = ("ðə pˈaɪpɚ plˈeɪd ɐ tʃˈɪɹfəl tˈuːn, wˌaɪl ðə tʃˈɪldɹən stˈɑːpt tə lˈɪsən. "
            "ðɛn ðeɪ ɔːl dˈæːnst təɡˈɛðɚ ɪntʊ ðə ɡˈoʊldən ˈiːvnɪŋ lˈaɪt.")


def ids_for(ipa: str, bos: bool, eos: bool) -> list[int]:
    out = MAP["^"] + MAP["_"] if bos else []
    for ch in ipa:
        if ch in MAP:
            out += MAP[ch] + MAP["_"]
    return out + (MAP["$"] if eos else [])


def save_wav(path: Path, audio: np.ndarray) -> None:
    pcm = np.round(np.clip(audio, -1, 1) * 32767).astype("<i2")
    with wave.open(str(path), "wb") as w:
        w.setnchannels(1); w.setsampwidth(2); w.setframerate(RATE); w.writeframes(pcm.tobytes())


class WindowedPiper:
    def __init__(self, scales, ctx_words, lookahead_words, seed=1234):
        o = ort.SessionOptions(); o.log_severity_level = 3; o.intra_op_num_threads = 4
        self.enc = ort.InferenceSession(str(ROOT / f"data.encoder.fixed{N_IDS}x{N_FRAMES}.onnx"), o, providers=["CPUExecutionProvider"])
        self.dec = ort.InferenceSession(str(ROOT / "data.decoder.onnx"), o, providers=["CPUExecutionProvider"])
        self.scales = np.asarray(scales, np.float32)
        self.ctx_words, self.look = ctx_words, lookahead_words
        rng = np.random.default_rng(seed)
        # Noise tied to global ID / frame positions so overlapping windows agree.
        self.noise_dp = rng.standard_normal((2, 4096)).astype(np.float32)
        self.noise_z = rng.standard_normal((192, 16384)).astype(np.float32)
        self.ids: list[int] = []      # all IDs received so far
        self.word_starts: list[int] = []  # ID index where each word starts
        self.committed = 0            # first ID not yet spoken
        self.committed_frames = 0     # global frame count spoken
        self.audio: list[np.ndarray] = []
        self.tail = None
        self.log: list[str] = []

    def add_word(self, ipa_word: str, first: bool, last: bool) -> None:
        self.word_starts.append(len(self.ids))
        self.ids += ids_for(ipa_word if first else " " + ipa_word, bos=first, eos=last)
        # A word's IDs start at its leading space; first word starts after "^ _".

    def _context_start(self) -> int:
        # Go back ctx_words word starts from the commit point.
        starts = [s for s in self.word_starts if s < self.committed]
        return starts[-self.ctx_words] if len(starts) >= self.ctx_words else 0

    def step(self, final: bool) -> bool:
        """Run one window; returns True if audio was committed."""
        if self.committed >= len(self.ids):
            return False
        ctx = self._context_start()
        window = self.ids[ctx:ctx + N_IDS]
        n = len(window)
        # Candidate commit points: word starts after the commit point, inside the window.
        boundaries = [s for s in self.word_starts if self.committed < s <= ctx + n]
        whole_after = lambda b: sum(1 for s in self.word_starts if s > b)
        if final and ctx + n == len(self.ids):
            target = len(self.ids)
        else:
            ok = [b for b in boundaries if whole_after(b) >= self.look and b < ctx + n]
            if not ok:
                return False
            target = ok[-1]
        if target <= self.committed:
            return False

        x = np.zeros((1, N_IDS), np.int64); x[0, :n] = window
        ctx_ids = self.committed - ctx
        t0 = time.perf_counter()
        z, dur, frames = self.enc.run(None, {
            "input": x, "input_lengths": np.array([n], np.int64), "scales": self.scales,
            "noise_dp": self.noise_dp[None, :, ctx:ctx + N_IDS],
            "noise_z": self._noise_z_window(dur_ctx_guess=None),
        })
        enc_ms = (time.perf_counter() - t0) * 1000
        cum = np.concatenate([[0], np.cumsum(dur[0, 0, :n])]).astype(int)
        f0, f1 = cum[ctx_ids], cum[target - ctx]
        valid = int(frames[0])
        if f1 > min(valid, N_FRAMES) - (0 if target == len(self.ids) else DEC_CTX):
            # Not enough frames left for decoder right context: commit less next time.
            if target == len(self.ids) or len(boundaries) < 2:
                f1 = min(f1, min(valid, N_FRAMES))
            else:
                return False
        t1 = time.perf_counter()
        limit = min(valid, N_FRAMES)
        # Decode XFADE extra frames past the commit point; they are crossfaded
        # into the start of the next window instead of being played now.
        ext = min(f1 + XFADE, limit)
        a, b = max(0, f0 - DEC_CTX), min(limit, ext + DEC_CTX)
        wav = self.dec.run(None, {"z": z[:, :, a:b]})[0].reshape(-1)
        out = wav[(f0 - a) * HOP:(ext - a) * HOP].copy()
        pcm, tail = out[:(f1 - f0) * HOP], out[(f1 - f0) * HOP:]
        if self.tail is not None and len(self.tail):
            k = min(len(self.tail), len(pcm))
            r = np.linspace(0, 1, k, endpoint=False, dtype=np.float32)
            pcm[:k] = self.tail[:k] * (1 - r) + pcm[:k] * r
        self.tail = tail
        dec_ms = (time.perf_counter() - t1) * 1000
        self.audio.append(pcm)
        self.log.append(f"window ids[{ctx}:{ctx+n}] ctx={ctx_ids} commit ids {self.committed}->{target} "
                        f"frames {f0}->{f1} of {valid}; enc {enc_ms:.0f} ms dec {dec_ms:.0f} ms; "
                        f"{len(pcm)/RATE:.2f} s")
        self.committed = target
        self.committed_frames += f1 - f0
        return True

    def _noise_z_window(self, dur_ctx_guess):
        # Window frame 0 corresponds roughly to the global frame where the context starts.
        # We approximate by the frames already spoken minus a small context allowance;
        # exact alignment only matters when noise is on.
        g0 = max(0, self.committed_frames - 2 * 12 * self.ctx_words)
        return self.noise_z[None, :, g0:g0 + N_FRAMES]


def full_reference(scales):
    o = ort.SessionOptions(); o.log_severity_level = 3
    s = ort.InferenceSession(str(ROOT / "data.onnx"), o, providers=["CPUExecutionProvider"])
    ids = np.asarray([ids_for(TEXT_IPA, True, True)], np.int64)
    return s.run(None, {"input": ids, "input_lengths": np.array([ids.shape[1]]), "scales": np.asarray(scales, np.float32)})[0].reshape(-1)


def main():
    noise = (sys.argv[1] if len(sys.argv) > 1 else "noise") == "noise"
    ctx_words = int(sys.argv[2]) if len(sys.argv) > 2 else 2
    look = int(sys.argv[3]) if len(sys.argv) > 3 else 1
    scales = [0.667, 1.0, 0.8] if noise else [0.0, 1.0, 0.0]
    words = TEXT_IPA.split(" ")
    p = WindowedPiper(scales, ctx_words, look)
    first_audio_word = None
    for i, w in enumerate(words):
        p.add_word(w, first=i == 0, last=i == len(words) - 1)
        p.log.append(f"+ word {i+1}: {w}")
        while p.step(final=False):
            if first_audio_word is None:
                first_audio_word = i + 1
    while p.step(final=True):
        pass
    audio = np.concatenate(p.audio)
    tag = f"{'noise' if noise else 'nonoise'}_c{ctx_words}_l{look}"
    save_wav(ROOT / f"words_stream_{tag}.wav", audio)
    ref = full_reference(scales)
    save_wav(ROOT / f"words_full_{tag}.wav", ref)
    print("\n".join(p.log))
    print(f"first audio after word {first_audio_word} of {len(words)}")
    print(f"stream {len(audio)/RATE:.2f} s vs full {len(ref)/RATE:.2f} s")
    if not noise:
        k = min(len(audio), len(ref))
        print(f"length diff {len(audio)-len(ref)} samples; rel diff (first {k}) "
              f"{np.linalg.norm(audio[:k]-ref[:k])/np.linalg.norm(ref[:k]):.4f}")


if __name__ == "__main__":
    main()