from typing import Any, List

import numpy as np
import numpy.typing as npt
import matplotlib
import matplotlib.pyplot as plt

from navsim.visualization.config import LIDAR_CONFIG
from navsim.common.enums import LidarIndex


def filter_lidar_pc(lidar_pc: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
    """
    Filter lidar point cloud according to global configuration
    :param lidar_pc: numpy array of shape (6,n)
    :return: filtered point cloud
    """

    pc = lidar_pc.T
    mask = (
        np.ones((len(pc)), dtype=bool)
        & (pc[:, LidarIndex.X] > LIDAR_CONFIG["x_lim"][0])
        & (pc[:, LidarIndex.X] < LIDAR_CONFIG["x_lim"][1])
        & (pc[:, LidarIndex.Y] > LIDAR_CONFIG["y_lim"][0])
        & (pc[:, LidarIndex.Y] < LIDAR_CONFIG["y_lim"][1])
        & (pc[:, LidarIndex.Z] > LIDAR_CONFIG["z_lim"][0])
        & (pc[:, LidarIndex.Z] < LIDAR_CONFIG["z_lim"][1])
    )
    pc = pc[mask]
    return pc.T


def get_lidar_pc_color(lidar_pc: npt.NDArray[np.float32], as_hex: bool = False) -> List[Any]:
    """
    Compute color map of lidar point cloud according to global configuration
    :param lidar_pc: numpy array of shape (6,n)
    :param as_hex: whether to return hex values, defaults to False
    :return: list of RGB or hex values
    """

    pc = lidar_pc.T
    if LIDAR_CONFIG["color_element"] == "none":
        colors_rgb = np.zeros((len(pc), 3), dtype=np.uin8)

    else:
        if LIDAR_CONFIG["color_element"] == "distance":
            color_intensities = np.linalg.norm(pc[:, LidarIndex.POSITION], axis=-1)
        else:
            color_element_map = {
                "x": LidarIndex.X,
                "y": LidarIndex.Y,
                "z": LidarIndex.Z,
                "intensity": LidarIndex.INTENSITY,
                "ring": LidarIndex.RING,
                "id": LidarIndex.ID,
            }
            color_intensities = pc[:, color_element_map[LIDAR_CONFIG["color_element"]]]

        min, max = color_intensities.min(), color_intensities.max()
        norm_intensities = [(value - min) / (max - min) for value in color_intensities]
        colormap = plt.get_cmap("viridis")
        colors_rgb = np.array([colormap(value) for value in norm_intensities])
        colors_rgb = (colors_rgb[:, :3] * 255).astype(np.uint8)

    assert len(colors_rgb) == len(pc)
    if as_hex:
        return [matplotlib.colors.to_hex(tuple(c / 255.0 for c in rgb)) for rgb in colors_rgb]

    return [tuple(value) for value in colors_rgb]