import numpy as np
import os
import time
from PIL import Image
from moviepy.editor import VideoClip, ImageSequenceClip
from moviepy.video.fx.all import fadein, fadeout
from tqdm import tqdm
from multiprocessing import Pool
from functools import partial

# Paths
image_one_path = '/home/daddy/morph/one.jpeg'
image_two_path = '/home/daddy/morph/two.jpeg'
output_path = '/home/daddy/morph/output'
output_filename = f'tornado_morph{int(time.time())}.mp4'


# Function to create each frame of the video
def make_frame(t, img1_array, img2_array, particle_positions, explosion_offsets,
               duration1, duration2, duration3, width, height, particle_size, total_duration):
    if t < duration1:
        # Display the first image
        frame_array = img1_array.copy()
    elif t < duration1 + duration2:
        # Morphing effect with particle explosion and movement to tornado
        progress = (t - duration1) / duration2  # Normalize progress to [0, 1]

        # Create an empty frame with transparency
        frame_array = img1_array.copy()

        for idx, (x, y) in enumerate(particle_positions):
            # Explosion offsets
            offset_x, offset_y = explosion_offsets[idx]

            if progress < 0.3:
                # Explosion phase
                explosion_progress = progress / 0.3  # Scale progress to [0, 1]
                current_x = x + explosion_progress * offset_x
                current_y = y + explosion_progress * offset_y
            else:
                # Movement towards tornado
                tornado_progress = (progress - 0.3) / 0.7  # Scale progress to [0, 1]

                # Calculate tornado position at current time
                tornado_x = width * tornado_progress
                tornado_y = height / 2 + (height / 4) * np.sin(2 * np.pi * tornado_progress * 3)  # Sine wave path

                # Current position moves towards tornado position
                current_x = x + tornado_progress * (tornado_x - x)
                current_y = y + tornado_progress * (tornado_y - y)

            # Convert positions to integers and ensure they stay within frame boundaries
            current_x = int(np.clip(current_x, 0, width - particle_size))
            current_y = int(np.clip(current_y, 0, height - particle_size))

            # Extract patch from the first image
            patch = img1_array[y:y + particle_size, x:x + particle_size].copy()

            # Place the current patch onto the frame
            frame_array[current_y:current_y + particle_size, current_x:current_x + particle_size] = patch

    else:
        # Reveal second image as tornado moves across the screen
        # Create a mask based on tornado position
        progress = (t - duration1 - duration2) / duration3  # Normalize progress to [0, 1]

        # Calculate tornado position at current time
        tornado_x = width * progress
        tornado_radius = width / 10  # Adjust as needed

        # Create a circular mask centered on tornado position
        y_indices, x_indices = np.ogrid[:height, :width]
        mask = ((x_indices - tornado_x) ** 2 + (y_indices - height / 2) ** 2) <= (tornado_radius ** 2)

        # Initialize frame with the first image
        frame_array = img1_array.copy()

        # Replace pixels within the mask with pixels from the second image
        frame_array[mask] = img2_array[mask]

    # Convert the NumPy array back to an image
    frame_image = Image.fromarray(frame_array, 'RGBA')
    # Return frame as an RGB image (MoviePy expects RGB format)
    return np.array(frame_image.convert('RGB'))


# Function for processing a range of frames (used in multiprocessing)
def process_frame_range(frame_numbers, img1_array, img2_array, particle_positions, explosion_offsets,
                        duration1, duration2, duration3, width, height, particle_size, fps, temp_folder, total_duration):
    for frame_number in frame_numbers:
        t = frame_number / fps
        frame = make_frame(t, img1_array, img2_array, particle_positions, explosion_offsets,
                           duration1, duration2, duration3, width, height, particle_size, total_duration)

        # Save frame to disk
        frame_image = Image.fromarray(frame)
        frame_path = os.path.join(temp_folder, f"frame_{frame_number:05d}.jpg")
        frame_image.save(frame_path, format='JPEG', quality=95)  # Adjust quality as needed


