#!/usr/bin/env python3
import argparse
import random
import os

import cv2
import librosa
import numpy as np
from moviepy.editor import *
from scipy.signal import butter, lfilter
from scipy.signal import find_peaks

# pip install opencv-python librosa numpy moviepy scipy


def detect_beats(audio_file_path, highcut=200, order=5, peak_distance=10, peak_height=0.01):
    # Load the audio file
    y, sr = librosa.load(audio_file_path)

    # Apply a high-pass filter to isolate the bass frequencies
    b, a = butter(order, highcut / (0.5 * sr), btype='high')
    y_filtered = lfilter(b, a, y)

    # Calculate the RMS energy of the filtered signal
    rms = librosa.feature.rms(y=y_filtered, frame_length=1024, hop_length=512)[0]

    # Normalize the RMS energy
    rms_normalized = rms / np.max(rms)

    # Detect the peaks in the RMS energy signal
    peaks, _ = find_peaks(rms_normalized, distance=peak_distance, height=peak_height)

    # Convert the peak indices to times
    beat_times = librosa.frames_to_time(peaks, sr=sr, hop_length=512)

    return beat_times


def create_slideshow(image_folder, audio_file, beat_times, max_duration=2, images=None):
    if images is None:
        images = [img for img in os.listdir(image_folder) if img.endswith(".jpg") or img.endswith(".png")]

    clips = []
    target_size = (1280, 720)

    for i, beat_time in enumerate(beat_times[:-1]):
        img_path = os.path.join(image_folder, images[i % len(images)])
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

        # Resize the image while maintaining aspect ratio and fitting within the target size
        height, width, _ = img.shape
        target_width, target_height = target_size

        scale_width = float(target_width) / float(width)
        scale_height = float(target_height) / float(height)
        scale_factor = min(scale_width, scale_height)

        new_width = int(width * scale_factor)
        new_height = int(height * scale_factor)

        img_resized = cv2.resize(img, (new_width, new_height), interpolation=cv2.INTER_AREA)

        # Add padding
        pad_top = max((target_height - new_height) // 2, 0)
        pad_bottom = max(target_height - new_height - pad_top, 0)
        pad_left = max((target_width - new_width) // 2, 0)
        pad_right = max(target_width - new_width - pad_left, 0)

        img_padded = cv2.copyMakeBorder(img_resized, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_CONSTANT, value=[0, 0, 0])

        duration = beat_times[i + 1] - beat_times[i]

        # If the duration between two beats is greater than the max_duration, repeat the image
        while duration > max_duration:
            clip = ImageClip(img_padded, duration=max_duration)
            clips.append(clip)
            duration -= max_duration

        clip = ImageClip(img_padded, duration=duration)
        clips.append(clip)

    slideshow = concatenate_videoclips(clips)
    return slideshow


def main():
    parser = argparse.ArgumentParser(description="Create a slideshow that matches the bass beats or lyrics of a song.")
    parser.add_argument("image_folder", help="Path to the folder containing the images for the slideshow.")
    parser.add_argument("audio_file_path", help="Path to the input audio file.")
    parser.add_argument("output_file_path", help="Path to the output video file.")
    parser.add_argument("--highcut", type=int, default=200, help="Cutoff frequency for the high-pass filter (default: 200 Hz).")
    parser.add_argument("--order", type=int, default=5, help="Order of the Butterworth filter (default: 5).")
    parser.add_argument("--peak-distance", type=int, default=10, help="Minimum number of samples between peaks (default: 10).")
    parser.add_argument("--peak-height", type=float, default=0.01, help="Minimum height of a peak in the RMS energy signal (default: 0.01).")
    parser.add_argument("--more-help", action="store_true", help="Show more help.")
    parser.add_argument("--randomize", "-r", action="store_true", help="Randomize the order of the images in the slideshow.")
    parser.add_argument("--image_order_file", "-f", help="Path to a text file containing the ordered list of image filenames (not file paths!).")
    args = parser.parse_args()

    if args.more_help:
        print("""highcut: The cutoff frequency for the high-pass filter applied to isolate the bass frequencies. The default value is 200 Hz, which means that the filter will keep frequencies below 200 Hz (bass frequencies) and attenuate higher frequencies. You can adjust this value to focus on different frequency ranges of the bass.

order: The order of the Butterworth filter used for the high-pass filtering. A higher order results in a steeper roll-off, which means a more aggressive filtering. The default value is 5, which should work well for most cases. You can increase or decrease this value to change the sharpness of the filter.

peak_distance: The minimum number of samples between peaks in the RMS energy signal. This parameter helps to avoid detecting multiple peaks that are too close to each other. The default value is 10, which means that two peaks must be at least 10 samples apart to be considered separate peaks. You can adjust this value to control the minimum distance between detected beats.

peak_height: The minimum height of a peak in the normalized RMS energy signal. This parameter helps to filter out peaks that are too small and might not correspond to actual bass beats. The default value is 0.01, which means that a peak must have a height of at least 1% of the maximum RMS energy value to be considered a beat. You can adjust this value to control the minimum strength of detected beats.

When fine-tuning these parameters, you might want to start by adjusting highcut and peak_height to focus on the desired bass frequency range and beat strength. Then, you can experiment with the order and peak_distance parameters to further refine the beat detection. Keep in mind that the optimal values for these parameters might vary depending on the specific characteristics of the audio file you are working with.""")
        quit()

    print('Processing beats...')
    beat_times = detect_beats(args.audio_file_path, highcut=args.highcut, order=args.order, peak_distance=args.peak_distance, peak_height=args.peak_height)
    audio_file = AudioFileClip(args.audio_file_path)

    images = [img for img in os.listdir(args.image_folder) if img.endswith(".jpg") or img.endswith(".png")]
    if args.image_order_file:
        with open(args.image_order_file, 'r') as f:
            ordered_images = [line.strip() for line in f.readlines()]
        images = [img for img in ordered_images if img in images]
    elif args.randomize:
        random.shuffle(images)

    if not images:
        print("No valid images found. Please check the image folder or the image_order_file.")
        return

    print('Creating slideshow...')
    slideshow = create_slideshow(args.image_folder, audio_file, beat_times, images=images)
    final_video = slideshow.set_audio(audio_file)

    print('Writing video...')
    final_video.write_videofile(args.output_file_path, fps=24, codec='libx264', audio_codec='aac')


if __name__ == "__main__":
    main()