from pathlib import Path
from sys import exit

import matplotlib.pyplot as plt
import numpy as np
from configargparse import (
    ArgParser,
    ArgumentDefaultsRawHelpFormatter,
    YAMLConfigFileParser,
)
from open3d.io import read_pinhole_camera_parameters, write_image
from open3d.utility import Vector3dVector
from open3d.visualization.rendering import MaterialRecord, OffscreenRenderer
from pointcloudset import Dataset
from rich.progress import track

from util import (
    calculate_average_frame_rate,
    create_video_from_images,
    existing_file,
    load_dataset,
)


def render_3d_images(
    dataset: Dataset,
    camera_config_input_json_path: Path,
    output_images_path: Path,
    image_pattern_prefix,
) -> list[Path]:
    camera_params = read_pinhole_camera_parameters("saved_camera_settings.json")

    width, height = 1920, 1080
    renderer = OffscreenRenderer(width, height)
    renderer.setup_camera(
        intrinsic_matrix=camera_params.intrinsic.intrinsic_matrix,
        extrinsic_matrix=camera_params.extrinsic,
        intrinsic_height_px=camera_params.intrinsic.height,
        intrinsic_width_px=camera_params.intrinsic.width,
    )
    renderer.scene.set_background([1, 1, 1, 1])

    def color_points_by_range(pcd):
        points = np.asarray(pcd.points)
        distances = np.linalg.norm(points, axis=1)
        max_distance = distances.max()
        min_distance = distances.min()
        normalized_distances = (distances - min_distance) / (
            max_distance - min_distance
        )
        colors = plt.get_cmap("jet")(normalized_distances)[:, :3]
        pcd.colors = Vector3dVector(colors)
        return pcd

    rendered_images = []

    for i, pc in track(
        enumerate(dataset, 1), description="Rendering images...", total=len(dataset)
    ):
        o3d_pc = pc.to_instance("open3d")
        o3d_pc = color_points_by_range(o3d_pc)
        renderer.scene.add_geometry("point_cloud", o3d_pc, MaterialRecord())
        image_path = output_images_path / f"{image_pattern_prefix}_{i:04d}.png"
        write_image(image_path.as_posix(), renderer.render_to_image())
        renderer.scene.remove_geometry("point_cloud")
        rendered_images.append(image_path)

    return rendered_images


def main() -> int:
    parser = ArgParser(
        config_file_parser_class=YAMLConfigFileParser,
        default_config_files=["render3d_config.yaml"],
        formatter_class=ArgumentDefaultsRawHelpFormatter,
        description="Render a 3d representation 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(
        "--camera-config-input-json-path",
        default="./saved_camera_settings.json",
        type=existing_file,
        help="path to json file containing camera settings (can be created with the create_camera_settings.py script)",
    )

    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(args.input_bag_path)

    images = render_3d_images(
        dataset,
        args.camera_config_input_json_path,
        args.tmp_files_path,
        args.output_images_prefix,
    )

    if args.output_video:
        input_images_pattern = (
            f"{args.tmp_files_path / args.output_images_prefix}_%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())