"""
Some utility functions e.g. for normalizing angles
Functions for detecting red lights are adapted from scenario runners
atomic_criteria.py
"""

import math
import carla
import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
import cv2
from collections import deque
from shapely.geometry import Polygon, Point
import shapely
import itertools
from .open3d_vis_util import draw_scenes
from copy import deepcopy


def normalize_angle(x):
    x = x % (2 * np.pi)  # force in range [0, 2 pi)
    if x > np.pi:  # move to [-pi, pi)
        x -= 2 * np.pi
    return x


def normalize_angle_degree(x):
    x = x % 360.0
    if x > 180.0:
        x -= 360.0
    return x


def rotate_point(point, angle):
    """
    rotate a given point by a given angle
    """
    x_ = (
        math.cos(math.radians(angle)) * point.x
        - math.sin(math.radians(angle)) * point.y
    )
    y_ = (
        math.sin(math.radians(angle)) * point.x
        + math.cos(math.radians(angle)) * point.y
    )
    return carla.Vector3D(x_, y_, point.z)


def get_traffic_light_waypoints(traffic_light, carla_map):
    """
    get area of a given traffic light
    """
    base_transform = traffic_light.get_transform()
    base_loc = traffic_light.get_location()
    base_rot = base_transform.rotation.yaw
    area_loc = base_transform.transform(traffic_light.trigger_volume.location)

    # Discretize the trigger box into points
    area_ext = traffic_light.trigger_volume.extent
    x_values = np.arange(
        -0.9 * area_ext.x, 0.9 * area_ext.x, 1.0
    )  # 0.9 to avoid crossing to adjacent lanes

    area = []
    for x in x_values:
        point = rotate_point(carla.Vector3D(x, 0, area_ext.z), base_rot)
        point_location = area_loc + carla.Location(x=point.x, y=point.y)
        area.append(point_location)

    # Get the waypoints of these points, removing duplicates
    ini_wps = []
    for pt in area:
        wpx = carla_map.get_waypoint(pt)
        # As x_values are arranged in order, only the last one has to be checked
        if (
            not ini_wps
            or ini_wps[-1].road_id != wpx.road_id
            or ini_wps[-1].lane_id != wpx.lane_id
        ):
            ini_wps.append(wpx)

    # Advance them until the intersection
    wps = []
    eu_wps = []
    for wpx in ini_wps:
        distance_to_light = base_loc.distance(wpx.transform.location)
        eu_wps.append(wpx)
        next_distance_to_light = distance_to_light + 1.0
        while not wpx.is_intersection:
            next_wp = wpx.next(0.5)[0]
            next_distance_to_light = base_loc.distance(next_wp.transform.location)
            if (
                next_wp
                and not next_wp.is_intersection
                and next_distance_to_light <= distance_to_light
            ):
                eu_wps.append(next_wp)
                distance_to_light = next_distance_to_light
                wpx = next_wp
            else:
                break

        if not next_distance_to_light <= distance_to_light and len(eu_wps) >= 4:
            wps.append(eu_wps[-4])
        else:
            wps.append(wpx)

    return area_loc, wps


def lidar_to_ego_coordinate(config, lidar):
    """
    Converts the LiDAR points given by the simulator into the ego agents
    coordinate system
    :param config: GlobalConfig, used to read out lidar orientation and location
    :param lidar: the LiDAR point cloud as provided in the input of run_step
    :return: lidar where the points are w.r.t. 0/0/0 of the car and the carla
    coordinate system.
    """
    yaw = np.deg2rad(config.lidar_rot[2])
    rotation_matrix = np.array(
        [
            [np.cos(yaw), -np.sin(yaw), 0.0],
            [np.sin(yaw), np.cos(yaw), 0.0],
            [0.0, 0.0, 1.0],
        ]
    )

    translation = np.array(config.lidar_pos)

    # The double transpose is a trick to compute all the points together.
    ego_lidar = (rotation_matrix @ lidar[1][:, :3].T).T + translation

    return ego_lidar


def algin_lidar(lidar, translation, yaw):
    """
    Translates and rotates a LiDAR into a new coordinate system.
    Rotation is inverse to translation and yaw
    :param lidar: numpy LiDAR point cloud (N,3)
    :param translation: translations in meters
    :param yaw: yaw angle in radians
    :return: numpy LiDAR point cloud in the new coordinate system.
    """

    rotation_matrix = np.array(
        [
            [np.cos(yaw), -np.sin(yaw), 0.0],
            [np.sin(yaw), np.cos(yaw), 0.0],
            [0.0, 0.0, 1.0],
        ]
    )
    ######## bug here!! This Eq. is only adapt to data augmentation, but is not adapt to align lidars from different timestep.
    aligned_lidar = (rotation_matrix.T @ (lidar - translation).T).T

    # aligned_lidar = (rotation_matrix.T @ (lidar).T).T - translation

    return aligned_lidar


