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
import numpy as np

matplotlib.use("Agg")
import matplotlib.pyplot as plt

from util import (
    angle,
    angle_width,
    positive_int,
    load_dataset,
    existing_path,
    create_video_from_images,
    calculate_average_frame_rate,
    get_colormap_with_special_missing_color,
)


def fill_sparse_data(data: DataFrame, horizontal_resolution: int) -> DataFrame:
    complete_original_ids = DataFrame(
        {
            "original_id": np.arange(
                0,
                (data["ring"].max() + 1) * horizontal_resolution,
                dtype=np.uint32,
            )
        }
    )
    data = complete_original_ids.merge(data, on="original_id", how="left")
    data["ring"] = data["original_id"] // horizontal_resolution
    data["horizontal_position"] = data["original_id"] % horizontal_resolution
    return data


def crop_lidar_data_to_roi(
    data: DataFrame,
    roi_angle_start: float,
    roi_angle_width: float,
    horizontal_resolution: int,
) -> tuple[DataFrame, int]:
    if roi_angle_width == 360:
        return data, horizontal_resolution

    roi_index_start = int(horizontal_resolution / 360 * roi_angle_start)
    roi_index_width = int(horizontal_resolution / 360 * roi_angle_width)
    roi_index_end = roi_index_start + roi_index_width

    if roi_index_end < horizontal_resolution:
        cropped_data = data.iloc[:, roi_index_start:roi_index_end]
    else:
        roi_index_end = roi_index_end - horizontal_resolution
        cropped_data = data.iloc[:, roi_index_end:roi_index_start]

    return cropped_data, roi_index_width


def create_projection_data(
    dataset: Dataset,
    horizontal_resolution: int,
    roi_angle_start: float,
    roi_angle_width: float,
) -> list[Path]:
    converted_lidar_frames = []

    for i, pc in track(
        enumerate(dataset, 1), description="Rendering images...", total=len(dataset)
    ):
        lidar_data = fill_sparse_data(pc.data, horizontal_resolution)
        lidar_data["normalized_range"] = 1 / np.sqrt(
            lidar_data["x"] ** 2 + lidar_data["y"] ** 2 + lidar_data["z"] ** 2
        )
        lidar_data = lidar_data.pivot(
            index="ring", columns="horizontal_position", values="normalized_range"
        )
        lidar_data, _ = crop_lidar_data_to_roi(
            lidar_data, roi_angle_start, roi_angle_width, horizontal_resolution
        )
        converted_lidar_frames.append(lidar_data.to_numpy())

    return np.stack(converted_lidar_frames, axis=0)


def create_2d_projection(
    df: DataFrame,
    output_file_path: Path,
    tmp_file_path: Path,
    colormap_name: str,
    missing_data_color: str,
    reverse_colormap: bool,
    horizontal_resolution: int,
    vertical_resolution: int,
):
    fig, ax = plt.subplots(
        figsize=(float(horizontal_resolution) / 100, float(vertical_resolution) / 100)
    )
    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(
        (horizontal_resolution, vertical_resolution), Image.LANCZOS
    )
    img_resized.save(output_file_path)
    tmp_file_path.unlink()


def render_2d_images(
    dataset: Dataset,
    output_path: Path,
    colormap_name: str,
    missing_data_color: str,
    reverse_colormap: bool,
    horizontal_resolution: int,
    vertical_scale: int,
    horizontal_scale: int,
    roi_angle_start: float,
    roi_angle_width: float,
) -> list[Path]:
    rendered_images = []

    for i, pc in track(
        enumerate(dataset, 1), description="Rendering images...", total=len(dataset)
    ):
        image_data = fill_sparse_data(pc.data, horizontal_resolution).pivot(
            index="ring", columns="horizontal_position", values="range"
        )

        image_data, output_horizontal_resolution = crop_lidar_data_to_roi(
            image_data, roi_angle_start, roi_angle_width, horizontal_resolution
        )

        normalized_data = (image_data - image_data.min().min()) / (
            image_data.max().max() - image_data.min().min()
        )
        image_path = create_2d_projection(
            normalized_data,
            output_path / f"frame_{i:04d}.png",
            output_path / "tmp.png",
            colormap_name,
            missing_data_color,
            reverse_colormap,
            horizontal_resolution=output_horizontal_resolution * horizontal_scale,
            vertical_resolution=(pc.data["ring"].max() + 1) * vertical_scale,
        )

        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-experiment-path",
        required=True,
        type=existing_path,
        help="path to experiment. (directly to bag file, to parent folder for mcap)",
    )
    parser.add_argument(
        "--pointcloud-topic",
        default="/ouster/points",
        type=str,
        help="topic in the ros/mcap bag file containing the point cloud data",
    )
    parser.add_argument(
        "--output-path",
        default=Path("./output"),
        type=Path,
        help="path rendered frames should 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-video",
        type=bool,
        default=True,
        help="if rendered frames should be outputted as a video",
    )
    parser.add_argument(
        "--output-pickle",
        default=True,
        type=bool,
        help="if the processed data should be saved as a pickle file",
    )
    parser.add_argument(
        "--skip-existing",
        default=True,
        type=bool,
        help="if true will skip rendering existing files",
    )
    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",
    )
    parser.add_argument(
        "--horizontal-resolution",
        default=2048,
        type=positive_int,
        help="number of horizontal lidar data points",
    )
    parser.add_argument(
        "--vertical-scale",
        default=1,
        type=positive_int,
        help="multiplier for vertical scale, for better visualization",
    )
    parser.add_argument(
        "--horizontal-scale",
        default=1,
        type=positive_int,
        help="multiplier for horizontal scale, for better visualization",
    )
    parser.add_argument(
        "--roi-angle-start",
        default=0,
        type=angle,
        help="angle where roi starts",
    )
    parser.add_argument(
        "--roi-angle-width",
        default=360,
        type=angle_width,
        help="width of roi in degrees",
    )

    args = parser.parse_args()

    output_path = args.output_path / args.input_experiment_path.stem
    output_path.mkdir(parents=True, exist_ok=True)

    # Create temporary folder for images, if outputting images we use the output folder itself as temp folder
    tmp_path = output_path / "frames" if args.output_images else output_path / "tmp"
    tmp_path.mkdir(parents=True, exist_ok=True)

    dataset = load_dataset(args.input_experiment_path, args.pointcloud_topic)

    images = render_2d_images(
        dataset,
        tmp_path,
        args.colormap_name,
        args.missing_data_color,
        args.reverse_colormap,
        args.horizontal_resolution,
        args.vertical_scale,
        args.horizontal_scale,
        args.roi_angle_start,
        args.roi_angle_width,
    )

    if args.output_pickle:
        output_pickle_path = (
            output_path / args.input_experiment_path.stem
        ).with_suffix(".pkl")
        processed_range_data = create_projection_data(
            dataset,
            args.horizontal_resolution,
            args.roi_angle_start,
            args.roi_angle_width,
        )
        processed_range_data.dump(output_pickle_path)

    if args.output_video:
        input_images_pattern = f"{tmp_path}/frame_%04d.png"
        create_video_from_images(
            input_images_pattern,
            (output_path / args.input_experiment_path.stem).with_suffix(".mp4"),
            calculate_average_frame_rate(dataset),
        )

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

    return 0


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