import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import wave
import contextlib
import cv2
import os
import shutil
from moviepy.editor import AudioFileClip, VideoFileClip

# --- Configuration ---
# Input WAV and output MP4 paths
WAV_PATH = "garbage.wav"
OUTPUT_MP4 = "garbage_VISUAL.mp4"

# Visualization parameters
FPS = 30
FIG_SIZE = (10, 10)
DPI = 120
NUM_LOOPS = 4 # Number of concentric visualizer loops

# Aesthetics
BG_COLOR = '#0B0014'
COLORMAP = cm.magma # Colormap for dynamic colors (e.g., magma, plasma, inferno)
PARTICLE_COLOR = '#FFFFFF'
TITLE_TEXT = 'GARBAGE'
TITLE_COLOR = (1.0, 1.0, 1.0, 0.08)

# --- Setup Directories ---
FRAMES_DIR = "temp_visualization_frames"
if os.path.exists(FRAMES_DIR):
    shutil.rmtree(FRAMES_DIR)
os.makedirs(FRAMES_DIR, exist_ok=True)

# --- Audio Processing ---
try:
    with contextlib.closing(wave.open(WAV_PATH, 'r')) as wf:
        num_channels = wf.getnchannels()
        sampwidth = wf.getsampwidth()
        framerate = wf.getframerate()
        num_frames = wf.getnframes()
        audio_data = wf.readframes(num_frames)
        
        dtype_map = {1: np.int8, 2: np.int16, 4: np.int32}
        if sampwidth not in dtype_map:
            raise ValueError("Unsupported sample width")
        audio_np = np.frombuffer(audio_data, dtype=dtype_map[sampwidth])

    if num_channels == 2:
        audio_np = audio_np.reshape(-1, 2).mean(axis=1)

    print("Audio file loaded successfully.")
    print(f"Duration: {num_frames / float(framerate):.2f}s, Framerate: {framerate}Hz")

except FileNotFoundError:
    print(f"Error: The file '{WAV_PATH}' was not found.")
    exit()
except Exception as e:
    print(f"An error occurred while reading the audio file: {e}")
    exit()

# --- Visualization Frame Generation ---
samples_per_frame = framerate // FPS
num_frames_to_generate = int(num_frames / samples_per_frame)
max_amplitude = np.max(np.abs(audio_np)) if len(audio_np) > 0 else 1.0

print(f"Generating {num_frames_to_generate} frames...")

