import parselmouth

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

# import math

sns.set()  # Use seaborn's default style to make attractive graphs
# Assume you have a variable with the file path
file_path = 'audio/Beginning and end.wav'  # This is your file path
# file_path = 'audio/We would like to go Blow up the balloon.wav'
# Store only the file name without the extension
file_name_without_extension = file_path.split('/')[-1].split('.')[0]
snd = parselmouth.Sound(file_path)
# snd = parselmouth.Sound(r"6902739Z009014628.wav")
def draw_pitch(pitch):
    # Extract selected pitch contour, and
    # replace unvoiced samples by NaN to not plot
    pitch_values = pitch.selected_array['frequency']
    pitch_values[pitch_values == 0] = np.nan
    plt.plot(pitch.xs(), pitch_values, 'o', markersize=2, color='w')
    plt.plot(pitch.xs(), pitch_values, 'o', markersize=2)
    plt.grid(False)
    plt.ylim(0, pitch.ceiling)
    plt.ylabel("fundamental frequency [Hz]")

# def draw_spectrogram(spectrogram, dynamic_range=70):
#     X, Y = spectrogram.x_grid(), spectrogram.y_grid()
#     sg_db = 10 * np.log10(spectrogram.values)
#     plt.pcolormesh(X, Y, sg_db, vmin=sg_db.max() - dynamic_range, cmap='afmhot')
#     plt.ylim([spectrogram.ymin, spectrogram.ymax])
#     plt.xlabel("time [s]")
#     plt.ylabel("frequency [Hz]")

# def draw_intensity(intensity):
#     plt.plot(intensity.xs(), intensity.values.T, linewidth=3, color='w')
#     plt.plot(intensity.xs(), intensity.values.T, linewidth=1)
#     plt.grid(False)
#     plt.ylim(0)
#     plt.ylabel("intensity [dB]")

class Tonality:
    base_diff = [0, 2, 4, 5, 7, 9, 11]
    tones = ['A', 'A#', 'B', 'C', 'C#', 'D', "D#", 'E', 'F', 'F#', 'G', 'G#']
    # tones = ['C']

    def __init__(self, tone):
        assert tone in self.tones
        self.base_freq = self.tone_to_freq(tone)

    def tone_to_freq(self, tone):
        A_freq = 440
        return A_freq * (2 ** (self.tones.index(tone) / 12))

    def get_freq(self):
        ret = []
        for r in range(-3, 2):
            base = self.base_freq * (2 ** r)
            ret.extend([base * (2 ** (diff / 12)) for diff in self.base_diff])
        return ret

def draw_standard(tone):
    for f in Tonality(tone).get_freq():
        plt.axline((0, f), (1, f), lw=1)

pitch = snd.to_pitch()
# If desired, pre-emphasize the sound fragment before calculating the spectrogram
# pre_emphasized_snd = snd.copy()
# pre_emphasized_snd.pre_emphasize()
# spectrogram = pre_emphasized_snd.to_spectrogram(window_length=0.03, maximum_frequency=8000)
# plt.figure()
# # draw_spectrogram(spectrogram)
# plt.twinx()

tone = 'C'
draw_standard(tone)
draw_pitch(pitch)
# plt.xlim([snd.xmin, snd.xmax])
# plt.show()  # or plt.savefig("spectrogram_0.03.pdf")

import json
import numpy as np

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

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

# Calculate resampled pitch values
resampled_pitch_values = []
for start in resampled_times:
    end = start + resample_interval
    mask = (time_values >= start) & (time_values < end)
    # Note the use of time_values instead of pitch.xs() to ensure we are using the correct time values
    avg_pitch = np.nanmean(pitch_values[mask]) if np.any(mask) else 0
    resampled_pitch_values.append(avg_pitch)

# Find the index of the first and last non-zero pitch value
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))

# Based on the found index, remove the starting and ending zero 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  # Here the time is relative to the first non-zero value
    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)

# Create a data list containing trimmed time and pitch values
data_list = [{"time": time, "pitch": pitch} for time, pitch in zip(formatted_time_values, trimmed_pitch_values)]


# Write a list of data to a JSON file
output_data = {"data": data_list, "text": file_name_without_extension}
with open("data/output2.json", "w") as json_file:
    json.dump(output_data, json_file, ensure_ascii=False)


# from flask import Flask, send_from_directory
# from flask_cors import CORS

# app = Flask(__name__)
# CORS(app)

# @app.route('/data/output.json')
# def serve_json():
#     return send_from_directory('data', 'output.json') 

# app.run(port=8000)