#!/usr/bin/env python3
"""
Generate talker codec tokens on PC using the full Python pipeline,
then save them for the tablet to decode (CP + VQ + Decoder).

This bypasses the broken talker ONNX/GGUF export and proves
that the rest of the tablet pipeline (CP + decoder) works.

Output:
  - talker_output.bin: binary file with cb0 tokens + hidden states
    Format: int32 n_tokens, then for each token:
      int32 cb0, float32[1024] hidden_state
"""
import sys
sys.path.insert(0, "/opt/Kazeia/qnn_venv/lib/python3.10/site-packages")
import warnings; warnings.filterwarnings("ignore")
import torch, numpy as np, struct

model_path = "/home/alf/.cache/huggingface/hub/models--Qwen--Qwen3-TTS-12Hz-0.6B-Base/snapshots/5d83992436eae1d760afd27aff78a71d676296fc"

from qwen_tts.core.models.modeling_qwen3_tts import Qwen3TTSForConditionalGeneration
from transformers import AutoTokenizer

print("Loading model...")
model = Qwen3TTSForConditionalGeneration.from_pretrained(model_path, dtype=torch.float32, device_map='cpu')
model.eval()
tokenizer = AutoTokenizer.from_pretrained(model_path)
spk = torch.from_numpy(np.load('models_qnn/qwen3-tts-embeddings/damien_spk_embedding.npy')).float()

TEXT = "Bonjour, je m'appelle Kazeia, je suis encore en phase de développement."

# Hook to capture hidden states from talker
hidden_states = []
cb0_tokens = []

orig_forward = model.talker.model.forward

def capture_forward(*args, **kwargs):
    result = orig_forward(*args, **kwargs)
    # result is a BaseModelOutputWithPast — hidden states are in last_hidden_state
    if hasattr(result, 'last_hidden_state'):
        h = result.last_hidden_state
        if h.shape[1] == 1:  # decode step
            hidden_states.append(h.detach().squeeze().cpu().numpy())
    return result

model.talker.model.forward = capture_forward

# Also hook codec_head to capture cb0 tokens
orig_codec_head = model.talker.codec_head

class HookedCodecHead(torch.nn.Module):
    def __init__(self, head):
        super().__init__()
        self.head = head
    def forward(self, x):
        logits = self.head(x)
        return logits

model.talker.codec_head = HookedCodecHead(orig_codec_head)

prompt = f"<|im_start|>assistant\n{TEXT}<|endoftext|>\n<|im_start|>assistant\n"
input_ids = tokenizer(prompt, return_tensors="pt").input_ids

print(f"Generating: '{TEXT}'")
print(f"Input tokens: {input_ids.shape[1]}")

torch.manual_seed(42)
with torch.no_grad():
    result = model.generate(
        input_ids=[input_ids],
        languages=["french"],
        speakers=[""],
        voice_clone_prompt={
            "x_vector_only_mode": [True],
            "icl_mode": [False],
            "ref_code": None,
            "ref_spk_embedding": [spk],
            "x_vector": [spk],
        },
        do_sample=True,
        temperature=0.9,
        top_k=50,
        repetition_penalty=1.05,
        max_new_tokens=200,
    )

# result is a list of [audio_codes] tensors
codes = result[0][0]  # [n_tokens, 16] (16 codebooks)
n_tokens = codes.shape[0]
print(f"Generated {n_tokens} tokens × {codes.shape[1]} codebooks")
print(f"Duration: {n_tokens/12:.2f}s")
print(f"First 5 cb0: {codes[:5, 0].tolist()}")
print(f"Hidden states captured: {len(hidden_states)}")

# Save ALL 16 codebooks directly — tablet doesn't need to run CP at all!
# Format: int32 n_tokens, int32 n_codebooks, then [n_tokens × n_codebooks] int32
OUT = "/opt/Kazeia/tablet_codes.bin"
with open(OUT, "wb") as f:
    f.write(struct.pack("<ii", n_tokens, codes.shape[1]))
    for t in range(n_tokens):
        for cb in range(codes.shape[1]):
            f.write(struct.pack("<i", int(codes[t, cb].item())))

print(f"\nSaved {n_tokens} × {codes.shape[1]} codes to {OUT}")
print(f"File size: {4*2 + n_tokens*codes.shape[1]*4} bytes")

# Also save as numpy for easy loading
np.save("/opt/Kazeia/tablet_codes.npy", codes.numpy())

# Decode locally for comparison WAV
print("\nDecoding locally...")
speech_tokenizer_path = "/home/alf/.cache/huggingface/hub/models--Qwen--Qwen3-TTS-Tokenizer-12Hz/snapshots/7dd38ad4e9bad454aae9cd937d0cd577604fe229"
from qwen_tts.core.tokenizer_25hz.modeling_qwen3_tts_tokenizer_v1 import Qwen3TTSTokenizerV1Config, Qwen3TTSTokenizerV1Model
import json
from safetensors.torch import load_file

with open(f"{speech_tokenizer_path}/config.json") as f:
    tok_cfg = json.load(f)
tok_config = Qwen3TTSTokenizerV1Config(**tok_cfg)
tok_model = Qwen3TTSTokenizerV1Model(tok_config)
tok_state = load_file(f"{speech_tokenizer_path}/model.safetensors")
tok_model.load_state_dict(tok_state, strict=False)
tok_model.eval()

# Decode codes → audio
with torch.no_grad():
    audio_codes = codes.unsqueeze(0).long()  # [1, n_tokens, 16]
    # The decode method needs xvectors and ref_mels, but for x_vector_only_mode
    # we can try direct decode
    # Actually, the tokenizer decode takes audio_codes and produces waveform
    dummy_xvec = spk.unsqueeze(0)
    dummy_mel = torch.zeros(1, 1, 128)
    output = tok_model.decode(audio_codes[:, :, 0], dummy_xvec, dummy_mel)

print(f"Decode output type: {type(output)}")
if hasattr(output, 'audio_values') or isinstance(output, tuple):
    audio = output[0] if isinstance(output, tuple) else output.audio_values
    if isinstance(audio, list):
        audio = audio[0]
    audio_np = audio.numpy() if hasattr(audio, 'numpy') else np.array(audio)
    print(f"Audio: {audio_np.shape}, range [{audio_np.min():.3f}, {audio_np.max():.3f}]")
    import soundfile as sf
    sf.write("/opt/Kazeia/kazeia_tablet_codes_decoded_pc.wav", audio_np, 24000)
    print("Saved decoded WAV")