from __future__ import annotations

from typing import Tuple

import numpy as np
import numpy.typing as npt
from PIL import Image
from matplotlib import cm
from nuplan.database.utils.geometry import view_points


def transform_points(points, transf_matrix: npt.NDArray[np.float64]):
    """
    Applies a homogeneous transform.
    :param transf_matrix: <np.float: 4, 4>. Homogeneous transformation matrix.
    """
    transf_matrix = transf_matrix.astype(np.float32)
    points[:3, :] = transf_matrix[:3, :3] @ points[:3] + transf_matrix[:3, 3].reshape((-1, 1))
    return points


def render_image(
        points, name,
        canvas_size: Tuple[int, int] = (1001, 1001),
        view: npt.NDArray[np.float64] = np.array([[10, 0, 0, 500], [0, 10, 0, 500], [0, 0, 10, 0]]),
        color_dim: int = 2,
):
    """
    Renders pointcloud to an array with 3 channels appropriate for viewing as an image. The image is color coded
    according the color_dim dimension of points (typically the height).
    :param canvas_size: (width, height). Size of the canvas on which to render the image.
    :param view: <np.float: n, n>. Defines an arbitrary projection (n <= 4).
    :param color_dim: The dimension of the points to be visualized as color. Default is 2 for height.
    :return: A Image instance.
    """
    # Apply desired transformation to the point cloud. (height is here considered independent of the view).
    heights = points[2, :]
    points = view_points(points[:3, :], view, normalize=False)
    points[2, :] = heights

    # Remove points that fall outside the canvas.
    mask = np.ones(points.shape[1], dtype=bool)  # type: ignore
    mask = np.logical_and(mask, points[0, :] < canvas_size[0] - 1)
    mask = np.logical_and(mask, points[0, :] > 0)
    mask = np.logical_and(mask, points[1, :] < canvas_size[1] - 1)
    mask = np.logical_and(mask, points[1, :] > 0)
    points = points[:, mask]

    # Scale color_values to be between 0 and 255.
    color_values = points[color_dim, :]
    color_values = 255.0 * (color_values - np.amin(color_values)) / (np.amax(color_values) - np.amin(color_values))

    # Rounds to ints and generate colors that will be used in the image.
    points = np.int16(np.round(points[:2, :]))
    color_values = np.int16(np.round(color_values))
    cmap = [cm.jet(i / 255, bytes=True)[:3] for i in range(256)]

    # Populate canvas, use maximum color_value for each bin
    render = np.tile(np.expand_dims(np.zeros(canvas_size, dtype=np.uint8), axis=2), [1, 1, 3])  # type: ignore
    color_value_array: npt.NDArray[np.float64] = -1 * np.ones(canvas_size, dtype=float)  # type: ignore
    for (col, row), color_value in zip(points.T, color_values.T):
        if color_value > color_value_array[row, col]:
            color_value_array[row, col] = color_value
            render[row, col] = cmap[color_value]

    Image.fromarray(render).save(f'/mnt/f/e2e/navsim_ours/debug/{name}.png')