import os
import cv2
import numpy as np

def compose_final_video(mask_path, frames_folder, extracted_frames_folder, output_path, fps=24):
    """
    Composes the final heatmap video by placing each heatmap frame
    on top of its corresponding original frame using the table region defined by the mask.
    """
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    if mask is None:
        raise ValueError(f"Mask not found at {mask_path}")

    _, binary_mask = cv2.threshold(mask, 200, 255, cv2.THRESH_BINARY)
    height, width = binary_mask.shape

    contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        raise ValueError("No white area found in mask")

    table_contour = max(contours, key=cv2.contourArea)
    x, y, w, h = cv2.boundingRect(table_contour)

    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    video_writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height))

    heatmap_files = sorted([f for f in os.listdir(frames_folder) if f.endswith(('.bmp', '.jpg', '.png'))])
    original_files = sorted([f for f in os.listdir(extracted_frames_folder) if f.endswith(('.bmp', '.jpg', '.png'))])

    if len(heatmap_files) != len(original_files):
        raise ValueError("Mismatch between number of heatmap frames and original frames")

    for heatmap_file, original_file in zip(heatmap_files, original_files):
        heatmap = cv2.imread(os.path.join(frames_folder, heatmap_file))
        original = cv2.imread(os.path.join(extracted_frames_folder, original_file))

        if heatmap is None or original is None:
            continue

        ph, pw = heatmap.shape[:2]
        if pw > w or ph > h:
            scale = min(w / pw, h / ph)
            heatmap = cv2.resize(heatmap, (int(pw * scale), int(ph * scale)))
            pw, ph = heatmap.shape[1], heatmap.shape[0]

        x_offset = x + (w - pw) // 2
        y_offset = y + (h - ph) // 2

        result = original.copy()
        result[y_offset:y_offset+ph, x_offset:x_offset+pw] = heatmap
        cv2.drawContours(result, [table_contour], -1, (0, 255, 0), 2)

        video_writer.write(result)

    video_writer.release()
    print("[Video] Final video saved to:", output_path)