from openride.core.numba.interpolation import interpolate_line_index_at_xy, linear_interpolation

import numba
import numpy as np

from typing import Tuple, Union

EPSILON = np.finfo(float).eps


@numba.njit(cache=True)
def vertices_to_distances(vertices:np.ndarray, bird_eye_view:bool=True) -> np.ndarray:
    distances_squared = (vertices[1:,0] - vertices[:-1,0]) ** 2 + (vertices[1:,1] - vertices[:-1,1]) ** 2
    if not bird_eye_view:
        distances_squared += (vertices[1:,2] - vertices[:-1,2]) ** 2
    distances = distances_squared ** 0.5
    distances = np.append(0, distances)
    return np.cumsum(distances)


@numba.njit(cache=True)
def index_to_distance(vertices:np.ndarray, index:float) -> float:
    distances = vertices_to_distances(vertices)
    if index >= vertices.shape[0]-1: return distances[-1]
    return linear_interpolation(distances, index)


@numba.njit(cache=True)
def distance_to_index(vertices:np.ndarray, distance:float) -> float:
    if distance <= 0.0: return 0.0
    distances = vertices_to_distances(vertices)
    index = np.searchsorted(distances - distance, 0)
    if index > distances.size - 1: return distances.size - 1
    index += (distance - distances[index])/(distances[index] - distances[index-1] + EPSILON)
    return index


@numba.njit(cache=True)
def left_or_right_from_line(vertices:np.ndarray, x:float, y:float) -> int:
    i = int(interpolate_line_index_at_xy(vertices, x, y))
    i0 = max([i, 0])
    i1 = min([i + 1, vertices.shape[0] - 1])
    return np.sign((vertices[i1,0] - vertices[i0,0]) * (y - vertices[i0,1]) - (vertices[i1,1] - vertices[i0,1]) * (x - vertices[i0,0]))


@numba.njit(cache=True)
def tangent_angle(vertices:np.ndarray, index:int) -> float:
    i1 = max([0, index-1])
    i2 = min([index+1, vertices.shape[0]-1])
    delta_x = vertices[i2,0] - vertices[i1,0]
    delta_y = vertices[i2,1] - vertices[i1,1]
    ratio = delta_y / (delta_x + 1e-9)
    angle = np.arctan(ratio)
    angle = angle + np.pi if delta_x < 0 else angle
    return angle % (2*np.pi)


@numba.njit(cache=True)
def tangent_angles(vertices:np.ndarray) -> np.ndarray:
    angles = np.zeros(vertices.shape[0])
    delta_x = vertices[2:,0] - vertices[:-2,0]
    delta_y = vertices[2:,1] - vertices[:-2,1]
    ratio = delta_y / (delta_x + 1e-9)
    angles[1:-1] = np.arctan(ratio)
    angles[1:-1][np.where(delta_x < 0)] += np.pi
    angles[1:-1] = angles[1:-1] % (2*np.pi)
    angles[0] = tangent_angle(vertices, 0)
    angles[-1] = tangent_angle(vertices, vertices.shape[0]-1)
    return angles


@numba.njit(cache=True)
def angle_difference(angle1:float, angle2:float) -> float:
    two_pi = 2.0 * np.pi
    return min((two_pi - abs(angle1 - angle2), abs(angle1 - angle2))) % two_pi

@numba.njit(cache=True)
def angle_differences(angles1:float, angles2:float) -> float:
    two_pi = 2.0 * np.pi
    return np.array([min((two_pi - abs(angles1[i] - angles2[i]), abs(angles1[i] - angles2[i]))) % two_pi for i in range(angles1.size)])


@numba.njit(cache=True)
def project_point_on_line(vertices:np.ndarray, x:float, y:float) -> Tuple[float, float, float]:
    """Projection is performed in bird eye view, but the coordinate interpolation is 3D."""
    index = interpolate_line_index_at_xy(vertices, x, y)
    vx = linear_interpolation(vertices[:,0], index)
    vy = linear_interpolation(vertices[:,1], index)
    vz = linear_interpolation(vertices[:,2], index)
    return vx, vy, vz


@numba.njit(cache=True)
def radius_of_curvature(vertices: np.ndarray) -> np.ndarray:
    # https://stackoverflow.com/questions/61945776/curve-radius-in-python
    R = np.full(vertices.shape[0], np.inf)

    dx = vertices[1:,0] - vertices[:-1,0] # First derivative of x (centered)
    dy = vertices[1:,1] - vertices[:-1,1] # First derivative of y (centered)
    ddx = dx[1:] - dx[:-1] # Second derivative of x (centered)
    ddy = dy[1:] - dy[:-1] # Second derivative of y (centered)

    # Remove edges to match shapes
    dx = dx[:-1]
    dy = dy[:-1]

    R[:-2] = (dx**2 + dy**2)**1.5 / (dx*ddy - dy*ddx + EPSILON)

    # Deal with edges
    R[-2:] = R[-3]

    return R


@numba.njit(cache=True)
def apply_lateral_deviations(vertices:np.ndarray, tangent_angles:np.ndarray, lateral_deviations:Union[np.ndarray, float]) -> np.ndarray:
    """Returns the route vertices after applying the lateral deviations"""
    dx = lateral_deviations * np.sin(tangent_angles)
    dy = lateral_deviations * np.cos(tangent_angles)
    shifted_vertices = np.copy(vertices)
    shifted_vertices[:,0] -= dx
    shifted_vertices[:,1] += dy
    return shifted_vertices


@numba.njit(cache=True)
def resample_linear(vertices:np.ndarray, N:int):
    distances = vertices_to_distances(vertices)
    new_distances = np.linspace(0, distances[-1], N)
    new_vertices = np.zeros((new_distances.size, 3))
    indices = np.zeros(new_distances.size)

    for i in range(new_distances.size):
        indices[i] = distance_to_index(vertices, new_distances[i])
        new_vertices[i,0] = linear_interpolation(vertices[:,0], indices[i])
        new_vertices[i,1] = linear_interpolation(vertices[:,1], indices[i])
        new_vertices[i,2] = linear_interpolation(vertices[:,2], indices[i])

    return new_vertices, indices