def inverse_conversion_2d(point, translation, yaw):
    """
    Performs a forward coordinate conversion on a 2D point
    :param point: Point to be converted
    :param translation: 2D translation vector of the new coordinate system
    :param yaw: yaw in radian of the new coordinate system
    :return: Converted point
    """
    rotation_matrix = np.array(
        [[np.cos(yaw), -np.sin(yaw)], [np.sin(yaw), np.cos(yaw)]]
    )

    converted_point = rotation_matrix.T @ (point - translation)
    return converted_point


def preprocess_compass(compass):
    """
    Checks the compass for Nans and rotates it into the default CARLA coordinate
    system with range [-pi,pi].
    :param compass: compass value provided by the IMU, in radian
    :return: yaw of the car in radian in the CARLA coordinate system.
    """
    if math.isnan(compass):  # simulation bug
        compass = 0.0
    # The minus 90.0 degree is because the compass sensor uses a different
    # coordinate system then CARLA. Check the coordinate_sytems.txt file
    compass = normalize_angle(compass - np.deg2rad(90.0))

    return compass


def get_relative_transform(ego_matrix, vehicle_matrix):
    """
    Returns the position of the vehicle matrix in the ego coordinate system.
    :param ego_matrix: ndarray 4x4 Matrix of the ego vehicle in global
    coordinates
    :param vehicle_matrix: ndarray 4x4 Matrix of another actor in global
    coordinates
    :return: ndarray position of the other vehicle in the ego coordinate system
    """
    relative_pos = vehicle_matrix[:3, 3] - ego_matrix[:3, 3]
    rot = ego_matrix[:3, :3].T
    relative_pos = rot @ relative_pos

    return relative_pos


def extract_yaw_from_matrix(matrix):
    """Extracts the yaw from a CARLA world matrix"""
    yaw = math.atan2(matrix[1, 0], matrix[0, 0])
    yaw = normalize_angle(yaw)
    return yaw


# Taken from https://stackoverflow.com/a/47381058/9173068
def trapez(y, y0, w):
    return np.clip(np.minimum(y + 1 + w / 2 - y0, -y + 1 + w / 2 + y0), 0, 1)


def weighted_line(r0, c0, r1, c1, w, rmin=0, rmax=np.inf):
    # The algorithm below works fine if c1 >= c0 and c1-c0 >= abs(r1-r0).
    # If either of these cases are violated, do some switches.
    if abs(c1 - c0) < abs(r1 - r0):
        # Switch x and y, and switch again when returning.
        xx, yy, val = weighted_line(
            c0, r0, c1, r1, w, rmin=rmin, rmax=rmax
        )  # pylint: disable=locally-disabled, arguments-out-of-order
        return (yy, xx, val)

    # At this point we know that the distance in columns (x) is greater
    # than that in rows (y). Possibly one more switch if c0 > c1.
    if c0 > c1:
        return weighted_line(
            r1, c1, r0, c0, w, rmin=rmin, rmax=rmax
        )  # pylint: disable=locally-disabled, arguments-out-of-order

    # The following is now always < 1 in abs
    if (c1 - c0) != 0.0:
        slope = (r1 - r0) / (c1 - c0)
    else:
        slope = 0.0

    # Adjust weight by the slope
    w *= np.sqrt(1 + np.abs(slope)) / 2

    # We write y as a function of x, because the slope is always <= 1
    # (in absolute value)
    x = np.arange(c0, c1 + 1, dtype=float)
    if (c1 - c0) != 0.0:
        y = x * slope + (c1 * r0 - c0 * r1) / (c1 - c0)
    else:
        y = np.zeros_like(x)

    # Now instead of 2 values for y, we have 2*np.ceil(w/2).
    # All values are 1 except the upmost and bottommost.
    thickness = np.ceil(w / 2)
    yy = np.floor(y).reshape(-1, 1) + np.arange(-thickness - 1, thickness + 2).reshape(
        1, -1
    )
    xx = np.repeat(x, yy.shape[1])
    vals = trapez(yy, y.reshape(-1, 1), w).flatten()

    yy = yy.flatten()

    # Exclude useless parts and those outside of the interval
    # to avoid parts outside of the picture
    mask = np.logical_and.reduce((yy >= rmin, yy < rmax, vals > 0))

    return (yy[mask].astype(int), xx[mask].astype(int), vals[mask])


def draw_line(
    img,
    start_row,
    start_column,
    end_row,
    end_column,
    color=(255, 255, 255),
    thickness=1,
    rmax=256,
):

    if start_row == end_row and start_column == end_column:
        rr, cc, val = start_row, start_column, 1.0
    else:
        rr, cc, val = weighted_line(
            r0=start_row,
            c0=start_column,
            r1=end_row,
            c1=end_column,
            w=thickness,
            rmax=rmax,
        )

    img[rr, cc, 0] = val * color[0] + (1.0 - val) * img[rr, cc, 0]
    img[rr, cc, 1] = val * color[1] + (1.0 - val) * img[rr, cc, 1]
    img[rr, cc, 2] = val * color[2] + (1.0 - val) * img[rr, cc, 2]
    return img


