import gradio as gr
import torch
import soundfile as sf
import numpy as np
import yaml
from inference import MasteringStyleTransfer
from utils import download_youtube_audio
from config import args

mastering_transfer = MasteringStyleTransfer(args)

def process_audio(input_audio, reference_audio, perform_ito):
    # Process the audio files
    output_audio, predicted_params, ito_output_audio, ito_predicted_params, _, sr, _ = mastering_transfer.process_audio(
        input_audio, reference_audio, reference_audio, {}, perform_ito
    )
    
    # Save the output audio
    sf.write("output_mastered.wav", output_audio.T, sr)
    if ito_output_audio is not None:
        sf.write("ito_output_mastered.wav", ito_output_audio.T, sr)
    
    # Generate parameter output strings
    param_output = mastering_transfer.get_param_output_string(predicted_params)
    ito_param_output = mastering_transfer.get_param_output_string(ito_predicted_params) if ito_predicted_params is not None else "ITO not performed"
    
    # Generate top 10 differences if ITO was performed
    top_10_diff = mastering_transfer.get_top_10_diff_string(predicted_params, ito_predicted_params) if ito_predicted_params is not None else "ITO not performed"
    
    return "output_mastered.wav", "ito_output_mastered.wav" if ito_output_audio is not None else None, param_output, ito_param_output, top_10_diff

def process_youtube(input_url, reference_url, perform_ito):
    input_audio = download_youtube_audio(input_url)
    reference_audio = download_youtube_audio(reference_url)
    return process_audio(input_audio, reference_audio, perform_ito)

with gr.Blocks() as demo:
    gr.Markdown("# Mastering Style Transfer Demo")
    with gr.Tab("Upload Audio"):
        input_audio = gr.Audio(label="Input Audio")
        reference_audio = gr.Audio(label="Reference Audio")
        perform_ito = gr.Checkbox(label="Perform ITO")
        submit_button = gr.Button("Process")
        output_audio = gr.Audio(label="Output Audio")
        ito_output_audio = gr.Audio(label="ITO Output Audio")
        param_output = gr.Textbox(label="Predicted Parameters", lines=10)
        ito_param_output = gr.Textbox(label="ITO Predicted Parameters", lines=10)
        top_10_diff = gr.Textbox(label="Top 10 Parameter Differences", lines=10)
        submit_button.click(process_audio, 
                            inputs=[input_audio, reference_audio, perform_ito], 
                            outputs=[output_audio, ito_output_audio, param_output, ito_param_output, top_10_diff])
    
    with gr.Tab("YouTube URLs"):
        input_url = gr.Textbox(label="Input YouTube URL")
        reference_url = gr.Textbox(label="Reference YouTube URL")
        perform_ito_yt = gr.Checkbox(label="Perform ITO")
        submit_button_yt = gr.Button("Process")
        output_audio_yt = gr.Audio(label="Output Audio")
        ito_output_audio_yt = gr.Audio(label="ITO Output Audio")
        param_output_yt = gr.Textbox(label="Predicted Parameters", lines=10)
        ito_param_output_yt = gr.Textbox(label="ITO Predicted Parameters", lines=10)
        top_10_diff_yt = gr.Textbox(label="Top 10 Parameter Differences", lines=10)
        submit_button_yt.click(process_youtube, 
                               inputs=[input_url, reference_url, perform_ito_yt], 
                               outputs=[output_audio_yt, ito_output_audio_yt, param_output_yt, ito_param_output_yt, top_10_diff_yt])

demo.launch()