from argparse import ArgumentTypeError
from datetime import timedelta
from pathlib import Path
from subprocess import run

from matplotlib import colormaps
from matplotlib.colors import Colormap
from pointcloudset import Dataset


def load_dataset(
    bag_file_path: Path, pointcloud_topic: str = "/ouster/points"
) -> Dataset:
    return Dataset.from_file(bag_file_path, topic=pointcloud_topic)


def save_dataset(dataset: Dataset, output_file_path: Path):
    if not output_file_path.is_dir():
        raise ArgumentTypeError(f"{output_file_path} has to be a valid folder!")
    dataset.to_file(output_file_path)


def calculate_average_frame_rate(dataset: Dataset):
    timestamps = dataset.timestamps
    time_deltas = [
        timestamps[i + 1] - timestamps[i] for i in range(len(timestamps) - 1)
    ]
    average_delta = sum(time_deltas, timedelta()) / len(time_deltas)
    average_frame_rate = 1 / average_delta.total_seconds()
    return average_frame_rate


def existing_file(path_string: str) -> Path:
    path = Path(path_string)
    if not path.exists():
        raise ArgumentTypeError(f"{path} does not exist!")
    if not path.is_file():
        raise ArgumentTypeError(f"{path} is not a valid file!")
    return path


def existing_folder(path_string: str) -> Path:
    path = Path(path_string)
    if not path.exists():
        raise ArgumentTypeError(f"{path} does not exist!")
    if not path.is_dir():
        raise ArgumentTypeError(f"{path} is not a valid folder!")
    return path


def existing_path(path_string: str) -> Path:
    path = Path(path_string)
    if not path.exists():
        raise ArgumentTypeError(f"{path} does not exist!")
    return path


def positive_int(number_str: str) -> int:
    number_val = int(number_str)
    if number_val < 0:
        raise ArgumentTypeError(f"{number_val} is not a positive integer!")
    return number_val


def angle(angle_str: str) -> float:
    angle_val = float(angle_str)
    if angle_val < 0 or angle_val >= 360:
        raise ArgumentTypeError(
            f"{angle_val} is not a valid angle! Needs to be in [0, 360)"
        )
    return angle_val


def angle_width(angle_str: str) -> float:
    angle_val = float(angle_str)
    if angle_val < 0 or angle_val > 360:
        raise ArgumentTypeError(
            f"{angle_val} is not a valid angle width! Needs to be in [0, 360]"
        )
    return angle_val


def get_colormap_with_special_missing_color(
    colormap_name: str, missing_data_color: str = "black", reverse: bool = False
) -> Colormap:
    colormap = (
        colormaps[colormap_name] if not reverse else colormaps[f"{colormap_name}_r"]
    )
    colormap.set_bad(missing_data_color)
    return colormap


def create_video_from_images(
    input_images_pattern: str, output_file: Path, frame_rate: int
) -> None:
    # Construct the ffmpeg command
    command = [
        "ffmpeg",
        "-y",
        "-framerate",
        str(frame_rate),
        "-i",
        input_images_pattern,
        "-c:v",
        "libx264",
        "-profile:v",
        "high",
        "-crf",
        "20",
        "-pix_fmt",
        "yuv420p",
        output_file.as_posix(),
    ]
    run(command, check=True)