import gradio as gr
import parselmouth
from parselmouth.praat import call
import numpy as np
import pandas as pd

# gradio_client 1.3.0 crashes when a JSON Schema has additionalProperties: false (a boolean)
import gradio_client.utils as _gcu
_orig = _gcu._json_schema_to_python_type
def _safe(schema, defs=None):
    if not isinstance(schema, dict):
        return "any"
    return _orig(schema, defs)
_gcu._json_schema_to_python_type = _safe

def find_rises_and_peaks_gradient(data, threshold=4):
    data['rise_point'] = 0
    data['peak_point'] = 0

    pitch_values = data['pitch'].values
    gradients = np.gradient(pitch_values)

    in_rise = False
    rise_start = 0
    successive_rise_count = 0
    min_successive_rise = 3
    checking_rise = False

    for i in range(1, len(gradients)):
        if np.isnan(pitch_values[i]):
            checking_rise = False
            in_rise = False
            successive_rise_count = 0
            continue

        if not checking_rise:
            checking_rise = True
            continue

        if gradients[i] >= threshold:
            if not in_rise:
                in_rise = True
                rise_start = i-1
            successive_rise_count += 1
        else:
            if in_rise:
                if successive_rise_count >= min_successive_rise:
                    data.at[rise_start, 'rise_point'] = 1
                    data.at[i-1, 'peak_point'] = 1
                in_rise = False
                successive_rise_count = 0

    return data

def get_pitch(audio_file):
    try:
        sound = parselmouth.Sound(audio_file)
        pitch = call(sound, "To Pitch", 0.0, 75, 500)
        pitch_values = pitch.selected_array['frequency'].copy()
        pitch_values[pitch_values==0] = np.nan
        df_pitch = pd.DataFrame(np.column_stack([pitch.xs(), pitch_values]),
                                columns=['time', 'pitch'])
        df_pitch = find_rises_and_peaks_gradient(df_pitch)
        return df_pitch.to_json(orient='records')
    except Exception as e:
        return "Error in pitch extraction: " + str(e)

demo = gr.Interface(fn=get_pitch, inputs=gr.Audio(type="filepath"), outputs="text", delete_cache=(600, 600))
demo.launch(server_name="0.0.0.0")