import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.widgets import Slider, Button, RadioButtons
import librosa
import subprocess
import sys
import os
from pathlib import Path

# Version 35 - Created by PolarPatch

def show_control_panel(audio_file):
    """
    Shows control panel with sliders to test settings.
    """
    print("🎛️  Opening control panel...")
    
    # Load audio for preview
    y, sr = librosa.load(audio_file, sr=22050)
    duration = len(y) / sr
    
    frame_length = 512
    hop_length = 256
    rms = librosa.feature.rms(y=y, frame_length=frame_length, hop_length=hop_length)[0]
    rms_norm = rms / np.max(rms)
    
    # Setup figure with control panel
    fig, ax = plt.subplots(figsize=(14, 8), facecolor='white')
    plt.subplots_adjust(bottom=0.45, right=0.85)
    
    base_y_range = 4
    base_x_range = base_y_range * (16/9)
    
    ax.set_xlim(0, base_x_range)
    ax.set_ylim(-base_y_range/2, base_y_range/2)
    ax.set_aspect('equal')
    ax.axis('off')
    ax.set_title('Preview - Adjust settings and click "Generate Video"', fontsize=12)
    fig.patch.set_facecolor('white')
    
    # Background colors
    bg_colors = {
        'White': '#FFFFFF',
        'Black': '#000000',
        'Green Screen': '#00FF00'
    }
    
    # Line colors (contrast with background)
    line_colors = {
        'White': 'black',
        'Black': 'white',
        'Green Screen': 'black'
    }
    
    # Default settings
    settings = {
        'trail_length': 55,
        'smoothing': 5,
        'intensity': 4.0,
        'power': 0.4,
        'view_width': 1.0,
        'view_height': 1.0,
        'bg_color': 'White'
    }
    
    # Slider axes
    ax_trail = plt.axes([0.15, 0.32, 0.55, 0.03])
    ax_smooth = plt.axes([0.15, 0.27, 0.55, 0.03])
    ax_intensity = plt.axes([0.15, 0.22, 0.55, 0.03])
    ax_power = plt.axes([0.15, 0.17, 0.55, 0.03])
    ax_width = plt.axes([0.15, 0.12, 0.55, 0.03])
    ax_height = plt.axes([0.15, 0.07, 0.55, 0.03])
    ax_button = plt.axes([0.35, 0.01, 0.2, 0.04])
    
    # Radio buttons for background color
    ax_radio = plt.axes([0.86, 0.15, 0.12, 0.15], facecolor='#f0f0f0')
    ax_radio.set_title('Background', fontsize=9)
    radio = RadioButtons(ax_radio, ('White', 'Black', 'Green Screen'), active=0)
    
    # Sliders
    s_trail = Slider(ax_trail, 'Trail length', 20, 160, valinit=55, valstep=5)
    s_smooth = Slider(ax_smooth, 'Smoothing', 1, 20, valinit=5, valstep=1)
    s_intensity = Slider(ax_intensity, 'Intensity', 1.0, 8.0, valinit=4.0, valstep=0.5)
    s_power = Slider(ax_power, 'Low sound boost', 0.2, 1.0, valinit=0.4, valstep=0.05)
    s_width = Slider(ax_width, 'View width', 0.3, 2.0, valinit=1.0, valstep=0.1)
    s_height = Slider(ax_height, 'View height', 0.3, 2.0, valinit=1.0, valstep=0.1)
    
    btn_generate = Button(ax_button, 'Generate Video')
    
    # Preview line
    from matplotlib.collections import LineCollection
    from matplotlib.patches import Circle
    
    lc = LineCollection([], linewidths=2.5, capstyle='round')
    ax.add_collection(lc)
    
    dot_radius = 0.035
    leading_dot = Circle((0, 0), radius=dot_radius, color='black', zorder=10)
    ax.add_patch(leading_dot)
    
    def update_preview(val=None):
        trail_length = int(s_trail.val)
        smoothing = int(s_smooth.val)
        intensity = s_intensity.val
        power = s_power.val
        width_scale = s_width.val
        height_scale = s_height.val
        bg_choice = radio.value_selected
        
        # Update background and line colors
        bg_hex = bg_colors[bg_choice]
        line_color = line_colors[bg_choice]
        
        ax.set_facecolor(bg_hex)
        fig.patch.set_facecolor(bg_hex)
        
        # Update title color for visibility
        title_color = 'white' if bg_choice == 'Black' else 'black'
        ax.set_title('Preview - Adjust settings and click "Generate Video"', fontsize=12, color=title_color)
        
        # Update view based on separate width/height
        scaled_x_range = base_x_range * width_scale
        scaled_y_range = base_y_range * height_scale
        
        ax.set_xlim(0, scaled_x_range)
        ax.set_ylim(-scaled_y_range/2, scaled_y_range/2)
        
        # Generate preview with middle part of audio
        mid_point = len(rms_norm) // 2
        sample_size = min(trail_length, len(rms_norm) // 2)
        sample_rms = rms_norm[mid_point:mid_point + sample_size]
        
        # Apply settings
        sample_processed = np.power(sample_rms, power) * intensity
        
        num_points = len(sample_processed)
        point_spacing = scaled_x_range / trail_length
        
        x_positions = scaled_x_range - (np.arange(num_points)[::-1] * point_spacing)
        y_positions = sample_processed - (scaled_y_range / 2)
        
        # Smoothing
        if num_points > smoothing and smoothing > 1:
            kernel = np.ones(smoothing) / smoothing
            y_smoothed = np.convolve(y_positions, kernel, mode='same')
            edge = smoothing // 2
            y_smoothed[:edge] = y_positions[:edge]
            y_smoothed[-edge:] = y_positions[-edge:]
            y_positions = y_smoothed
        
        # Update plot
        if num_points > 2:
            points = np.array([x_positions, y_positions]).T.reshape(-1, 1, 2)
            segments = np.concatenate([points[:-1], points[1:]], axis=1)
            
            fade_length = max(10, int(len(segments) * 0.25))
            alphas = np.ones(len(segments))
            if len(segments) > fade_length:
                fade_alphas = np.linspace(0.0, 1.0, fade_length)
                alphas[:fade_length] = fade_alphas
            else:
                alphas = np.linspace(0.0, 1.0, len(segments))
            
            # Convert line color to RGB
            if line_color == 'black':
                rgb = (0, 0, 0)
            else:
                rgb = (1, 1, 1)
            
            colors = [(rgb[0], rgb[1], rgb[2], alpha) for alpha in alphas]
            lc.set_segments(segments)
            lc.set_colors(colors)
        
        # Scale dot with average of width/height
        avg_scale = (width_scale + height_scale) / 2
        leading_dot.set_radius(0.035 * avg_scale)
        leading_dot.set_center((x_positions[-1], y_positions[-1]))
        leading_dot.set_color(line_color)
        fig.canvas.draw_idle()
    
    def on_radio_change(label):
        update_preview()
    
    def generate_video(event):
        settings['trail_length'] = int(s_trail.val)
        settings['smoothing'] = int(s_smooth.val)
        settings['intensity'] = s_intensity.val
        settings['power'] = s_power.val
        settings['view_width'] = s_width.val
        settings['view_height'] = s_height.val
        settings['bg_color'] = radio.value_selected
        plt.close(fig)
    
    s_trail.on_changed(update_preview)
    s_smooth.on_changed(update_preview)
    s_intensity.on_changed(update_preview)
    s_power.on_changed(update_preview)
    s_width.on_changed(update_preview)
    s_height.on_changed(update_preview)
    radio.on_clicked(on_radio_change)
    btn_generate.on_clicked(generate_video)
    
    update_preview()
    plt.show()
    
    return settings

def visualize_audio_reactive(audio_file, output_final='output_with_audio.mp4', fps=30, 
                             smoothing=5, trail_length=55, intensity=4.0, power=0.4, 
                             view_width=1.0, view_height=1.0, bg_color='White'):
    print("🎵 Audio-Reactive Visualizer v35 by PolarPatch")
    print("=" * 50)
    
    temp_video = 'temp_visualization_no_audio.mp4'
    
    print("\n📊 Step 1/2: Creating visualization...")
    create_audio_reactive_visualization(audio_file, temp_video, fps, smoothing, 
                                        trail_length, intensity, power, view_width, view_height, bg_color)
    
    print("\n🔊 Step 2/2: Adding audio...")
    add_audio_to_video(temp_video, audio_file, output_final)
    
    if os.path.exists(temp_video):
        os.remove(temp_video)
        print("🧹 Cleaned up")
    
    print(f"\n✅ Done! Video: {output_final}")

def create_audio_reactive_visualization(audio_file, output_video, fps, smoothing, 
                                        trail_length, intensity, power, view_width, view_height, bg_color):
    y, sr = librosa.load(audio_file, sr=22050)
    duration = len(y) / sr
    
    # Background colors
    bg_colors = {
        'White': '#FFFFFF',
        'Black': '#000000',
        'Green Screen': '#00FF00'
    }
    
    # Line colors (contrast with background)
    line_colors = {
        'White': 'black',
        'Black': 'white',
        'Green Screen': 'black'
    }
    
    bg_hex = bg_colors[bg_color]
    line_color = line_colors[bg_color]
    
    print(f"   ⏱️  Duration: {duration:.1f} seconds")
    print(f"   🔧 Trail: {trail_length}, Smoothing: {smoothing}, Intensity: {intensity}")
    print(f"   🔧 Power: {power}, Width: {view_width}, Height: {view_height}")
    print(f"   🎨 Background: {bg_color}")
    
    frame_length = 512
    hop_length = 256
    
    rms = librosa.feature.rms(y=y, frame_length=frame_length, hop_length=hop_length)[0]
    
    rms = rms / np.max(rms)
    rms = np.power(rms, power)
    
    total_frames = int(duration * fps)
    frame_indices = np.linspace(0, len(rms)-1, total_frames)
    amplitudes = np.interp(frame_indices, np.arange(len(rms)), rms)
    
    window = 2
    amplitudes = np.convolve(amplitudes, np.ones(window)/window, mode='same')
    
    fig, ax = plt.subplots(figsize=(16, 9), facecolor=bg_hex)
    
    # Separate width and height scaling
    base_y_range = 4
    base_x_range = base_y_range * (16/9)
    
    x_range = base_x_range * view_width
    y_range = base_y_range * view_height
    
    ax.set_xlim(0, x_range)
    ax.set_ylim(-y_range/2, y_range/2)
    ax.set_aspect('equal')
    ax.axis('off')
    ax.set_facecolor(bg_hex)
    fig.patch.set_facecolor(bg_hex)
    
    amplitude_history = []
    point_spacing = x_range / trail_length
    
    from matplotlib.collections import LineCollection
    from matplotlib.patches import Circle
    lc = LineCollection([], linewidths=2.5, capstyle='round')
    ax.add_collection(lc)
    
    # Scale dot with average
    avg_scale = (view_width + view_height) / 2
    dot_radius = 0.035 * avg_scale
    leading_dot = Circle((0, 0), radius=dot_radius, color=line_color, zorder=10)
    ax.add_patch(leading_dot)
    
    # Convert line color to RGB
    if line_color == 'black':
        rgb = (0, 0, 0)
    else:
        rgb = (1, 1, 1)
    
    def animate(frame):
        amplitude = amplitudes[frame] * intensity
        
        amplitude_history.append(amplitude)
        
        if len(amplitude_history) > trail_length:
            amplitude_history.pop(0)
        
        num_points = len(amplitude_history)
        
        if num_points >= 1:
            x_positions = x_range - (np.arange(num_points)[::-1] * point_spacing)
            y_positions = np.array(amplitude_history) - (y_range / 2)
            
            if num_points > smoothing and smoothing > 1:
                kernel = np.ones(smoothing) / smoothing
                y_smoothed = np.convolve(y_positions, kernel, mode='same')
                edge = smoothing // 2
                y_smoothed[:edge] = y_positions[:edge]
                y_smoothed[-edge:] = y_positions[-edge:]
                y_positions = y_smoothed
            
            leading_dot.set_center((x_positions[-1], y_positions[-1]))
            leading_dot.set_visible(True)
            
            if num_points > 2:
                points = np.array([x_positions, y_positions]).T.reshape(-1, 1, 2)
                segments = np.concatenate([points[:-1], points[1:]], axis=1)
                
                fade_length = max(10, int(len(segments) * 0.25))
                
                alphas = np.ones(len(segments))
                if len(segments) > fade_length:
                    fade_alphas = np.linspace(0.0, 1.0, fade_length)
                    alphas[:fade_length] = fade_alphas
                else:
                    alphas = np.linspace(0.0, 1.0, len(segments))
                
                colors = [(rgb[0], rgb[1], rgb[2], alpha) for alpha in alphas]
                
                lc.set_segments(segments)
                lc.set_colors(colors)
            else:
                lc.set_segments([])
        
        if frame % 90 == 0:
            progress = (frame / total_frames) * 100
            print(f"   ⏳ {progress:.0f}%")
        
        return lc, leading_dot
    
    anim = animation.FuncAnimation(
        fig, animate,
        frames=total_frames,
        interval=1000/fps,
        blit=True
    )
    
    Writer = animation.writers['ffmpeg']
    writer = Writer(fps=fps, bitrate=3000, codec='libx264')
    anim.save(output_video, writer=writer, dpi=100)
    
    plt.close()

def add_audio_to_video(video_file, audio_file, output_file):
    cmd = [
        'ffmpeg',
        '-i', video_file,
        '-i', audio_file,
        '-c:v', 'copy',
        '-c:a', 'aac',
        '-strict', 'experimental',
        '-shortest',
        '-y',
        output_file
    ]
    
    try:
        subprocess.run(cmd, check=True, capture_output=True)
        print("   ✓ Audio added!")
    except subprocess.CalledProcessError as e:
        print(f"   ❌ Error: {e}")

def main():
    if len(sys.argv) < 2:
        print("╔════════════════════════════════════════════════╗")
        print("║  Audio-Reactive Visualizer v35                 ║")
        print("║  Created by PolarPatch                         ║")
        print("╚════════════════════════════════════════════════╝")
        print()
        print("Usage:")
        print("  python audio_visualizer_reactive35.py <audiofile> [output.mp4]")
        print()
        print("Example:")
        print("  python audio_visualizer_reactive35.py song.mp3")
        print()
        print("🎛️  A control panel will open where you can adjust:")
        print("   - Trail length (how much is displayed)")
        print("   - Smoothing (smoothness of the line)")
        print("   - Intensity (amplitude of the graph)")
        print("   - Low sound boost (amplification of quiet sounds)")
        print("   - View width (horizontal zoom)")
        print("   - View height (vertical zoom)")
        print("   - Background color (White, Black, Green Screen)")
        print()
        sys.exit(1)
    
    audio_file = sys.argv[1]
    
    if not Path(audio_file).exists():
        print(f"❌ Error: Cannot find file '{audio_file}'")
        sys.exit(1)
    
    output_video = sys.argv[2] if len(sys.argv) > 2 else Path(audio_file).stem + '_reactive.mp4'
    
    # Show control panel and get settings
    settings = show_control_panel(audio_file)
    
    # Generate video with selected settings
    visualize_audio_reactive(
        audio_file, 
        output_video, 
        fps=30,
        smoothing=settings['smoothing'],
        trail_length=settings['trail_length'],
        intensity=settings['intensity'],
        power=settings['power'],
        view_width=settings['view_width'],
        view_height=settings['view_height'],
        bg_color=settings['bg_color']
    )

if __name__ == "__main__":
    main()