def draw_box(img, box, color=(255, 255, 255), pixel_per_meter=4, thickness=1):
    translation = np.array([[box[0], box[1]]])
    width = box[2]
    height = box[3]
    yaw = box[4]
    rotation_matrix = np.array(
        [[np.cos(yaw), -np.sin(yaw)], [np.sin(yaw), np.cos(yaw)]]
    )
    speed = box[5] * pixel_per_meter
    speed_coords = np.array([[0.0, 0.0], [0.0, speed]])
    corners = np.array(
        [[-width, -height], [width, -height], [width, height], [-width, height]]
    )
    corner_global = (rotation_matrix @ corners.T).T + translation
    speed_coords_global = (rotation_matrix @ speed_coords.T).T + translation
    corner_global = corner_global.astype(np.int64)
    speed_coords_global = speed_coords_global.astype(np.int64)

    # Only the center is guaranteed to be within the image. Need to clip the corner points.
    max_row = img.shape[0]
    max_column = img.shape[1]
    corner_global[:, 0] = np.clip(corner_global[:, 0], a_min=0, a_max=max_row - 1)
    corner_global[:, 1] = np.clip(corner_global[:, 1], a_min=0, a_max=max_column - 1)
    speed_coords_global[:, 0] = np.clip(
        speed_coords_global[:, 0], a_min=0, a_max=max_row - 1
    )
    speed_coords_global[:, 1] = np.clip(
        speed_coords_global[:, 1], a_min=0, a_max=max_column - 1
    )

    img = draw_line(
        img,
        start_row=corner_global[0, 0],
        start_column=corner_global[0, 1],
        end_row=corner_global[1, 0],
        end_column=corner_global[1, 1],
        color=color,
        thickness=thickness,
        rmax=max_row,
    )
    img = draw_line(
        img,
        start_row=corner_global[1, 0],
        start_column=corner_global[1, 1],
        end_row=corner_global[2, 0],
        end_column=corner_global[2, 1],
        color=color,
        thickness=thickness,
        rmax=max_row,
    )
    img = draw_line(
        img,
        start_row=corner_global[2, 0],
        start_column=corner_global[2, 1],
        end_row=corner_global[3, 0],
        end_column=corner_global[3, 1],
        color=color,
        thickness=thickness,
        rmax=max_row,
    )
    img = draw_line(
        img,
        start_row=corner_global[3, 0],
        start_column=corner_global[3, 1],
        end_row=corner_global[0, 0],
        end_column=corner_global[0, 1],
        color=color,
        thickness=thickness,
        rmax=max_row,
    )
    img = draw_line(
        img,
        start_row=speed_coords_global[0, 0],
        start_column=speed_coords_global[0, 1],
        end_row=speed_coords_global[1, 0],
        end_column=speed_coords_global[1, 1],
        color=color,
        thickness=thickness,
        rmax=max_row,
    )

    return img


class PIDController(object):
    """
    PID controller that converts waypoints to steer, brake and throttle commands
    """

    def __init__(self, k_p=1.0, k_i=0.0, k_d=0.0, n=20):
        self.k_p = k_p
        self.k_i = k_i
        self.k_d = k_d

        self.window = deque([0 for _ in range(n)], maxlen=n)

    def step(self, error):
        self.window.append(error)

        if len(self.window) >= 2:
            integral = np.mean(self.window)
            derivative = self.window[-1] - self.window[-2]
        else:
            integral = 0.0
            derivative = 0.0

        return self.k_p * error + self.k_i * integral + self.k_d * derivative


def gaussian_focal_loss(pred, gaussian_target, alpha=2.0, gamma=4.0, reduction="mean"):
    """Adapted from mmdetection
    Args:
        pred (torch.Tensor): The prediction.
        gaussian_target (torch.Tensor): The learning target of the prediction
            in gaussian distribution.
        alpha (float, optional): A balanced form for Focal Loss.
            Defaults to 2.0.
        gamma (float, optional): The gamma for calculating the modulating
            factor. Defaults to 4.0.
    """
    eps = 1e-12
    pos_weights = gaussian_target.eq(1)
    neg_weights = (1 - gaussian_target).pow(gamma)
    pos_loss = -(pred + eps).log() * (1 - pred).pow(alpha) * pos_weights
    neg_loss = -(1 - pred + eps).log() * pred.pow(alpha) * neg_weights
    loss = pos_loss + neg_loss

    if reduction == "mean":
        loss = loss.mean()
    elif reduction == "sum":
        loss = loss.sum()
    # All other reductions will be no reduction.
    return loss


def bb_vehicle_to_image_system(box, pixels_per_meter, min_x, min_y):
    """
    Changed a bounding box from the vehicle x front, y right coordinate system
    to the x back, y right coordinate system of an image, where the center of
    the car is in the center of the image.
    :return:
    """
    # Multiply position and extent by pixels_per_meter to convert the unit from meters to pixels
    box[:4] = box[:4] * pixels_per_meter
    # Pixel coordinates is y front, x right. CARLA is x front, y right.
    # So we need to swap the axes to convert the coordinates.
    box[0], box[1] = box[1], box[0]
    box[2], box[3] = box[3], box[2]
    # Compute pixel location that represents 0/0 in the image
    translation = np.array([-(min_x * pixels_per_meter), -(min_y * pixels_per_meter)])
    # Shift the coordinates so that the ego_vehicle is at the center of the image
    box[:2] = box[:2] + translation
    box[4] = -box[4]
    return box


