import whisperx
import torch
import os
from textgrid import TextGrid, IntervalTier
import soundfile as sf
# torch.backends.cuda.matmul.allow_tf32 = False
# torch.backends.cudnn.allow_tf32 = False

def transcribe_audio(audio_path, language=None):
    # 设置设备
    device = "cuda" if torch.cuda.is_available() else "cpu"
    batch_size = 16  # 根据你的显存调整批次大小
    compute_type = "float16" if torch.cuda.is_available() else "float32"

    # 加载ASR模型
    model = whisperx.load_model("large-v3-turbo", device, compute_type=compute_type)

    # 转录音频
    if language is None:
        result = model.transcribe(audio_path, batch_size=batch_size)
    else:
        result = model.transcribe(audio_path, batch_size=batch_size, language=language)

    # 加载对齐模型并对齐转录结果

    model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=device)
    result = whisperx.align(result["segments"], model_a, metadata, audio_path, device)


    # 获取音频文件的最大时长
    with sf.SoundFile(audio_path) as f:
        max_duration = len(f) / f.samplerate

    # 打印转录结果
    # 初始化合并后的词段列表
    merged_segments = []
    if result.get("word_segments"):
        prev_segment = result["word_segments"][0]
        for segment in result["word_segments"][1:]:
            if segment["start"] >= max_duration:
                continue
            if prev_segment["end"] >= segment["start"]:  # 检查前后两个seg是否相邻
                # 合并相邻词段

                merged_segment = {
                    "start": prev_segment["start"],
                    "end": segment["end"],
                    "word": prev_segment["word"] + segment["word"]
                }
                prev_segment = merged_segment
            else:
                merged_segments.append(prev_segment)
                prev_segment = segment
        merged_segments.append(prev_segment)
    # print(merged_segments)

    # 获取音频文件所在目录的父目录
    parent_dir = os.path.dirname(os.path.dirname(audio_path))
    # 拼接 test_results 目录路径
    test_results_dir = os.path.join(parent_dir, 'test_results')
    # 确保 test_results 目录存在
    os.makedirs(test_results_dir, exist_ok=True)
    # 从音频文件名获取基础名称
    base_name = os.path.basename(os.path.splitext(audio_path)[0])
    # 拼接完整的 TextGrid 文件路径，保存到 test_results 目录
    textgrid_path = os.path.join(test_results_dir, base_name + '_whisperx.TextGrid')

    # 创建 TextGrid 对象
    tg = TextGrid()
    # 创建 IntervalTier 对象
    tier = IntervalTier(name='words', minTime=0, maxTime=max_duration)



    # 添加每个合并后的词段到 IntervalTier
    for segment in merged_segments:
        if segment["start"] >= max_duration:
            continue
        if segment["end"] > max_duration:
            segment["end"] = max_duration
        if not segment["word"].strip():
            continue
        tier.add(segment['start'], segment['end'], segment['word'])


    # 将 IntervalTier 添加到 TextGrid
    tg.append(tier)
    # 保存 TextGrid 文件
    tg.write(textgrid_path)
    return result

if __name__ == "__main__":
    audio_path = r'C:\Users\User\Desktop\Praasper\data\test_audio.wav'  # 替换为你的音频文件路径
    transcribe_audio(audio_path)