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

# Initialize MasteringStyleTransfer
args = type('Args', (), {
    "model_path": "models/mastering_converter.pt",
    "encoder_path": "models/effects_encoder.pt",
    "sample_rate": 44100,
    "path_to_config": "networks/configs.yaml"
})()
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)
    
    return "output_mastered.wav", "ito_output_mastered.wav" if ito_output_audio is not None else None

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")
        submit_button.click(process_audio, inputs=[input_audio, reference_audio, perform_ito], outputs=[output_audio, ito_output_audio])
    
    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")
        submit_button_yt.click(process_youtube, inputs=[input_url, reference_url, perform_ito_yt], outputs=[output_audio_yt, ito_output_audio_yt])

demo.launch()