def bb_image_to_vehicle_system(box, pixels_per_meter, min_x, min_y):
    """
    Changed a bounding box from the vehicle x front, y right coordinate system
    to the x back, y right coordinate system of an image, where the center of
    the car is in the center of the image.
    :return:
    """
    box[4] = -box[4]
    # Compute pixel location that represents 0/0 in the image
    translation = np.array([-(min_x * pixels_per_meter), -(min_y * pixels_per_meter)])
    # Shift the coordinates so that the ego_vehicle is at [0,0]
    box[:2] = box[:2] - translation
    # Pixel coordinates is y front, x right. CARLA is x front, y right.
    # So we need to swap the axes to convert the coordinates.
    box[0], box[1] = box[1], box[0]
    box[2], box[3] = box[3], box[2]
    # Divide position and extent by pixels_per_meter to convert the unit from pixels to meters
    box[:4] = box[:4] / pixels_per_meter
    return box


def non_maximum_suppression(bounding_boxes, iou_treshhold):
    filtered_boxes = []
    bounding_boxes = np.array(
        list(itertools.chain.from_iterable(bounding_boxes)), dtype=object
    )

    if bounding_boxes.size == 0:  # If no bounding boxes are detected can't do NMS
        return filtered_boxes

    confidences_indices = np.argsort(bounding_boxes[:, -1])
    while len(confidences_indices) > 0:
        idx = confidences_indices[-1]
        current_bb = bounding_boxes[idx]
        filtered_boxes.append(current_bb)
        # Remove last element from the list
        confidences_indices = confidences_indices[:-1]

        if len(confidences_indices) == 0:
            break

        for idx2 in deepcopy(confidences_indices):
            if (
                iou_bbs(current_bb, bounding_boxes[idx2]) > iou_treshhold
            ):  # Remove BB from list
                confidences_indices = confidences_indices[confidences_indices != idx2]

    return filtered_boxes


def rect_polygon(x, y, width, height, angle):
    """Return a shapely Polygon describing the rectangle with centre at
    (x, y) and the given width and height, rotated by angle quarter-turns.

    """
    p = Polygon(
        [(-width, -height), (width, -height), (width, height), (-width, height)]
    )
    # Shapely is very inefficient at these operations, worth rewriting
    return shapely.affinity.translate(
        shapely.affinity.rotate(p, angle, use_radians=True), x, y
    )


def iou_bbs(bb1, bb2):
    a = rect_polygon(bb1[0], bb1[1], bb1[2], bb1[3], bb1[4])
    b = rect_polygon(bb2[0], bb2[1], bb2[2], bb2[3], bb2[4])
    intersection_area = a.intersection(b).area
    union_area = a.union(b).area
    iou = intersection_area / union_area
    return iou


def dot_product(vector1, vector2):
    return vector1.x * vector2.x + vector1.y * vector2.y + vector1.z * vector2.z


def cross_product(vector1, vector2):
    return carla.Vector3D(
        x=vector1.y * vector2.z - vector1.z * vector2.y,
        y=vector1.z * vector2.x - vector1.x * vector2.z,
        z=vector1.x * vector2.y - vector1.y * vector2.x,
    )


def get_separating_plane(r_pos, plane, obb1, obb2):
    """Checks if there is a seperating plane
    rPos Vec3
    plane Vec3
    obb1  Bounding Box
    obb2 Bounding Box
    """
    return abs(dot_product(r_pos, plane)) > (
        abs(dot_product((obb1.rotation.get_forward_vector() * obb1.extent.x), plane))
        + abs(dot_product((obb1.rotation.get_right_vector() * obb1.extent.y), plane))
        + abs(dot_product((obb1.rotation.get_up_vector() * obb1.extent.z), plane))
        + abs(dot_product((obb2.rotation.get_forward_vector() * obb2.extent.x), plane))
        + abs(dot_product((obb2.rotation.get_right_vector() * obb2.extent.y), plane))
        + abs(dot_product((obb2.rotation.get_up_vector() * obb2.extent.z), plane))
    )


