"""Isolated Chatterbox GPU comparison; no change to the live hint system."""
from pathlib import Path
import json,time,random
import numpy as np
import soundfile as sf
import torch
from chatterbox.tts_turbo import ChatterboxTurboTTS
root=Path("/mnt/d/Projects/escape-gallery/room_glitch/labs/glitch_odin/.training/teddy-v1/samples")
random.seed(42);np.random.seed(42);torch.manual_seed(42)
torch.set_num_threads(4)
assert torch.cuda.is_available()
x=torch.randn(32,32,device="cuda");x@x;torch.cuda.synchronize()
start=time.perf_counter()
model=ChatterboxTurboTTS.from_pretrained(device="cuda")
load_seconds=time.perf_counter()-start
start=time.perf_counter()
model.prepare_conditionals(str(root/"chatterbox-reference.wav"))
torch.cuda.synchronize();reference_seconds=time.perf_counter()-start
texts=[
"Something has been moving in the arcade after closing time. Find the locked junction box, and see whether your friends have discovered its key.",
"The switch is up. Give the machines four seconds to wake, then bring the matching token to your cabinet.",
"Tell your partner which symbol is flashing. They can move the controls, but only you can see the target."]
parts=[];timings=[]
for i,text in enumerate(texts):
    start=time.perf_counter()
    wav=model.generate(text)
    torch.cuda.synchronize()
    elapsed=time.perf_counter()-start
    samples=wav.squeeze().detach().cpu().numpy()
    sf.write(root/f"chatterbox-hint-{i+1}.wav",samples,model.sr)
    duration=len(samples)/model.sr
    timings.append(dict(text=text,generation_seconds=elapsed,audio_seconds=duration,real_time_factor=elapsed/duration))
    parts.append(samples)
    print(json.dumps(timings[-1]),flush=True)
joined=[]
for i,part in enumerate(parts):
    if i: joined.append(np.zeros(int(model.sr*.45),dtype=np.float32))
    joined.append(part)
sf.write(root/"teddy-chatterbox-turbo.wav",np.concatenate(joined),model.sr)
report=dict(engine="Chatterbox Turbo",reference="chatterbox-reference.wav",model_load_seconds=load_seconds,reference_seconds=reference_seconds,sample_rate=model.sr,device=torch.cuda.get_device_name(0),torch_version=torch.__version__,chunks=timings)
(root/"chatterbox-report.json").write_text(json.dumps(report,indent=2),encoding="utf-8")
print("Comparison complete",flush=True)