from configargparse import ArgParser, YAMLConfigFileParser, ArgumentDefaultsRawHelpFormatter
from sys import exit
from pathlib import Path
from pointcloudset import Dataset
from rich.progress import track
from pandas import DataFrame
from PIL import Image
import matplotlib.pyplot as plt

from util import (
    load_dataset_from_bag,
    existing_file,
    create_video_from_images,
    calculate_average_frame_rate,
    get_colormap_with_special_missing_color,
)


def create_2d_projection(
    df: DataFrame,
    output_file_path: Path,
    tmp_file_path: Path,
    colormap_name: str,
    missing_data_color: str,
    reverse_colormap: bool,
):
    fig, ax = plt.subplots(figsize=(20.48, 5.12))
    ax.imshow(
        df,
        cmap=get_colormap_with_special_missing_color(colormap_name, missing_data_color, reverse_colormap),
        aspect="auto",
    )
    ax.axis("off")
    fig.subplots_adjust(left=0, right=1, top=1, bottom=0)
    plt.savefig(tmp_file_path, dpi=100, bbox_inches="tight", pad_inches=0)
    plt.close()
    img = Image.open(tmp_file_path)
    img_resized = img.resize((2048, 512), Image.LANCZOS)
    img_resized.save(output_file_path)


def render_2d_images(
    dataset: Dataset,
    output_images_path: Path,
    image_pattern_prefix: str,
    tmp_files_path: Path,
    colormap_name: str,
    missing_data_color: str,
    reverse_colormap: bool,
) -> list[Path]:
    rendered_images = []

    for i, pc in track(enumerate(dataset, 1), description="Rendering images...", total=len(dataset)):
        pc.data["horizontal_position"] = pc.data["original_id"] % 2048
        image_data = pc.data.pivot(index="ring", columns="horizontal_position", values="range")
        normalized_data = (image_data - image_data.min().min()) / (image_data.max().max() - image_data.min().min())
        image_path = create_2d_projection(
            normalized_data,
            output_images_path / f"{image_pattern_prefix}_frame_{i:04d}.png",
            tmp_files_path / "tmp.png",
            colormap_name,
            missing_data_color,
            reverse_colormap,
        )

        rendered_images.append(image_path)

    return rendered_images


def main() -> int:
    parser = ArgParser(
        config_file_parser_class=YAMLConfigFileParser,
        default_config_files=["render2d_config.yaml"],
        formatter_class=ArgumentDefaultsRawHelpFormatter,
        description="Render a 2d projection of a point cloud",
    )
    parser.add_argument("--render-config-file", is_config_file=True, help="yaml config file path")
    parser.add_argument("--input-bag-path", required=True, type=existing_file, help="path to bag file")
    parser.add_argument(
        "--tmp-files-path", default=Path("./tmp"), type=Path, help="path temporary files will be written to"
    )
    parser.add_argument(
        "--output-images", type=bool, default=True, help="if rendered frames should be outputted as images"
    )
    parser.add_argument(
        "--output-images-path", default=Path("./output"), type=Path, help="path rendered frames should be written to"
    )
    parser.add_argument(
        "--output-video", type=bool, default=True, help="if rendered frames should be outputted as a video"
    )
    parser.add_argument(
        "--output-video-path",
        default=Path("./output/2d_render.mp4"),
        type=Path,
        help="path rendered video should be written to",
    )
    parser.add_argument("--output-images-prefix", default="2d_render", type=str, help="filename prefix for output")
    parser.add_argument("--colormap-name", default="viridis", type=str, help="name of matplotlib colormap to be used")
    parser.add_argument(
        "--missing-data-color", default="black", type=str, help="name of color to be used for missing data"
    )
    parser.add_argument("--reverse-colormap", default=True, type=bool, help="if colormap should be reversed")

    args = parser.parse_args()

    if args.output_images:
        args.output_images_path.mkdir(parents=True, exist_ok=True)
        args.tmp_files_path = args.output_images_path
    else:
        args.tmp_files_path.mkdir(parents=True, exist_ok=True)

    if args.output_video:
        args.output_video_path.parent.mkdir(parents=True, exist_ok=True)

    dataset = load_dataset_from_bag(args.input_bag_path)

    images = render_2d_images(
        dataset,
        args.tmp_files_path,
        args.output_images_prefix,
        args.tmp_files_path,
        args.colormap_name,
        args.missing_data_color,
        args.reverse_colormap,
    )

    if args.output_video:
        input_images_pattern = f"{args.tmp_files_path / args.output_images_prefix}_frame_%04d.png"
        create_video_from_images(input_images_pattern, args.output_video_path, calculate_average_frame_rate(dataset))

    if not args.output_images:
        for image in images:
            image.unlink()

    return 0


if __name__ == "__main__":
    exit(main())