#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
from pathlib import Path

from seavad.main import SeaVAD
from pydub import AudioSegment

vad = SeaVAD(sample_rate=16000, sample_width=2, vad_on_needed=2, vad_off_needed=5)

test_audio = "/seasalt-dev-01/chmei/SeaVAD/data/test_audio/16k_en_zh.wav"
out_dir = f"/seasalt-dev-01/chmei/SeaVAD/data/exp/{Path(test_audio).stem}"
os.makedirs(out_dir, exist_ok=True)

with open(test_audio, "rb") as f:
    f.seek(44)
    data = f.read()

seg = vad.get_vad_segments(data)

print(seg)

# Generate audio segments from the `seg` returned from get_vad_segments()
audio = AudioSegment.from_wav("/seasalt-dev-01/chmei/SeaVAD/data/test_audio/16k_en_zh.wav")
for i, (start, end) in enumerate(seg):
    segment = audio[start*1000:end*1000]
    segment.export(f"{out_dir}/segment_{i}.wav", format="wav")