def generate_video_from_frames(total_frames, fps, temp_folder, output_video_path):
    frame_paths = [os.path.join(temp_folder, f"frame_{i:05d}.jpg") for i in range(total_frames)]

    # Use MoviePy to create the video from frames
    clip = ImageSequenceClip(frame_paths, fps=fps)

    # Apply fade-in and fade-out effects
    clip = fadein(clip, 1.5).fx(fadeout, 1.5)

    # Write the video file
    clip.write_videofile(output_video_path, codec='libx264', audio=False)


# Main program
if  __name __== "__main__":
    # Input parameters with default values
    try:
        duration1 = int(input("Enter duration for first image display (seconds, default 3): ") or 3)
        duration2 = int(input("Enter duration for explosion and tornado formation (seconds, default 7): ") or 7)
        duration3 = int(input("Enter duration for tornado revealing second image (seconds, default 5): ") or 5)
        particle_size = int(input("Enter particle size (pixels, default 4): ") or 4)
    except ValueError:
        print("Invalid input. Please enter integer values.")
        exit(1)

    # Load images and ensure they are in RGBA format
    img1 = Image.open(image_one_path).convert('RGBA').resize((1024, 1024))
    img2 = Image.open(image_two_path).convert('RGBA').resize((1024, 1024))

    # Convert images to NumPy arrays
    img1_array = np.array(img1)
    img2_array = np.array(img2)

    # Get image dimensions
    width, height = img1.size

    # Generate particle positions
    particle_positions = [(x, y) for y in range(0, height - particle_size + 1, particle_size)
                          for x in range(0, width - particle_size + 1, particle_size)]

    # Prepare explosion offsets
    explosion_offsets = [(np.random.uniform(-width, width), np.random.uniform(-height, height))
                         for _ in particle_positions]

    total_duration = duration1 + duration2 + duration3
    fps = 30  # Increased FPS for smoother animation
    total_frames = int(total_duration * fps)

    # Create a temporary folder for frames
    temp_folder = os.path.join(output_path, "temp_frames")
    os.makedirs(temp_folder, exist_ok=True)

    # Prepare frame numbers
    frame_numbers = list(range(total_frames))

    # Divide frame numbers into chunks for multiprocessing
    num_processes = 2  # Adjust based on your device's capabilities
    chunk_size = len(frame_numbers) // num_processes
    frame_chunks = [frame_numbers[i:i + chunk_size] for i in range(0, len(frame_numbers), chunk_size)]

    # Partial function with fixed arguments
    process_frame_range_partial = partial(
        process_frame_range,
        img1_array=img1_array,
        img2_array=img2_array,
        particle_positions=particle_positions,
        explosion_offsets=explosion_offsets,
        duration1=duration1,
        duration2=duration2,
        duration3=duration3,
        width=width,
        height=height,
        particle_size=particle_size,
        fps=fps,
        temp_folder=temp_folder,
        total_duration=total_duration
    )

    # Process frames in parallel
    with Pool(num_processes) as pool:
        list(tqdm(pool.imap(process_frame_range_partial, frame_chunks), total=len(frame_chunks)))

    # Generate the final video from the processed frames
    generate_video_from_frames(total_frames, fps, temp_folder, os.path.join(output_path, output_filename))

    # Clean up temporary frames
    for frame_file in os.listdir(temp_folder):
        os.remove(os.path.join(temp_folder, frame_file))
    os.rmdir(temp_folder)

    print(f"Video saved to {os.path.join(output_path, output_filename)}")

This didn't come close, I want the particles to whip or like they are all being sucked from the right to the left then a tornado forms about half as tall as the image on the left centered, it's visible the tarp part of the image is the bigger part of the tornado and it gets smaller as it goes towards the bottom of the image, the particles of the one.jpeg are what the tornadoe is created from (if it don't take all the partlics to form the tornado those needed pieces disappear, behind the tornado the image starts to form and it's left visible but only after the tornado goes over that part of the image, eventually going from left to right showing image two.png. And that is when the timer goes for how long it shows that image then fade out 1.5 s. Any questions?