import gradio as gr
import mido
import matplotlib.pyplot as plt
from pydub import AudioSegment
from pydub.playback import play
import io

def visualize_midi(midi_file):
    # Load MIDI file
    mid = mido.MidiFile(midi_file.name)

    # Prepare to plot
    plt.figure(figsize=(10, 4))
    ticks = []
    notes = []

    # Extract note and tick data
    for i, track in enumerate(mid.tracks):
        time = 0
        for msg in track:
            time += msg.time
            if not msg.is_meta and msg.type == 'note_on':
                ticks.append(time)
                notes.append(msg.note)

    plt.scatter(ticks, notes, c='blue', marker='o')
    plt.xlabel('Ticks')
    plt.ylabel('Note')
    plt.title('MIDI Visualization')
    plt.grid(True)

    # Save plot to a buffer
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    plt.close()

    return buf

def play_midi(midi_file):
    # Convert MIDI to audio using pydub
    audio = AudioSegment.from_file(midi_file.name, format='mid')
    play(audio)
    return "Playing MIDI..."

# Create Gradio blocks interface
with gr.Blocks() as demo:
    midi_input = gr.File(label="Upload MIDI File")
    midi_visual = gr.Image(type="plot", label="MIDI Visualization")
    play_button = gr.Button("Play MIDI")

    midi_input.change(visualize_midi, inputs=midi_input, outputs=midi_visual)
    play_button.click(play_midi, inputs=midi_input, outputs=None)

demo.launch()