import whisper
from textgrid import TextGrid, IntervalTier
import os
from itertools import product
import torch
import gc
import soundfile as sf

def transcribe_audio(audio_path, language=None, model_name='large-v3-turbo', word_timestamps=False):
    """
    使用Whisper模型转录指定音频文件，并将结果输出为TextGrid格式。

    参数:
    audio_path (str): 音频文件的路径。
    model_name (str): Whisper模型的名称，默认为 'large-v3-turbo'。
    word_timestamps (bool): 是否包含单词级时间戳，默认为True。

    返回:
    TextGrid: 包含转录结果的TextGrid对象。
    """
    device = 'cuda'
    model = whisper.load_model(model_name, device=device)

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


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

    # 创建一个空列表来存储转录结果
    transcription_list = []
    
    # 将每个分段添加到层级中
    if not result['segments']:
        return transcription_list
        
    target_seg_tiername = "word_segments" if "word_segments" in result else "segments"

    textname = "word" if "word" in result[target_seg_tiername][0] else "text"
    
    merged_segments = []
    prev_segment = result[target_seg_tiername][0]
    for segment in result[target_seg_tiername][1:]:
        if segment["start"] >= max_duration:
            continue
        if prev_segment['end'] >= segment['start']:  # 检查是否相邻
            prev_segment = {
                'start': prev_segment['start'],
                'end': segment['end'],
                'text': prev_segment[textname] + segment[textname]
            }
        else:
            merged_segments.append(prev_segment)
            prev_segment = segment
    merged_segments.append(prev_segment)

    # 创建 TextGrid 对象
    tg = TextGrid()
    # 创建 IntervalTier 对象用于存储转录结果
    tier = IntervalTier(name='transcription', minTime=0, maxTime=max_duration)

    for segment in merged_segments:
        if segment["start"] >= max_duration:
            continue
        if segment["end"] > max_duration:
            segment["end"] = max_duration
        if not segment[textname].strip():
            continue
        # print(segment['start'], segment['end'], segment[textname])
        tier.add(segment['start'], segment['end'], segment[textname])
    tg.append(tier)

    # 保存 TextGrid 文件，保存到 test_results 目录
    # 获取音频文件所在目录的父目录
    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 + f'_whisper_word_timestamp_{word_timestamps}.TextGrid')
    tg.write(textgrid_path)
    
    return transcription_list, result


if __name__ == '__main__':
    audio_path = r"E:\Corpus\ma\audio\32-4.wav"
    word_timestamps = [True, False]

    results = {}

    for word_timestamps in [False, True]:
        if word_timestamps:
            mode = 'with_timestamps'
        else:
            mode = 'no_timestamps'

        textgrid, result = transcribe_audio(audio_path, language="zh", word_timestamps=word_timestamps)

        results['whisper'] = results.get('whisper', {})
        results['whisper'][mode] = [(segment['start'], segment['end'], segment['text']) for segment in result['segments']]

        print(results['whisper'][mode])