for i in range(num_frames_to_generate):
    start = i * samples_per_frame
    end = start + samples_per_frame
    frame_audio = audio_np[start:end]

    if frame_audio.size == 0:
        continue

    # --- Create the Radial Visualization ---
    fig = plt.figure(figsize=FIG_SIZE, facecolor=BG_COLOR)
    ax = plt.subplot(111, polar=True, facecolor=BG_COLOR)

    fig.text(0.5, 0.5, TITLE_TEXT, fontsize=120, color=TITLE_COLOR, 
             ha='center', va='center', weight='bold')

    # Calculate RMS for intensity and get the dynamic color for this frame
    rms = np.sqrt(np.mean(np.square(frame_audio.astype(np.float64))))
    normalized_rms = rms / max_amplitude if max_amplitude > 0 else 0
    if not np.isfinite(normalized_rms):
        normalized_rms = 0
    
    # The color is determined by the loudness of the frame
    dynamic_color = COLORMAP(normalized_rms)
    # A slightly brighter color for the main line
    glow_color = COLORMAP(min(normalized_rms * 1.2, 1.0))

    num_samples = len(frame_audio)
    theta = np.linspace(0, 2 * np.pi, num_samples)
    
    # --- Add Particle Effects ---
    num_particles = int(150 * normalized_rms)
    particle_thetas = np.random.uniform(0, 2 * np.pi, num_particles)
    particle_radii = np.random.uniform(0, 1.5, num_particles) * (1 + normalized_rms)
    particle_sizes = (np.random.rand(num_particles) * 30) * (0.5 + normalized_rms)
    ax.scatter(particle_thetas, particle_radii, s=particle_sizes, color=PARTICLE_COLOR, alpha=0.4 * normalized_rms, zorder=2)

    # --- Plot Multiple Concentric Loops ---
    for j in range(NUM_LOOPS):
        loop_base_radius = 0.4 * (j + 1)
        
        # Modulate radius based on audio waveform
        radius_modulation = (np.abs(frame_audio) / max_amplitude) * 0.2 * (1 + normalized_rms) if max_amplitude > 0 else np.zeros_like(frame_audio)
        radii = loop_base_radius + radius_modulation
        
        # Ensure smooth wrap-around
        loop_theta = np.append(theta, theta[0])
        loop_radii = np.append(radii, radii[0])

        # Plot glow effect - increased alpha for more intensity
        ax.plot(loop_theta, loop_radii, color=glow_color, linewidth=5 + j*2, alpha=0.7 * (0.3 + normalized_rms), zorder=3)
        # Plot main line - increased linewidth and always visible
        ax.plot(loop_theta, loop_radii, color=dynamic_color, linewidth=2.5, alpha=0.9, zorder=4)

    # --- Final Touches on the Plot ---
    ax.set_ylim(0, 2.5) # Increased limit for multiple loops
    ax.grid(False)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.spines['polar'].set_visible(False)

    plt.tight_layout()
    frame_path = os.path.join(FRAMES_DIR, f"frame_{i:05d}.png")
    plt.savefig(frame_path, dpi=DPI, facecolor=BG_COLOR)
    plt.close(fig)

    if (i + 1) % 10 == 0:
        print(f"   ... {i+1}/{num_frames_to_generate} frames rendered")


# --- Video Compilation ---
print("\nAll frames generated. Compiling video...")

first_frame_path = os.path.join(FRAMES_DIR, "frame_00000.png")
if not os.path.exists(first_frame_path):
    print("Error: No frames were generated. Cannot create video.")
    exit()

frame_example = cv2.imread(first_frame_path)
h, w, _ = frame_example.shape
temp_video_path = os.path.join(FRAMES_DIR, "visual.avi")

# Use MJPEG codec which is more compatible
fourcc = cv2.VideoWriter_fourcc(*'MJPG')
video = cv2.VideoWriter(temp_video_path, fourcc, FPS, (w, h))

for i in range(num_frames_to_generate):
    frame_path = os.path.join(FRAMES_DIR, f"frame_{i:05d}.png")
    if os.path.exists(frame_path):
        frame = cv2.imread(frame_path)
        video.write(frame)
video.release()

# --- Combine Video and Audio using ffmpeg-python ---
print("Combining video and audio...")

try:
    import ffmpeg
    
    # Use ffmpeg-python to combine video and audio
    video_stream = ffmpeg.input(temp_video_path)
    audio_stream = ffmpeg.input(WAV_PATH)
    
    output = ffmpeg.output(
        video_stream, 
        audio_stream, 
        OUTPUT_MP4,
        vcodec='libx264',
        acodec='aac',
        shortest=None,
        **{'b:v': '5000k'}  # Video bitrate for quality
    )
    
    # Overwrite output file if it exists
    output = ffmpeg.overwrite_output(output)
    
    # Run the ffmpeg command
    ffmpeg.run(output, capture_stdout=True, capture_stderr=True)
    
    print(f"\nSuccessfully created video: {OUTPUT_MP4}")

except ImportError:
    print("\nError: ffmpeg-python library not found.")
    print("Please install it with: pip install ffmpeg-python")
    print("You can find the silent video at:", temp_video_path)
    
except Exception as e:
    print(f"\nAn error occurred during the final video compilation: {e}")
    import traceback
    traceback.print_exc()
    print("You can find the silent video at:", temp_video_path)

finally:
    print("Cleaning up temporary files...")
    if os.path.exists(FRAMES_DIR):
        shutil.rmtree(FRAMES_DIR)