import os
import sys
current_dir = os.path.dirname(os.path.abspath(__file__))
print('add current dir to sys.path', current_dir)
sys.path.append(current_dir)
from sparktts.models.audio_tokenizer import BiCodecTokenizer
from transformers import AutoTokenizer, AutoModelForCausalLM
import soundfile as sf
import numpy as np
import torch
from utilities import generate_embeddings_batch

def generate_speech_batch(model, tokenizer, texts, bicodec, prompt_text=None, prompt_audio=None, 
                         max_new_tokens=3000, do_sample=True, top_k=50, top_p=0.95, 
                         temperature=1.0, device="cuda:0"):
    """
    批量生成语音的函数
    
    Args:
        model: 语言模型
        tokenizer: 文本分词器
        texts: 要生成语音的文本列表
        bicodec: BiCodecTokenizer 实例
        prompt_text: 提示文本（可选）
        prompt_audio: 提示音频数组（可选）
        max_new_tokens: 最大生成token数
        do_sample: 是否使用采样
        top_k: top-k采样参数
        top_p: top-p采样参数
        temperature: 温度参数
        device: 设备
    
    Returns:
        list: 生成的音频波形列表
    """
    # 设置eos_token_id - 根据训练代码，eos_token_id = model.config.vocab_size - 1
    eos_token_id = model.config.vocab_size - 1
    print(f"EOS token ID: {eos_token_id}")
    
    # 生成输入嵌入
    embeddings, attention_mask = generate_embeddings_batch(
        model=model,
        tokenizer=tokenizer,
        texts=texts,
        bicodec=bicodec,
        prompt_text=prompt_text,
        prompt_audio=prompt_audio
    )
    
    print("开始批量生成语音...")
    print(f"输入嵌入形状: {embeddings['input_embs'].shape}")
    print(f"注意力掩码形状: {attention_mask.shape}")
    print(f'embeddings dtype: {embeddings["input_embs"].dtype}')
    batch_size = len(texts)
    global_tokens = embeddings['global_tokens']  # [batch_size, global_token_len]
    
    # 设置模型为评估模式
    model.eval()
    
    with torch.no_grad():
        # 使用模型的generate方法进行批量生成
        generated_outputs = model.generate(
            inputs_embeds=embeddings['input_embs'],
            attention_mask=attention_mask,
            max_new_tokens=max_new_tokens,
            do_sample=do_sample,
            top_k=top_k,
            top_p=top_p,
            temperature=temperature,
            eos_token_id=eos_token_id,
            pad_token_id=tokenizer.pad_token_id if hasattr(tokenizer, 'pad_token_id') else tokenizer.eos_token_id,
            use_cache=True
        )
    
    print(f"生成的token形状: {generated_outputs.shape}")
    print(f'generated_tokens :{generated_outputs.tolist()}')
    # 处理每个样本的生成结果
    wavs = []
    eos_index = torch.where(generated_outputs[:,] == eos_token_id)[1]
    print(f'eos_index :{eos_index}')
    for i in range(batch_size):
        # 获取当前样本的生成结果
        sample_outputs = generated_outputs[i]
        print(f'sample_outputs :{sample_outputs.tolist()}')
        #find the first eos_token_id
        eos_token_id_index = eos_index[i]
        print(f'eos_token_id_index :{eos_token_id_index}')
        sample_outputs = sample_outputs[:eos_token_id_index]
        print(f'sample_outputs :{sample_outputs.tolist()}')
        
        print(f"样本 {i} - Semantic tokens shape: {sample_outputs.shape}")
        print(f"样本 {i} - Global tokens shape: {global_tokens[i:i+1].shape}")
        
        # 使用BiCodec解码生成音频
        with torch.no_grad():
            wav = bicodec.detokenize(global_tokens[i:i+1], sample_outputs.unsqueeze(0))
        
        print(f"样本 {i} - 生成的音频形状: {wav.shape}")
        wavs.append(wav)  # 移除批次维度
    
    return wavs

device = 'cuda:2'

audio_tokenizer = BiCodecTokenizer(model_dir=current_dir, device=device)

print(audio_tokenizer)

tokenizer = AutoTokenizer.from_pretrained(current_dir, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(current_dir, trust_remote_code=True)
print(tokenizer)
print(model)

model = model.bfloat16().to(device)
model.eval()

prompt_text = "我们并不是通过物理移动手段找到星河的。"
prompt_audio_file = os.path.join(current_dir, 'kafka.wav')
prompt_audio, sampling_rate = sf.read(prompt_audio_file)

print(f"Loaded prompt audio from {prompt_audio_file}")
print(f"Original sampling rate: {sampling_rate}Hz")
print(f"Audio shape: {prompt_audio.shape}")

target_sample_rate = audio_tokenizer.config['sample_rate']
if sampling_rate != target_sample_rate:
    print(f"Resampling from {sampling_rate}Hz to {target_sample_rate}Hz...")
    from librosa import resample
    prompt_audio = resample(prompt_audio, orig_sr=sampling_rate, target_sr=target_sample_rate)
    prompt_audio = np.array(prompt_audio, dtype=np.float32)
    print(f"Resampled audio shape: {prompt_audio.shape}")
else:
    print(f"Audio sampling rate already matches target ({target_sample_rate}Hz)")

texts = [
    "一九五二年二月十日，志愿军大英雄张积慧击落美军双料王牌飞行员戴维斯，在自己飞机坠毁处距离戴维斯坠机处不足五百米的情况下，取得了世界空战史不可能复制的奇迹。伟大的张积慧。",
    "在数字浪潮汹涌的今天，数智技术正以前所未有的力量重塑着社会的每一个角落。",
    "为了点燃青少年对科技的热情，培养他们的创新思维与动手能力",
    "杏花岭区巨轮街道社区教育学校携手中车社区教育分校，与太原市科学技术协会联手，于暑期精心策划了一场别开生面的青少年数智技术服务港探索之旅，吸引了众多社区青少年的积极参与。"
    "一踏入数智技术服务港的大门，一股浓厚的科技气息便扑面而来。",
    "科普课堂上，“简易红绿灯”科学实验更是将抽象的电路原理与日常生活紧密相连。",
    "实验开始前，老师生动地介绍了实验物品，并引导青少年思考红绿灯的工作原理，激发了他们浓厚的探索兴趣。",
    "在老师的指导下，青少年们开始动手组装电路，将红绿灯的各个部件连接起来。",
    "他们小心翼翼地调整电路，确保每个部件都正确连接，红灯、绿灯、黄灯依次亮起，仿佛在讲述一个关于交通规则的故事。",
    "实验过程中，青少年们不仅学到了电路知识，还体验到了动手实践的乐趣。",
    "他们纷纷表示，这次实验不仅让他们对科技有了更深的理解，还培养了他们的创新思维和动手能力。",
    "数智技术服务港，让科技触手可及，让创新无处不在。"

]

wavs = generate_speech_batch(model, tokenizer, texts, audio_tokenizer, prompt_audio=prompt_audio, device=device)
for i, wav in enumerate(wavs):
    sf.write(f'output_{i}.wav', wav, target_sample_rate)