import gradio as gr

import numpy as np
import soundfile as sf
import pandas as pd

import glob
import os
from io import BytesIO
import zipfile

from tqdm import tqdm

def synthesize_clicks(annotation_file, bar_duration_ms, beat_duration_ms, bar_freq, beat_freq):
    bar_duration_samples = int(bar_duration_ms * 44.1)
    beat_duration_samples = int(beat_duration_ms * 44.1)

    click_bar = np.sin(2 * np.pi * bar_freq * np.arange(bar_duration_samples) / 44100)
    click_bar = click_bar * np.linspace(1, 0, bar_duration_samples)
    click_beat = np.sin(2 * np.pi * beat_freq * np.arange(beat_duration_samples) / 44100)
    click_beat = click_beat * np.linspace(1, 0, beat_duration_samples)

    annotation = pd.read_csv(annotation_file, header = None, sep = '\t')
    annotation = annotation[annotation[0] >= 0].reset_index(drop = True)
    bars = annotation[annotation[2] == 'bar']
    beats = annotation[annotation[2] == 'beat']
    
    duration_samples = int((annotation[0].iloc[-1] + 1) * 44100)

    audio = np.zeros(duration_samples)

    for i in range(len(bars)):
        start = int(bars[0].iloc[i] * 44100)
        audio[start:start + bar_duration_samples] += click_bar
    
    for i in range(len(beats)):
        start = int(beats[0].iloc[i] * 44100)
        audio[start:start + beat_duration_samples] += click_beat

    return audio

def metronomify(files, _, bar_duration, beat_duration, bar_freq, beat_freq):
    with zipfile.ZipFile('metronome.zip', 'w') as z:
        for file in files:
            audio = synthesize_clicks(file, bar_duration, beat_duration, bar_freq, beat_freq)
            with BytesIO() as f:
                sf.write(f, audio, 44100, format='flac')
                f.seek(0)
                data = f.read()
                z.writestr(os.path.splitext(os.path.basename(file))[0] + '.flac', data)
    return 'metronome.zip'

# audio_file = os.path.join('./audio', os.path.splitext(os.path.basename(annotation_file))[0] + '.flac')
# sf.write(audio_file, audio, 44100)

file_uploads = gr.File(label = 'Annotation files', file_count = 'multiple', file_types = ['.txt','.csv'])
format_info = gr.Markdown('*Expects headerless tab separated files with 4 columns: onset, offset (ignored), beat/bar, ticks (ignored).*')
bar_duration = gr.Number(label = 'Bar duration (ms)', value = 50)
beat_duration = gr.Number(label = 'Beat duration (ms)', value = 50)
bar_freq = gr.Number(label = 'Bar frequency (Hz)', value = 1000)
beat_freq = gr.Number(label = 'Beat frequency (Hz)', value = 500)

zip_out = gr.File(label = 'Metronome zip')

demo = gr.Interface(metronomify, [file_uploads, format_info, bar_duration, beat_duration, bar_freq, beat_freq], zip_out, title = 'Metronomify')
demo.launch()