import wave
import gradio as gr


def get_audio_duration(audio_file):
    with wave.open(audio_file, 'rb') as audio:
        frames = audio.getnframes()
        rate = audio.getframerate()
        duration = frames / float(rate)
    return duration

def get_training_info(audio_file):
    if audio_file is None:
        return 'the audio file pls'
    
    duration = get_audio_duration(audio_file)
    sample_rate = wave.open(audio_file, 'rb').getframerate()

    training_info = {
        (0, 2): (150, 'TITAN'),
        (2, 3): (200, 'TITAN'),
        (3, 5): (250, 'TITAN'),
        (5, 10): (300, 'normal'),
        (10, 25): (500, 'Normal'),
        (25, 45): (700, 'Normal'),
        (45, 60): (1000, 'Normal')
    }

    for (min_duration, max_duration), (epochs, pretrain) in training_info.items():
        if min_duration <= duration < max_duration:
            break
    else:
        return 'Duration is not within the specified range.'

    return f'You should use the **{pretrain}** pretrain with **{epochs}** epochs at **{sample_rate/1000}kHz** sample rate. Good luck with your training!'

with gr.Blocks(theme=gr.themes.Base(primary_hue="sky", secondary_hue="blue"), title="RVC TRAINING HELPER") as demo:
    with gr.Tab("Main Settings"):
        audio_p = gr.Audio(type="filepath", label="Your Audio here")
        wtar = gr.Button("Start!")
        audio_q = gr.Textbox(scale=3, label="Your Output here")
        wtar.click(get_training_info, inputs=[audio_p], outputs=[audio_q])
    with gr.Tab("Credits"):
        gr.Markdown("### This code Originaly by [TheStinger](https://huggingface.co./TheStinger)")

        
    demo.launch(debug=True)