def check_obb_intersection(obb1, obb2):
    """
    Checks whether two bounding boxes intersect
    Rather complex looking because it is the general algorithm for 3D oriented bounding boxes.
    """
    r_pos = obb2.location - obb1.location
    return not (
        get_separating_plane(r_pos, obb1.rotation.get_forward_vector(), obb1, obb2)
        or get_separating_plane(r_pos, obb1.rotation.get_right_vector(), obb1, obb2)
        or get_separating_plane(r_pos, obb1.rotation.get_up_vector(), obb1, obb2)
        or get_separating_plane(r_pos, obb2.rotation.get_forward_vector(), obb1, obb2)
        or get_separating_plane(r_pos, obb2.rotation.get_right_vector(), obb1, obb2)
        or get_separating_plane(r_pos, obb2.rotation.get_up_vector(), obb1, obb2)
        or get_separating_plane(
            r_pos,
            cross_product(
                obb1.rotation.get_forward_vector(), obb2.rotation.get_forward_vector()
            ),
            obb1,
            obb2,
        )
        or get_separating_plane(
            r_pos,
            cross_product(
                obb1.rotation.get_forward_vector(), obb2.rotation.get_right_vector()
            ),
            obb1,
            obb2,
        )
        or get_separating_plane(
            r_pos,
            cross_product(
                obb1.rotation.get_forward_vector(), obb2.rotation.get_up_vector()
            ),
            obb1,
            obb2,
        )
        or get_separating_plane(
            r_pos,
            cross_product(
                obb1.rotation.get_right_vector(), obb2.rotation.get_forward_vector()
            ),
            obb1,
            obb2,
        )
        or get_separating_plane(
            r_pos,
            cross_product(
                obb1.rotation.get_right_vector(), obb2.rotation.get_right_vector()
            ),
            obb1,
            obb2,
        )
        or get_separating_plane(
            r_pos,
            cross_product(
                obb1.rotation.get_right_vector(), obb2.rotation.get_up_vector()
            ),
            obb1,
            obb2,
        )
        or get_separating_plane(
            r_pos,
            cross_product(
                obb1.rotation.get_up_vector(), obb2.rotation.get_forward_vector()
            ),
            obb1,
            obb2,
        )
        or get_separating_plane(
            r_pos,
            cross_product(
                obb1.rotation.get_up_vector(), obb2.rotation.get_right_vector()
            ),
            obb1,
            obb2,
        )
        or get_separating_plane(
            r_pos,
            cross_product(obb1.rotation.get_up_vector(), obb2.rotation.get_up_vector()),
            obb1,
            obb2,
        )
    )


def command_to_one_hot(command):
    if command < 0:
        command = 4
    command -= 1
    if command not in [0, 1, 2, 3, 4, 5]:
        command = 3
    cmd_one_hot = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    cmd_one_hot[command] = 1.0

    return np.array(cmd_one_hot)


class InfoDummy(object):
    """Info dummy that tries to mimic TIMMs info class"""

    def __init__(self, info):
        super().__init__()
        self.info = info


def calculate_intrinsic_matrix(fov, height, width):
    """Intrinsics and extrinsics for a single camera.
    adapted from MILE
    https://github.com/wayveai/mile/blob/247280758b40ae999a5de14a8423f0d4db2655ac/mile/data/dataset.py#L194
    """

    # Intrinsics
    f = width / (2.0 * np.tan(fov * np.pi / 360.0))
    cx = width / 2.0
    cy = height / 2.0
    intrinsics = np.array([[f, 0.0, cx], [0.0, f, cy], [0.0, 0.0, 1.0]])

    return intrinsics


def normalize_imagenet(x):
    """Normalize input images according to ImageNet standards.
    Args:
        x (tensor): input images
    """
    x = x.clone()
    x[:, 0] = ((x[:, 0] / 255.0) - 0.485) / 0.229
    x[:, 1] = ((x[:, 1] / 255.0) - 0.456) / 0.224
    x[:, 2] = ((x[:, 2] / 255.0) - 0.406) / 0.225
    return x


class CarlaActorDummy(object):
    """
    Actor dummy structure used to simulate a CARLA actor for data augmentation
    """

    world = None
    bounding_box = None
    transform = None
    id = None

    def __init__(
        self, world, bounding_box, transform, id
    ):  # pylint: disable=locally-disabled, redefined-builtin
        self.world = world
        self.bounding_box = bounding_box
        self.transform = transform
        self.id = id

    def get_world(self):
        return self.world

    def get_transform(self):
        return self.transform

    def get_bounding_box(self):
        return self.bounding_box


def convert_depth(data):
    """
    Computes the normalized depth from a CARLA depth map.
    """
    data = data.astype(np.float32)

    normalized = np.dot(data, [65536.0, 256.0, 1.0])
    normalized /= 256 * 256 * 256 - 1
    # in_meters = 1000 * normalized
    # clip to 50 meters
    normalized = np.clip(normalized, a_min=0.0, a_max=0.05)
    normalized = normalized * 20.0  # Rescale map to lie in [0,1]

    return normalized


