import gradio as gr
import numpy as np
import sounddevice as sd

def generate_tone(freq, duration=1.0, sample_rate=44100, volume=0.1):
    t = np.linspace(0, duration, int(sample_rate * duration), False)
    return (volume * np.sin(2 * np.pi * freq * t)).astype(np.float32)

def play_tone(freq, volume):
    tone = generate_tone(freq, volume=volume)
    sd.play(tone, samplerate=44100, blocking=True)
    return f"Played tone at {freq} Hz"

def calculate_sequence_frequencies(tinnitus_freq):
    return [
        int(tinnitus_freq * 0.77),
        int(tinnitus_freq * 0.9),
        int(tinnitus_freq * 1.1),
        int(tinnitus_freq * 1.4)
    ]

def play_sequence(tinnitus_freq, volume):
    frequencies = calculate_sequence_frequencies(tinnitus_freq)
    sequence = []
    for freq in frequencies:
        tone = generate_tone(freq, duration=0.5, volume=volume)
        sequence.extend(tone)
        sequence.extend(np.zeros(int(44100 * 0.1)))  # 0.1s pause between tones
    
    sd.play(np.array(sequence), samplerate=44100, blocking=True)
    return f"Played sequence: {frequencies} Hz"

def update_frequencies(tinnitus_freq):
    frequencies = calculate_sequence_frequencies(tinnitus_freq)
    return f"frequencies used in sequence: {', '.join(map(str, frequencies))}"

with gr.Blocks(title="ACRN Tinnitus Protocol") as demo:
    gr.Markdown(
    """
    # ACRN Tinnitus Protocol
    
    This is my attempt at implementing the Acoustic Coordinated Reset Neuromodulation tinnitus treatment protocol using [this paper](https://link-to-paper) as a guide.
    
    - First lower the volume on your device, so it is not too loud to start.
    - Start the tone by pressing the "Play Tone" button.
    - Adjust the frequency slider until it matches your tinnitus tone. You can also type in the frequency if you know it already.
    - Adjust the volume until it is a little bit louder than your tinnitus tone.
    - Switch from "Tone" to "Sequence" mode
    
    Inspired by [this thread on tinnitustalk.com](https://link-to-thread) and [this reddit thread](https://link-to-reddit-thread).
    """
    )

    with gr.Row():
        tone_btn = gr.Button("Tone")
        sequence_btn = gr.Button("Sequence")
    
    with gr.Row():
        freq_slider = gr.Slider(minimum=100, maximum=20000, value=12694, step=1, label="Frequency")
        freq_number = gr.Number(value=12694, label="Frequency")
    
    volume_slider = gr.Slider(minimum=0.01, maximum=1.0, value=0.1, step=0.01, label="Volume")
    
    output = gr.Textbox(label="Output")
    freq_output = gr.Markdown()

    def update_freq(value, slider):
        return gr.Number.update(value=value) if slider else gr.Slider.update(value=value)

    freq_slider.change(update_freq, inputs=[freq_slider, gr.State(True)], outputs=freq_number)
    freq_number.change(update_freq, inputs=[freq_number, gr.State(False)], outputs=freq_slider)
    
    freq_slider.change(update_frequencies, inputs=[freq_slider], outputs=[freq_output])
    freq_number.change(update_frequencies, inputs=[freq_number], outputs=[freq_output])

    tone_btn.click(play_tone, inputs=[freq_slider, volume_slider], outputs=[output])
    sequence_btn.click(play_sequence, inputs=[freq_slider, volume_slider], outputs=[output])

demo.launch()