core_planner.py

import numpy as np

def compute_binomial_coefficients(n):
    coefficients = np.ones(n + 1)
    for i in range(1, n):
        coefficients[i] = coefficients[i - 1] * (n - i + 1) / i
    return coefficients

def bezier_curve(points, t, binomial_coefficients, num_control_points):
    n = num_control_points - 1
    t_powers = np.array([(1 - t)**(n - i) * t**i for i in range(n + 1)])
    return np.dot(t_powers, binomial_coefficients[:, None] * points)

def generate_bezier_points(control_points, num_points, binomial_coefficients, num_control_points):
    t_values = np.linspace(0, 1, num_points)
    return np.array([bezier_curve(control_points, t, binomial_coefficients, num_control_points) for t in t_values])

def quaternion_to_euler(quaternion):
    w, x, y, z = quaternion

    t0 = 2.0 * (w * x + y * z)
    t1 = 1.0 - 2.0 * (x * x + y * y)
    roll = np.arctan2(t0, t1)

    t2 = 2.0 * (w * y - z * x)
    t2 = 1.0 if t2 > 1.0 else t2
    t2 = -1.0 if t2 < -1.0 else t2
    pitch = np.arcsin(t2)

    t3 = 2.0 * (w * z + x * y)
    t4 = 1.0 - 2.0 * (y * y + z * z)
    yaw = np.arctan2(t3, t4)

    return roll, pitch, yaw

def euler_to_quaternion(roll, pitch, yaw):
    cy = np.cos(yaw * 0.5)
    sy = np.sin(yaw * 0.5)
    cp = np.cos(pitch * 0.5)
    sp = np.sin(pitch * 0.5)
    cr = np.cos(roll * 0.5)
    sr = np.sin(roll * 0.5)

    qw = cr * cp * cy + sr * sp * sy
    qx = sr * cp * cy - cr * sp * sy
    qy = cr * sp * cy + sr * cp * sy
    qz = cr * cp * sy - sr * sp * cy

    return np.array([qw, qx, qy, qz])

def quaternion_to_rotation_matrix(quaternion):
    if len(quaternion) != 4:
        print(len(quaternion))
        raise ValueError("Quaternion must have exactly 4 components")

    w, x, y, z = quaternion

    rotation_matrix = [
        [1 - 2*y**2 - 2*z**2, 2*x*y - 2*w*z, 2*x*z + 2*w*y],
        [2*x*y + 2*w*z, 1 - 2*x**2 - 2*z**2, 2*y*z - 2*w*x],
        [2*x*z - 2*w*y, 2*y*z + 2*w*x, 1 - 2*x**2 - 2*y**2]
    ]

    return rotation_matrix

def rotation_matrix_to_quaternion(rotation_matrix):
    trace = rotation_matrix[0, 0] + rotation_matrix[1, 1] + rotation_matrix[2, 2]
    if trace > 0:
        s = 0.5 / np.sqrt(trace + 1.0)
        w = 0.25 / s
        x = (rotation_matrix[2, 1] - rotation_matrix[1, 2]) * s
        y = (rotation_matrix[0, 2] - rotation_matrix[2, 0]) * s
        z = (rotation_matrix[1, 0] - rotation_matrix[0, 1]) * s
    else:
        if rotation_matrix[0, 0] > rotation_matrix[1, 1] and rotation_matrix[0, 0] > rotation_matrix[2, 2]:
            s = 2.0 * np.sqrt(1.0 + rotation_matrix[0, 0] - rotation_matrix[1, 1] - rotation_matrix[2, 2])
            w = (rotation_matrix[2, 1] - rotation_matrix[1, 2]) / s
            x = 0.25 * s
            y = (rotation_matrix[0, 1] + rotation_matrix[1, 0]) / s
            z = (rotation_matrix[0, 2] + rotation_matrix[2, 0]) / s
        elif rotation_matrix[1, 1] > rotation_matrix[2, 2]:
            s = 2.0 * np.sqrt(1.0 + rotation_matrix[1, 1] - rotation_matrix[0, 0] - rotation_matrix[2, 2])
            w = (rotation_matrix[0, 2] - rotation_matrix[2, 0]) / s
            x = (rotation_matrix[0, 1] + rotation_matrix[1, 0]) / s
            y = 0.25 * s
            z = (rotation_matrix[1, 2] + rotation_matrix[2, 1]) / s
        else:
            s = 2.0 * np.sqrt(1.0 + rotation_matrix[2, 2] - rotation_matrix[0, 0] - rotation_matrix[1, 1])
            w = (rotation_matrix[1, 0] - rotation_matrix[0, 1]) / s
            x = (rotation_matrix[0, 2] + rotation_matrix[2, 0]) / s
            y = (rotation_matrix[1, 2] + rotation_matrix[2, 1]) / s
            z = 0.25 * s

    return np.array([w, x, y, z])

def log_map(q):
    theta = np.arccos(q[0])
    sin_theta = np.sin(theta)
    if sin_theta > 1e-6:
        return q[1:] * (theta / sin_theta)
    else:
        return q[1:]

def exp_map(v):
    theta = np.linalg.norm(v)
    if theta > 1e-6:
        sin_theta = np.sin(theta)
        q = np.concatenate(([np.cos(theta)], (v / theta) * sin_theta))
    else:
        q = np.array([1, 0, 0, 0])
    return q