def create_projection_grid(config):
    """
    Creates a voxel grid around the car with each voxel containing the pixel index indicating the pixel
    it would land on if you project it into the camera of the car with a pinhole camera model.
    Also returns a valid mask indicating which voxels are visible from the camera.
    Because the coordinates are in normalized display coordinates, the image can also be a down-sampled version.
    :return: grid: voxel grid around the car. Each voxel contains the index of the corresponding camera pixel (x, y, 0).
             Coordinates are in normalized display coordinates [-1, 1].
             (-1,-1) is the top left pixel, (1,1) is the bottom right pixel .
             all_valid: The same voxel grid containing a bool that indicates whether the voxel is visible from the
             camera.
    """
    meters_per_pixel = 1.0 / config.pixels_per_meter
    # + half a pixel because we want the center of the voxel.
    depths = torch.arange(config.min_x, config.max_x, meters_per_pixel) + (
        meters_per_pixel * 0.5
    )
    widths = torch.arange(config.min_y, config.max_y, meters_per_pixel) + (
        meters_per_pixel * 0.5
    )
    meters_per_pixel_height = (
        meters_per_pixel * config.bev_grid_height_downsample_factor
    )
    heights = torch.arange(
        config.min_z_projection, config.max_z_projection, meters_per_pixel_height
    ) + (meters_per_pixel_height * 0.5)

    depths, widths, heights = torch.meshgrid(depths, widths, heights, indexing="ij")
    test_cloud = torch.stack(
        (depths, widths, heights), dim=0
    )  # CARLA coordinate system
    _, d, w, h = test_cloud.shape  # channel, depth, width, height
    # If you rotate the camera adjust the rotation matrix here
    assert (
        config.camera_rot_0[0]
        == config.camera_rot_0[1]
        == config.camera_rot_0[2]
        == 0.0
    )
    rotation_matrix = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
    t = torch.tensor(config.camera_pos).unsqueeze(1)
    test_cloud2 = (rotation_matrix.T @ test_cloud.view(3, -1)) - (rotation_matrix.T @ t)

    # Convert CARLA coordiante system x front, y right, z up to pinhole coordinate system: x right, y down, z front
    # test_cloud2 = torch.stack((test_cloud2[1], test_cloud2[2], test_cloud2[0]))
    test_cloud2 = torch.stack((test_cloud2[1], -test_cloud2[2], test_cloud2[0]))

    # Apply intrinsic camera matrix
    intrinsic_matrix = torch.from_numpy(
        calculate_intrinsic_matrix(
            fov=config.camera_fov,
            height=config.camera_height,
            width=config.camera_width,
        )
    ).to(dtype=torch.float32)
    test_cloud2 = intrinsic_matrix @ test_cloud2

    depths = test_cloud2[2:3]
    grid = torch.zeros_like(test_cloud2).to(dtype=torch.float32)
    # Project to image coordinates using pinhole camera model
    # The depth grid is designed, so that the smallest number is +-0.125. Prevent division by 0 if you change it.
    grid[:2] = test_cloud2[:2] / depths

    # Note that the points themselfs are in pinhole camera coordinates, but the index in the grid represents the voxel
    # in the 3D volume
    grid = grid.view(3, d, w, h)

    width_valid = grid[0:1] >= 0.0
    witdh_valid2 = grid[0:1] < config.camera_width
    width_valid = torch.logical_and(width_valid, witdh_valid2)

    height_valid = grid[1:2] >= 0.0
    height_valid2 = grid[1:2] < config.camera_height
    height_valid = torch.logical_and(height_valid, height_valid2)

    depths = depths.view(1, d, w, h)
    depth_valid = depths > 0.0

    all_valid = torch.logical_and(width_valid, height_valid)
    all_valid = torch.logical_and(all_valid, depth_valid)

    # Normalizes pixel values to [-1, 1] normalized display coordinates
    grid[0:1] = (grid[0:1] / (0.5 * config.camera_width - 0.5)) - 1.0
    grid[1:2] = (grid[1:2] / (0.5 * config.camera_height - 0.5)) - 1.0

    grid = torch.reshape(grid, [1, 3, d, w, h, 1])
    grid = torch.transpose(grid, 1, 5).squeeze(1)

    # valid_points = test_cloud.permute(1, 2, 3, 0)[all_valid.squeeze(0)].clone()
    # draw_scenes(valid_points)
    return grid, all_valid.to(dtype=torch.float32)


def calculate_ego2image_proj(config):
    """
    Creates a voxel grid around the car with each voxel containing the pixel index indicating the pixel
    it would land on if you project it into the camera of the car with a pinhole camera model.
    Also returns a valid mask indicating which voxels are visible from the camera.
    Because the coordinates are in normalized display coordinates, the image can also be a down-sampled version.
    :return: grid: voxel grid around the car. Each voxel contains the index of the corresponding camera pixel (x, y, 0).
             Coordinates are in normalized display coordinates [-1, 1].
             (-1,-1) is the top left pixel, (1,1) is the bottom right pixel .
             all_valid: The same voxel grid containing a bool that indicates whether the voxel is visible from the
             camera.
    """
    assert (
        config.camera_rot_0[0]
        == config.camera_rot_0[1]
        == config.camera_rot_0[2]
        == 0.0
    )
    rotation_matrix = torch.tensor(
        [
            [1.0, 0.0, 0.0, 0.0],
            [0.0, 1.0, 0.0, 0.0],
            [0.0, 0.0, 1.0, 0.0],
            [0.0, 0.0, 0.0, 1.0],
        ]
    )
    t = torch.tensor([*config.camera_pos, 1]).unsqueeze(1)
    ego2camera = rotation_matrix.T
    ego2camera[:-1, -1:] = -(rotation_matrix.T @ t)[:3]

    left2right = torch.tensor(
        [
            [0.0, 1.0, 0.0, 0.0],
            [0.0, 0.0, -1.0, 0.0],
            [1.0, 0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0, 1.0],
        ]
    )
    intri_matrix = torch.eye(4)

    intri_matrix[:3, :3] = torch.from_numpy(
        calculate_intrinsic_matrix(
            fov=config.camera_fov,
            height=config.camera_height,
            width=config.camera_width,
        )
    ).to(dtype=torch.float32)
    ego2image = intri_matrix @ left2right @ ego2camera

    return ego2image


