import cv2
import pandas as pd
from PIL import Image
from tqdm import tqdm

from pipeline.clip_wrapper import MODEL_DIM, ClipWrapper
from pipeline.download_videos import DATA_DIR, REPO_ROOT, VIDEO_DIR

FRAME_EXTRACT_RATE_SECONDS = 5  # Extract a frame every 5 seconds
IMAGES_DIR = DATA_DIR / "images"
DATAFRAME_PATH = DATA_DIR / "dataset.parquet"


def process_videos() -> None:
    "Runs clip on video frames, saves results to a parquet file"
    clip_wrapper = ClipWrapper()
    results = []
    for video_path in tqdm(list(VIDEO_DIR.glob("*.mp4")), desc="Processing videos"):
        video_id = video_path.stem
        extracted_images_dir = IMAGES_DIR / video_id
        extracted_images_dir.mkdir(exist_ok=True, parents=True)
        complete_file = extracted_images_dir / "complete"
        if complete_file.exists():
            continue
        for clip_vector, image, timestamp_secs, frame_idx in get_clip_vectors(
            video_path, clip_wrapper
        ):
            image_path = extracted_images_dir / f"{frame_idx}.jpg"
            image.save(image_path)
            results.append(
                [
                    video_id,
                    frame_idx,
                    timestamp_secs,
                    str(image_path.relative_to(REPO_ROOT)),
                    *clip_vector,
                ]
            )
        complete_file.touch()
    df = pd.DataFrame(
        results,
        columns=["video_id", "frame_idx", "timestamp", "image_path"]
        + [f"dim_{i}" for i in range(MODEL_DIM)],
    )
    print(f"Saving data to {DATAFRAME_PATH}")
    df.to_parquet(DATAFRAME_PATH, index=False)


def get_clip_vectors(video_path, clip_wrapper):
    cap = cv2.VideoCapture(str(video_path))
    num_video_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    extract_every_n_frames = FRAME_EXTRACT_RATE_SECONDS * fps
    for frame_idx in tqdm(range(num_video_frames), desc="Running CLIP on video"):
        ret, frame = cap.read()
        if frame_idx % extract_every_n_frames != 0:
            continue
        image = Image.fromarray(frame[..., ::-1])
        clip_vector = clip_wrapper.images2vec([image]).squeeze().numpy()
        timestamp_secs = frame_idx / fps
        yield clip_vector, image, timestamp_secs, frame_idx
    cap.release()


if __name__ == "__main__":
    process_videos()