import numpy as np
import soundfile as sf
import gradio as gr

def binauralize(audio_file, simulate_rotation, rotation_speed):
    """
    Simulate a binaural (stereo) effect by applying a dynamic panning effect 
    to an input audio file. No HRIR files are required.
    
    Parameters:
        audio_file (str): Path to input audio file (mono or stereo).
        simulate_rotation (bool): If True, apply a dynamic rotation (panning) effect.
        rotation_speed (float): Speed of the rotation effect (in Hz).
        
    Returns:
        output_file (str): Path to the output stereo audio file.
        status (str): Status message.
    """
    try:
        # Load input audio file
        audio, sr = sf.read(audio_file)
    except Exception as e:
        return None, f"Error reading input audio file: {e}"
    
    # If the audio is stereo, convert to mono by averaging channels
    if audio.ndim > 1:
        audio = np.mean(audio, axis=1)
    
    # Create a time vector for the audio length
    t = np.arange(len(audio)) / sr
    
    if simulate_rotation:
        # Compute a time-varying angle for a full cycle (2π) at the desired rotation speed.
        angle = 2 * np.pi * rotation_speed * t
        # Constant power panning: left uses cosine, right uses sine.
        left = np.cos(angle) * audio
        right = np.sin(angle) * audio
    else:
        # If rotation is not enabled, duplicate the audio to both channels.
        left = audio
        right = audio
    
    # Combine the channels into a stereo signal.
    binaural_audio = np.stack((left, right), axis=-1)
    
    # Normalize to prevent clipping.
    max_val = np.max(np.abs(binaural_audio))
    if max_val > 0:
        binaural_audio = binaural_audio / max_val
    
    # Save the output to a WAV file.
    output_file = "output_binaural.wav"
    try:
        sf.write(output_file, binaural_audio, sr)
    except Exception as e:
        return None, f"Error writing output audio file: {e}"
    
    return output_file, "Binaural conversion complete!"

# Create an enhanced UI using Gradio Blocks and Tabs.
with gr.Blocks(title="SonicOrbit", css="""
    /* Custom CSS to enhance spacing and font styling */
    .title { font-size: 2.5em; font-weight: bold; text-align: center; margin-bottom: 0.5em; }
    .subtitle { font-size: 1.2em; text-align: center; margin-bottom: 1em; }
    .footer { text-align: center; font-size: 0.9em; margin-top: 2em; color: #555; }
    """) as demo:
    
    gr.Markdown("<div class='title'>SonicOrbit</div>")
    gr.Markdown("<div class='subtitle'>Binaural 360 Audio Converter with Dynamic Rotation</div>")
    
    with gr.Tabs():
        with gr.Tab("Converter"):
            with gr.Row():
                input_audio = gr.Audio(type="filepath", label="Upload Audio (Mono or Stereo)")
            with gr.Row():
                simulate_rotation = gr.Checkbox(label="Simulate Rotation", value=True)
                rotation_speed = gr.Slider(0.01, 1.0, value=0.1, step=0.01, label="Rotation Speed (Hz)")
            convert_button = gr.Button("Convert Audio")
            with gr.Row():
                # Changed type from "file" to "filepath"
                output_audio = gr.Audio(type="filepath", label="Binaural Audio Output")
                status_text = gr.Textbox(label="Status", interactive=False)
            
            convert_button.click(
                fn=binauralize,
                inputs=[input_audio, simulate_rotation, rotation_speed],
                outputs=[output_audio, status_text]
            )
            
        with gr.Tab("Instructions"):
            gr.Markdown("""
            ### How to Use SonicOrbit
            1. **Upload Audio:**  
               Upload a mono or stereo audio file. If you upload a stereo file, it will be converted to mono by averaging the channels.
            2. **Simulate Rotation:**  
               Enable this option to apply a dynamic panning effect that simulates a rotating sound source.
            3. **Rotation Speed:**  
               Adjust the slider to set the speed of the rotation effect (in Hertz). A higher value rotates the audio field faster.
            4. **Convert Audio:**  
               Click the **Convert Audio** button to process your audio file. The output is a binaural (stereo) audio file with the simulated 360° effect.
               
            Enjoy your immersive 3D audio experience!
            """)
    
    gr.Markdown("<div class='footer'>© 2025 SonicOrbit. All rights reserved.</div>")

if __name__ == "__main__":
    demo.launch()