import streamlit as st
import os
import time
import matchering as mg
from pedalboard import Pedalboard, HighpassFilter, Compressor, Limiter, Reverb, Gain
from pedalboard.io import AudioFile
import numpy as np
from scipy.signal import butter, lfilter

class AudioProcessor:
    def __init__(self, target_file, reference_file):
        self.target_file = target_file
        self.reference_file = reference_file
        self.mastered_24bit = 'my_song_master_24bit_.wav'
    
    def mix_and_generate(self):
        # Let's keep only warning outputs here, muting everything else
        mg.log(warning_handler=print)

        # Perform the mixing and generating to 24bit
        mg.process(
            target=self.target_file,
            reference=self.reference_file,
            results=[
                mg.pcm16("my_song_master_16bit_.wav"),
                mg.pcm24(self.mastered_24bit),
            ],
            preview_target=mg.pcm16("preview_my_song.flac"),
            preview_result=mg.pcm16("preview_my_song_master.flac"),
        )
        print("Mixing and generating to 24bit completed.")
    
    def stereo_widen(self, audio, width=1.2):
        left_channel = audio[0::2] * width
        right_channel = audio[1::2] * width
        widened_audio = np.empty_like(audio)
        widened_audio[0::2] = left_channel
        widened_audio[1::2] = right_channel
        return widened_audio

    def reduce_piano_volume(self, audio, sample_rate, freq_low=200, freq_high=2000, reduction_db=-18):
        nyquist = 0.5 * sample_rate
        low = freq_low / nyquist
        high = freq_high / nyquist
        b, a = butter(1, [low, high], btype='band')
        filtered_audio = lfilter(b, a, audio)
        gain_reduction = 10 ** (reduction_db / 20)
        reduced_audio = audio - (filtered_audio * gain_reduction)
        return reduced_audio
    
    def determine_parameters(self, audio):
        avg_volume = np.mean(np.abs(audio))
        highpass_cutoff = 100 + (avg_volume * 50)
        compressor_threshold = -20 + (avg_volume * 10)
        reverb_room_size = min(max(avg_volume, 0.1), 0.5)
        return highpass_cutoff, compressor_threshold, reverb_room_size
    
    def master_audio(self):
        # Load the mixed and generated 24bit audio file
        print(f"Loading 24bit audio file: {self.mastered_24bit}")
        with AudioFile(self.mastered_24bit) as f:
            audio = f.read(f.frames)
            sample_rate = f.samplerate

        # Automatically determine parameters
        highpass_cutoff, compressor_threshold, reverb_room_size = self.determine_parameters(audio)
        print(f"Determined Parameters - Highpass Cutoff: {highpass_cutoff}, Compressor Threshold: {compressor_threshold}, Reverb Room Size: {reverb_room_size}")

        # Create a pedalboard with the desired effects
        board = Pedalboard([
            HighpassFilter(cutoff_frequency_hz=highpass_cutoff),
            Compressor(threshold_db=compressor_threshold, ratio=4),
            Limiter(threshold_db=-0.1),
            Reverb(room_size=reverb_room_size, wet_level=0.2),
            Gain(gain_db=3),
        ])

        # Process the audio with the pedalboard
        processed_audio = board(audio, sample_rate)

        # Apply custom stereo widening
        processed_audio = self.stereo_widen(processed_audio)

        # Reduce piano volume
        processed_audio = self.reduce_piano_volume(processed_audio, sample_rate)

        # Save the processed audio to a new file
        output_file = 'final_mastered_audio.wav'
        print(f"Saving final mastered audio to: {output_file}")
        with AudioFile(output_file, 'w', sample_rate, processed_audio.shape[0]) as f:
            f.write(processed_audio)

        print("Audio mastering completed.")
        return output_file

# Streamlit app
st.title('Moonarch Music Analysis: Mixing & Mastering')
st.header("Experience the Future of Sound")
st.write("""
Everything you need to create and release your music, including samples, plugins, unlimited distribution, and the world's best AI mastering engine.
""")

# Display the Start button
if st.button("Start using Moonarch"):
    st.write("Welcome to Moonarch! Let's start creating amazing music.")

# Add upload option
uploaded_file = st.file_uploader("Upload a song", type=["mp3", "wav"])

if uploaded_file is not None:
    st.write("File uploaded successfully.")
    
    # Placeholder for progress bar
    progress_bar = st.progress(0)
    
    # Save uploaded file to a temporary location
    audio_file_path = os.path.join(os.getcwd(), f'{uploaded_file.name}')
    with open(audio_file_path, 'wb') as f:
        f.write(uploaded_file.getbuffer())
    
    # Simulate file processing
    for percent_complete in range(100):
        time.sleep(0.01)
        progress_bar.progress(percent_complete + 1)
    
    st.write("File processing complete.")
    
    # Reference file for mixing
    reference_file = 'audio_example.mp3'

    # Initialize the processor
    processor = AudioProcessor(target_file=audio_file_path, reference_file=reference_file)

    if st.button('Master Audio'):
        with st.spinner('Mixing, generating, and mastering audio...'):
            processor.mix_and_generate()
            final_output = processor.master_audio()
            
            st.success('Audio mastering complete!')
            st.audio(final_output)
            
            with open(final_output, 'rb') as f:
                st.download_button(
                    label="Download Mastered Audio",
                    data=f,
                    file_name='final_mastered_audio.wav',
                    mime='audio/wav'
                )