class PerspectiveDecoder(nn.Module):
    """
    Decodes a low resolution perspective grid to a full resolution output. E.g. semantic segmentation, depth
    """

    def __init__(
        self,
        in_channels,
        out_channels,
        inter_channel_0,
        inter_channel_1,
        inter_channel_2,
        scale_factor_0,
        scale_factor_1,
    ):
        super().__init__()
        self.scale_factor_0 = scale_factor_0
        self.scale_factor_1 = scale_factor_1

        self.deconv1 = nn.Sequential(
            nn.Conv2d(in_channels, inter_channel_0, 3, 1, 1),
            nn.ReLU(True),
            nn.Conv2d(inter_channel_0, inter_channel_1, 3, 1, 1),
            nn.ReLU(True),
        )
        self.deconv2 = nn.Sequential(
            nn.Conv2d(inter_channel_1, inter_channel_2, 3, 1, 1),
            nn.ReLU(True),
            nn.Conv2d(inter_channel_2, inter_channel_2, 3, 1, 1),
            nn.ReLU(True),
        )
        self.deconv3 = nn.Sequential(
            nn.Conv2d(inter_channel_2, inter_channel_2, 3, 1, 1),
            nn.ReLU(True),
            nn.Conv2d(inter_channel_2, out_channels, 3, 1, 1),
        )

    def forward(self, x):
        x = self.deconv1(x)
        x = F.interpolate(
            x, scale_factor=self.scale_factor_0, mode="bilinear", align_corners=False
        )
        x = self.deconv2(x)
        x = F.interpolate(
            x, scale_factor=self.scale_factor_1, mode="bilinear", align_corners=False
        )
        x = self.deconv3(x)

        return x


