from pytube import YouTube
import numpy as np
from decord import VideoReader, cpu
import imageio


def download_youtube_video(url: str):
    yt = YouTube(url)

    streams = yt.streams.filter(file_extension="mp4")
    file_path = streams[0].download()
    return file_path


def sample_frame_indices(clip_len, frame_sample_rate):
    converted_len = int(clip_len * frame_sample_rate)
    start_idx = 0
    end_idx = converted_len
    indices = np.linspace(start_idx, end_idx, num=clip_len)
    indices = np.clip(indices, start_idx, end_idx - 1).astype(np.int64)
    return indices


def sample_frames_from_video_file(file_path: str, num_frames: int = 16):
    videoreader = VideoReader(file_path, num_threads=1, ctx=cpu(0))

    # sample frames
    videoreader.seek(0)
    indices = sample_frame_indices(clip_len=num_frames, frame_sample_rate=4)
    frames = videoreader.get_batch(indices).asnumpy()
    return frames


def convert_frames_to_gif(frames):
    SAVE_PATH = "frames.gif"
    converted_frames = frames.astype(np.uint8)
    imageio.mimsave(SAVE_PATH, converted_frames, fps=8)
    return SAVE_PATH