def cached_manifold_interpolation(q0, q1, num_steps):
    v0 = log_map(q0)
    v1 = log_map(q1)
    cached_points = []

    for i in range(num_steps):
        t = i / (num_steps - 1)
        v_t = (1 - t) * v0 + t * v1
        cached_points.append(exp_map(v_t))

    return cached_points


def calculate_control_points(start, end, num_control_points, offset):
    control_points = [start * (1 - t) + end * t + np.array([0, offset * (-1)**(i // 2), 0]) if i % 2 == 0 else start * (1 - t) + end * t
                      for i, t in enumerate(np.linspace(0, 1, num_control_points))]
    start_quat = euler_to_quaternion(0, 0, 0)
    end_quat = euler_to_quaternion(np.pi, 0, 0)
    orientations = np.array(
        [start_quat] + [[1, 0, 0, 0]] * (num_control_points - 2) + [end_quat])
    return np.hstack((np.array(control_points), orientations))

def cost_function(control_points, obstacle, quadrotor_center, global_path_unit_vector):
    epsilon = 1e-6
    omega = 100000
    num_control_points = len(control_points)

    if np.linalg.norm(control_points[0, :3] - quadrotor_center) > 1e-6:
        return float('inf')

    curve_points = generate_bezier_points(
        control_points[:, :3], 30, compute_binomial_coefficients(num_control_points - 1), num_control_points)

    distances = obstacle.signed_distance(curve_points)
    obstacle_cost = np.sum(distances[distances < 0] ** 2)

    repulsion = np.sum(1 / (distances[distances > 0] + epsilon))

    n = num_control_points - 1
    tangent_start = n * \
        (control_points[1, :3] -
         control_points[0, :3])
    dot_product = np.dot(tangent_start, global_path_unit_vector)
    backward_penalty = max(0, -dot_product) * omega

    projections = np.dot(curve_points - quadrotor_center,
                         global_path_unit_vector)
    monotonicity_penalty = np.sum(np.maximum(
        0, projections[:-1] - projections[1:]) ** 2)

    total_cost = obstacle_cost + repulsion + \
        backward_penalty + monotonicity_penalty
    return total_cost

def adjust_control_points(control_points, obstacle, num_control_points):
    t_values = np.linspace(0, 1, 30)
    spatial_coordinates = control_points[:, :3]
    curve_points = np.array([bezier_curve(
        spatial_coordinates, t, compute_binomial_coefficients(num_control_points - 1), num_control_points) for t in t_values])

    distances_to_obstacle = np.linalg.norm(
        curve_points - obstacle.center, axis=1)
    closest_point_index = np.argmin(distances_to_obstacle)
    closest_point = curve_points[closest_point_index]

    points_per_segment = 30 // (len(control_points) - 1)
    segment_index = closest_point_index // points_per_segment

    if segment_index >= len(control_points) - 1:
        segment_index = len(control_points) - 2
    elif segment_index < 1:
        segment_index = 1

    displacement_vector = obstacle.center - closest_point
    control_points[segment_index,
                   :3] -= displacement_vector
    return control_points

def optimize_control_points(control_points, obstacle, quadrotor_center, global_path_unit_vector, num_control_points, lr, iter):
    for _ in range(iter):
        t = np.linspace(0, 1, 30)
        spatial_coordinates = control_points[:, :3]
        curve_points = np.array([bezier_curve(
            spatial_coordinates, t_val, compute_binomial_coefficients(num_control_points - 1), num_control_points) for t_val in t])

        cost = cost_function(curve_points, obstacle, quadrotor_center, global_path_unit_vector)
        if cost > 0:
            distances = obstacle.signed_distance(curve_points)
            intersecting_points = curve_points[distances < 0]
            directions = intersecting_points - obstacle.center
            directions /= np.linalg.norm(directions, axis=1)[:, None]
            gradient = np.sum(directions, axis=0)

            control_points[1:len(control_points)-1,
                           :3] -= lr * gradient
    return control_points

def ensure_no_intersection(control_points, obstacle, num_control_points, skip_radius):
    while True:
        t = np.linspace(0, 1, 30)
        spatial_coordinates = control_points[:, :3]
        curve_points = np.array([bezier_curve(
            spatial_coordinates, t_val, compute_binomial_coefficients(num_control_points - 1), num_control_points) for t_val in t])

        if np.linalg.norm(obstacle.center - curve_points[0]) < skip_radius or np.linalg.norm(obstacle.center - curve_points[-1]) < skip_radius:
            break

        if np.any(obstacle.signed_distance(curve_points) < 0):
            control_points = adjust_control_points(
                control_points, obstacle, num_control_points)
        else:
            break
    return control_points

def smooth_control_points(control_points, previous_control_points, smooth_factor):
    return (1 - smooth_factor) * previous_control_points + smooth_factor * control_points

def optimize_and_smooth_control_points(control_points, obstacles, quadrotor_center, global_path_unit_vector, num_control_points, lr, iter, smooth_factor, skip_radius, previous_control_points=None):
    for obstacle in obstacles:
        control_points = optimize_control_points(
            control_points, obstacle, quadrotor_center, global_path_unit_vector, num_control_points, lr, iter)
        control_points = ensure_no_intersection(
            control_points, obstacle, num_control_points, skip_radius)

    if previous_control_points is not None:
        control_points = smooth_control_points(
            control_points, previous_control_points, smooth_factor)
    return control_points