import parselmouth
import numpy as np
import os
import json

def process_audio_file(file_path):
    file_name_without_extension = file_path.split('/')[-1].split('.')[0]
    gender, text = file_name_without_extension.split('-')
    snd = parselmouth.Sound(file_path)
    pitch = snd.to_pitch()

    time_values = pitch.xs()
    pitch_values = pitch.selected_array['frequency']

    resample_interval = 0.01  # 10ms
    start_time = np.min(time_values)
    end_time = np.max(time_values)
    resampled_times = np.arange(start_time, end_time, resample_interval)

    resampled_pitch_values = []
    for start in resampled_times:
        end = start + resample_interval
        mask = (time_values >= start) & (time_values < end)
        avg_pitch = np.nanmean(pitch_values[mask]) if np.any(mask) else 0
        resampled_pitch_values.append(avg_pitch)

    first_non_zero_index = next((i for i, pitch in enumerate(resampled_pitch_values) if pitch > 0), None)
    last_non_zero_index = next((len(resampled_pitch_values) - i for i, pitch in enumerate(reversed(resampled_pitch_values)) if pitch > 0), len(resampled_pitch_values))

    trimmed_pitch_values = resampled_pitch_values[first_non_zero_index:last_non_zero_index]
    trimmed_times = resampled_times[first_non_zero_index:last_non_zero_index]

    formatted_time_values = []
    for i, time in enumerate(trimmed_times):
        relative_time = i * resample_interval
        hours, remainder = divmod(relative_time, 3600)
        minutes, seconds = divmod(remainder, 60)
        milliseconds = (seconds % 1) * 1000
        formatted_time = "{:02}:{:02}:{:02}:{:02}".format(int(hours), int(minutes), int(seconds), int(milliseconds//10))
        formatted_time_values.append(formatted_time)

    data_list = [{"time": time, "pitch": pitch} for time, pitch in zip(formatted_time_values, trimmed_pitch_values)]

    output_data = {"data": data_list, "text": text}
    output_folder = os.path.join("output", gender)
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    output_file = os.path.join(output_folder, text + ".json")
    with open(output_file, "w") as json_file:
        json.dump(output_data, json_file, ensure_ascii=False)

def process_directory(directory_path):
    for file in os.listdir(directory_path):
        if file.endswith(".wav"):
            file_path = os.path.join(directory_path, file)
            process_audio_file(file_path)

process_directory('audio/2')