import os
import cv2
import numpy as np
from PIL import Image

width = 200
height = 400
fps = 30  # Frames per second
total_frames = 240  # Increased to actually see the animation move

# Ensure output directory exists
os.makedirs("output", exist_ok=True)

# Define the VideoWriter
# 'mp4v' is a widely supported MP4 codec
fourcc = cv2.VideoWriter_fourcc(*"mp4v")
video_path = "output/rainbow.mp4"
video = cv2.VideoWriter(video_path, fourcc, fps, (width, height))

colors = (
    (0xFF, 0x00, 0x00),
    (0xFF, 0x7F, 0x00),
    (0xFF, 0xFF, 0x00),
    (0x00, 0xFF, 0x00),
    (0x00, 0xFF, 0xFF),
    (0x00, 0x00, 0xFF),
    (0x8B, 0x00, 0xFF),
    (0xFF, 0x00, 0xFF),
)

img = Image.new("RGB", (width, height))
pixels = img.load()

frame_num = 0
while frame_num < total_frames:
    for y in range(height):
        for x in range(width):
            # Adjusted the bit-shift so the animation moves visibly frame-by-frame
            pixels[x, y] = colors[((x + y + (frame_num >> 1)) >> 4) & 7]

    # 1. Convert PIL Image to a NumPy array
    frame_np = np.array(img)

    # 2. Convert RGB (PIL format) to BGR (OpenCV format)
    frame_bgr = cv2.cvtColor(frame_np, cv2.COLOR_RGB2BGR)

    # 3. Write the frame to the video file
    video.write(frame_bgr)

    frame_num += 1

# Crucial: Release the video writer to finalize and save the file
video.release()
print(f"Video saved successfully to {video_path}")