def draw_probability_boxes(
    img, speed_prob, target_speeds, color=(128, 128, 128), color_selected=(255, 165, 0)
):
    speed_index = np.argmax(speed_prob)
    colors = [color for _ in range(len(speed_prob))]
    colors[speed_index] = color_selected
    start_x = 0
    start_y = (
        img.shape[0] - 155 - 150
    )  # 1024-155-150  # start_x and start_y specify position of upper left corner of box
    width_bar = 20 * 4
    width_space = 10
    cv2.rectangle(
        img, (start_x, start_y), (1024, start_y + 155), (255, 255, 255), cv2.FILLED
    )

    for idx, s in enumerate(speed_prob):
        start = start_x + idx * (width_space + width_bar)
        cv2.rectangle(
            img,
            (start, start_y + 130),
            (start + width_bar, start_y + 130 - int(s * 100)),
            colors[idx],
            cv2.FILLED,
        )
        cv2.putText(
            img,
            f"{s:.2f}",
            (int(start + 0.33 * width_bar), start_y + 127 - int(s * 100)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.4,
            (0, 0, 0),
            1,
            cv2.LINE_AA,
        )

        # 3.6 is conversion from m/s to km/h
        cv2.putText(
            img,
            f"{int(round(target_speeds[idx] * 3.6)):02d}",
            (int(start + 0.33 * width_bar), start_y + 150),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (0, 0, 0),
            1,
            cv2.LINE_AA,
        )

    cv2.putText(
        img,
        "km/h",
        (start + width_bar + width_space, start_y + 150),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.7,
        (0, 0, 0),
        1,
        cv2.LINE_AA,
    )


def plant_quant_to_box(config, pred_bounding_boxes):
    """Convert a plant auxiliary class to an x,y location of a box"""
    pred_bb_x = F.softmax(pred_bounding_boxes[0][0], dim=1)
    pred_bb_y = F.softmax(pred_bounding_boxes[1][0], dim=1)
    pred_bb_ext_x = F.softmax(pred_bounding_boxes[2][0], dim=1)
    pred_bb_ext_y = F.softmax(pred_bounding_boxes[3][0], dim=1)
    pred_bb_yaw = F.softmax(pred_bounding_boxes[4][0], dim=1)
    pred_bb_speed = F.softmax(pred_bounding_boxes[5][0], dim=1)

    pred_bb_x = torch.argmax(pred_bb_x, dim=1)
    pred_bb_y = torch.argmax(pred_bb_y, dim=1)
    pred_bb_ext_x = torch.argmax(pred_bb_ext_x, dim=1)
    pred_bb_ext_y = torch.argmax(pred_bb_ext_y, dim=1)
    pred_bb_yaw = torch.argmax(pred_bb_yaw, dim=1)
    pred_bb_speed = torch.argmax(pred_bb_speed, dim=1)

    x_step = (config.max_x - config.min_x) / pow(2, config.plant_precision_pos)
    y_step = (config.max_y - config.min_y) / pow(2, config.plant_precision_pos)
    extent_step = (30) / pow(2, config.plant_precision_pos)
    yaw_step = (2 * np.pi) / pow(2, config.plant_precision_angle)
    speed_step = (config.plant_max_speed_pred / 3.6) / pow(
        2, config.plant_precision_speed
    )

    pred_bb_x = pred_bb_x * x_step - config.max_x
    pred_bb_y = pred_bb_y * y_step - config.max_y
    pred_bb_ext_x = pred_bb_ext_x * extent_step
    pred_bb_ext_y = pred_bb_ext_y * extent_step
    pred_bb_yaw = pred_bb_yaw * yaw_step - np.pi
    pred_bb_speed = pred_bb_speed * speed_step
    pred_bb_center = torch.stack(
        (
            pred_bb_x,
            pred_bb_y,
            pred_bb_ext_x,
            pred_bb_ext_y,
            pred_bb_yaw,
            pred_bb_speed,
        ),
        dim=1,
    )

    return pred_bb_center


def circle_line_segment_intersection(
    circle_center, circle_radius, pt1, pt2, full_line=True, tangent_tol=1e-9
):
    """Find the points at which a circle intersects a line-segment.  This can happen at 0, 1, or 2 points.

    :param circle_center: The (x, y) location of the circle center
    :param circle_radius: The radius of the circle
    :param pt1: The (x, y) location of the first point of the segment
    :param pt2: The (x, y) location of the second point of the segment
    :param full_line: True to find intersections along full line - not just in the segment.
                      False will just return intersections within the segment.
    :param tangent_tol: Numerical tolerance at which we decide the intersections are close enough to consider it a
                        tangent
    :return Sequence[Tuple[float, float]]: A list of length 0, 1, or 2, where each element is a point at which the
                                           circle intercepts a line segment.

    Note: We follow: http://mathworld.wolfram.com/Circle-LineIntersection.html
    Credit: https://stackoverflow.com/a/59582674/9173068
    """

    if np.linalg.norm(pt1 - pt2) < 0.000000001:
        print("Problem")

    (p1x, p1y), (p2x, p2y), (cx, cy) = pt1, pt2, circle_center
    (x1, y1), (x2, y2) = (p1x - cx, p1y - cy), (p2x - cx, p2y - cy)
    dx, dy = (x2 - x1), (y2 - y1)
    dr = (dx**2 + dy**2) ** 0.5
    big_d = x1 * y2 - x2 * y1
    discriminant = circle_radius**2 * dr**2 - big_d**2

    if discriminant < 0:  # No intersection between circle and line
        return []
    else:  # There may be 0, 1, or 2 intersections with the segment
        # This makes sure the order along the segment is correct
        intersections = [
            (
                cx
                + (big_d * dy + sign * (-1 if dy < 0 else 1) * dx * discriminant**0.5)
                / dr**2,
                cy + (-big_d * dx + sign * abs(dy) * discriminant**0.5) / dr**2,
            )
            for sign in ((1, -1) if dy < 0 else (-1, 1))
        ]
        if (
            not full_line
        ):  # If only considering the segment, filter out intersections that do not fall within the segment
            fraction_along_segment = [
                (xi - p1x) / dx if abs(dx) > abs(dy) else (yi - p1y) / dy
                for xi, yi in intersections
            ]
            intersections = [
                pt
                for pt, frac in zip(intersections, fraction_along_segment)
                if 0 <= frac <= 1
            ]
        # If line is tangent to circle, return just one point (as both intersections have same location)
        if len(intersections) == 2 and abs(discriminant) <= tangent_tol:
            return [intersections[0]]
        else:
            return intersections


def crop_array(config, images_i):  # images_i must have dimensions (H,W,C) or (H,W)
    """
    Crop rgb images to the desired height and width
    """
    if config.crop_image:
        # crops rgb/depth/semantics from the bottom to cropped_height and symetrically from both sides to cropped_width
        assert config.cropped_height <= images_i.shape[0]
        assert config.cropped_width <= images_i.shape[1]
        side_crop_amount = (images_i.shape[1] - config.cropped_width) // 2
        if len(images_i.shape) > 2:  # for rgb, we have 3 channels
            return images_i[
                0 : config.cropped_height,
                side_crop_amount : images_i.shape[1] - side_crop_amount,
                :,
            ]
        else:  # for depth and semantics, there is no channel dimension
            return images_i[
                0 : config.cropped_height,
                side_crop_amount : images_i.shape[1] - side_crop_amount,
            ]
    else:
        return images_i
