import numpy as np
import cv2
from numba import njit, prange

def prob_to_logodds(p):
    return np.log(p/(1-p))  

def logodds_to_prob(l):
    return 1 - 1/(1 + np.exp(l))

def lidar_noise(ray_ranges):
    sigma = ray_ranges * 0.01 + 0.001 # individual sigma values based on distance

    noise = np.random.normal(0, sigma)

    ray_ranges = np.maximum(0, ray_ranges + noise)

    return ray_ranges

def remove_out_of_range(ranges, angles, coords, lower_bound, upper_bound):
    mask = (ranges > lower_bound) & (ranges < upper_bound)
    ranges = ranges[mask]
    angles = angles[mask]
    coords = coords[mask]
    return ranges, angles, coords

def correct_lidar_data(coords, pose, lidar_coords_from_robot):
    coords = coords.T

    x_lidar = pose[0] + np.cos(pose[2]) * lidar_coords_from_robot[0] - np.sin(pose[2]) * lidar_coords_from_robot[1]
    y_lidar = pose[1] + np.sin(pose[2]) * lidar_coords_from_robot[0] + np.cos(pose[2]) * lidar_coords_from_robot[1]
    pose_lidar = np.array([x_lidar, y_lidar, pose[2]])

    dx = coords[0] - pose_lidar[0]
    dy = coords[1] - pose_lidar[1]
    
    corrected_angles = np.arctan2(dy, dx) - pose_lidar[2]
    corrected_ranges = np.sqrt(dx**2 + dy**2)
    return corrected_ranges, corrected_angles

def normalize_grid_logodds(grid, capped_value):
    grid_normalized = grid + capped_value # adjust to positive values only
    grid_normalized *= 255 / (2 * capped_value) # normalize between 0 and 255
    np.clip(grid_normalized, 0, 255, out=grid_normalized)
    grid_normalized = grid_normalized.astype(np.uint8) # convert to integers
    grid_normalized = 255 - grid_normalized # flip color map
    return grid_normalized

def normalize_grid_probability(grid):
    grid_normalized = logodds_to_prob(grid) # convert to probabilities
    grid_normalized = (grid_normalized * 255).astype(np.uint8) # normalize between 0 and 255
    grid_normalized = 255 - grid_normalized # flip color map
    return grid_normalized

def find_arrow_coords(pose, resolution, x_min, y_min):
    y, x, yaw = pose # x & y are swapped
    yaw = -1*yaw + np.pi/2 # yaw adjustment for map frame
    length = 1

    # adjust pose to map frame
    x0 = int((x - x_min) / resolution)
    y0 = int((y - y_min) / resolution)

    # find arrow tip point in map frame
    x1 = int(x0 + length / resolution * np.cos(yaw))
    y1 = int(y0 + length / resolution * np.sin(yaw))
    return (x0, y0), (x1, y1)

def draw_robot(grid_colored, pose, res, x_min, y_min):
    # find arrow points in map frame
    start_pt, end_pt = find_arrow_coords(pose, res, x_min, y_min)

    # draw direction arrow
    cv2.arrowedLine(
        grid_colored,
        pt1=start_pt,
        pt2=end_pt,
        color=(0, 0, 255),
        thickness=int(np.ceil(0.1 / res)),
        tipLength=0.03 / res
    )
    # draw robot circle
    cv2.circle(
        grid_colored,
        center=start_pt,
        radius=int(0.2 / res),
        color=(0, 0, 255),
        thickness=-1
    )
    return grid_colored
