import gradio as gr
import pretty_midi
import numpy as np
import tempfile
import librosa
import soundfile as sf
import os

class SimpleMP3Humanizer:
    def __init__(self):
        self.style_presets = {
            "pop": [0, 33, 25, 1],      # Drums, Bass, Guitar, Piano
            "electronic": [0, 39, 81, 89],  # Drums, Synth Bass, Lead, Pad
            "rock": [0, 33, 30, 27],    # Drums, Bass, Distortion Guitar, Clean Guitar
            "cinematic": [0, 48, 61, 5] # Drums, Strings, Horn, Piano
        }
    
    def mp3_to_humanized_mp3(self, mp3_path, style="pop", intensity=0.7):
        """Convert MP3 to humanized MP3 in one step"""
        try:
            # Load the MP3
            y, sr = librosa.load(mp3_path, sr=22050, mono=True)
            duration = len(y) / sr
            
            # Create MIDI
            midi = pretty_midi.PrettyMIDI()
            
            # Add instruments
            for program in self.style_presets[style]:
                is_drum = (program == 0)
                instrument = pretty_midi.Instrument(program=program, is_drum=is_drum)
                midi.instruments.append(instrument)
            
            # Create simple music based on audio
            self.create_simple_music(midi, y, sr, duration, style, intensity)
            
            # Add human feel
            self.add_human_touch(midi, intensity)
            
            # Convert to audio
            audio_data = midi.synthesize()
            
            return audio_data, sr
            
        except Exception as e:
            raise Exception(f"Conversion failed: {str(e)}")
    
    def create_simple_music(self, midi, y, sr, duration, style, intensity):
        """Create basic musical structure"""
        # Create beats
        num_beats = max(8, min(32, int(duration * 2)))  # 2 beats per second
        beat_times = np.linspace(0, duration, num_beats)
        
        instruments = midi.instruments
        
        # Add drums
        drums = next((inst for inst in instruments if inst.is_drum), None)
        if drums:
            for i, time in enumerate(beat_times):
                # Kick on strong beats
                if i % 4 == 0:
                    drums.notes.append(self.create_note(36, 90, time, 0.3))
                # Snare on off-beats
                if i % 4 == 2:
                    drums.notes.append(self.create_note(38, 80, time, 0.2))
                # Hi-hats for electronic/pop
                if style in ["electronic", "pop"]:
                    drums.notes.append(self.create_note(42, 70, time, 0.1))
        
        # Add bass
        bass = next((inst for inst in instruments if not inst.is_drum and 32 <= inst.program <= 39), None)
        if bass:
            bass_notes = [36, 38, 41, 43, 45, 48]  # Simple bass line
            for i, time in enumerate(beat_times[::2]):  # Every other beat
                if i < len(bass_notes):
                    bass.notes.append(self.create_note(bass_notes[i], 85, time, 0.8))
        
        # Add melody
        melody_instruments = [inst for inst in instruments if not inst.is_drum and inst.program not in range(32, 40)]
        if melody_instruments:
            melody = melody_instruments[0]
            melody_notes = [60, 62, 64, 65, 67, 69, 71, 72]  # C major scale
            for i, time in enumerate(beat_times[::4]):  # Every 4 beats
                if i < len(melody_notes):
                    melody.notes.append(self.create_note(melody_notes[i], 80, time, 1.0))
    
    def create_note(self, pitch, velocity, start, duration):
        """Helper to create a note"""
        return pretty_midi.Note(
            velocity=velocity,
            pitch=pitch,
            start=start,
            end=start + duration
        )
    
    def add_human_touch(self, midi, intensity):
        """Add humanization to the music"""
        for instrument in midi.instruments:
            for note in instrument.notes:
                # Random timing
                note.start += np.random.normal(0, 0.02 * intensity)
                note.start = max(0, note.start)
                
                # Random velocity
                note.velocity += int(np.random.normal(0, 10 * intensity))
                note.velocity = max(40, min(127, note.velocity))
                
                # Random duration for non-drums
                if not instrument.is_drum:
                    note.end += np.random.normal(0, 0.05 * intensity)
                    note.end = max(note.start + 0.1, note.end)

def convert_mp3(input_mp3, style, intensity):
    """Main conversion function"""
    if input_mp3 is None:
        return None, "Please upload an MP3 file"
    
    converter = SimpleMP3Humanizer()
    
    try:
        # Convert to humanized MP3
        audio_data, sr = converter.mp3_to_humanized_mp3(input_mp3, style, intensity)
        
        # Save as temporary MP3 file
        output_path = tempfile.mktemp(suffix='_humanized.mp3')
        sf.write(output_path, audio_data, sr)
        
        return output_path, "✅ Conversion successful! Your humanized song is ready below."
        
    except Exception as e:
        return None, f"❌ Error: {str(e)}"

# Simple and compatible interface
with gr.Blocks(theme=gr.themes.Soft(), title="MP3 Humanizer") as demo:
    gr.Markdown("""
    # 🎵 MP3 Humanizer
    **Upload AI Music → Get Human Version → Download MP3**
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 1. Upload Your AI Song")
            input_audio = gr.Audio(
                sources=["upload"],
                type="filepath",
                label="Upload MP3 File"
            )
            
            gr.Markdown("### 2. Choose Settings")
            style = gr.Radio(
                ["pop", "electronic", "rock", "cinematic"],
                value="pop",
                label="Music Style"
            )
            
            intensity = gr.Slider(
                0.1, 1.0, value=0.7,
                label="Human Feel Intensity"
            )
            
            convert_btn = gr.Button(
                "✨ Convert to Human Version",
                variant="primary",
                size="lg"
            )
        
        with gr.Column(scale=1):
            gr.Markdown("### 3. Download Result")
            output_audio = gr.Audio(
                label="Your Humanized Song",
                type="filepath",
                interactive=False
            )
            
            status = gr.Textbox(
                label="Status",
                interactive=False
            )
    
    # Simple instructions
    with gr.Accordion("📖 How to Use", open=True):
        gr.Markdown("""
        1. **Upload** your AI-generated MP3 file
        2. **Choose** your preferred music style
        3. **Adjust** the human feel slider
        4. **Click Convert** and wait a few seconds
        5. **Play the preview** to hear your humanized song
        6. **Click the download icon** in the audio player to save your MP3
        
        That's it! You'll get a complete MP3 file with drums, bass, melody, and human-sounding timing.
        """)
    
    # Conversion process
    def process_conversion(input_mp3, style, intensity):
        output_path, message = convert_mp3(input_mp3, style, intensity)
        return output_path, message
    
    convert_btn.click(
        fn=process_conversion,
        inputs=[input_audio, style, intensity],
        outputs=[output_audio, status]
    )

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