# 3rd party lib
import sys

import numpy as np
from shapely import LineString


def find_corner_points(coords, angle_threshold=30):
    """Find indices of corner points where the angle is below the threshold."""
    dx = np.diff(coords[:, 0])
    dy = np.diff(coords[:, 1])

    angles = np.arctan2(dy, dx)

    # Compute absolute angle differences in-place
    angle_diffs = np.abs(angles[1:] - angles[:-1]) * 180 / np.pi
    angle_diffs = np.minimum(angle_diffs, 360 - angle_diffs)  # Ensure within [0, 180]

    # Find indices where the angle is below the threshold
    corner_indices = np.flatnonzero(angle_diffs > angle_threshold) + 1  # +1 shifts indices forward

    # Always include first and last points
    return np.r_[0, corner_indices, len(coords) - 1]  # `np.r_` avoids explicit `concatenate`


def interpolate_segment_with_original_points(segment_coords: np.ndarray, step_size: float) -> np.ndarray:
    """Interpolate a single segment while keeping original points."""
    segment = LineString(segment_coords)

    # Compute cumulative distances along the segment
    dxdy = np.diff(segment_coords, axis=0)
    seg_lengths = np.hypot(dxdy[:, 0], dxdy[:, 1])
    cum_lengths = np.r_[0, np.cumsum(seg_lengths)]  # `np.r_` for efficiency
    total_length = cum_lengths[-1]

    # Ensure original points are included
    num_steps = max(1, int(np.rint(total_length / step_size)))
    distances = np.linspace(0, total_length, num_steps)

    return interpolate_line(segment, distances)


def interpolate_curve(line: LineString, step_size: float, angle_threshold=30):
    """Segment the LineString at corners and interpolate each segment while keeping original points."""
    coords = np.array(line.coords)
    corner_indices = find_corner_points(coords, angle_threshold)

    # Create interpolated segments, ensuring no duplicate points at the boundaries
    segments = [
        interpolate_segment_with_original_points(coords[start:end + 1], step_size)
        for start, end in zip(corner_indices[:-1], corner_indices[1:])
    ]

    # Stack all segments but remove the last point of the last segment to avoid duplicate
    all_segments = np.vstack(segments)
    if len(all_segments) > 1:
        return all_segments[:-1]  # Remove the last point to avoid duplicate
    return all_segments  # Return the segment if there's only one


def interpolate_line(line: LineString, distances: np.ndarray) -> np.ndarray:
    coords = np.array(line.coords)

    # If closed, append first coordinate
    if line.is_closed:
        coords = np.vstack([coords, coords[0]])

    dxdy = np.diff(coords, axis=0)
    seg_lengths = np.hypot(dxdy[:, 0], dxdy[:, 1], dtype=np.float32)

    cum_lengths = np.r_[0, np.cumsum(seg_lengths)]  # Avoids explicit concatenate

    indices = np.searchsorted(cum_lengths, distances, side='right') - 1
    indices = np.clip(indices, 0, len(coords) - 2)

    seg_start = coords[indices]
    seg_end = coords[indices + 1]
    seg_dists = np.maximum(cum_lengths[indices + 1] - cum_lengths[indices], 1e-10)  # Prevent division by zero

    t = (distances - cum_lengths[indices]) / seg_dists
    t = np.clip(t, 0, 1)

    return seg_start + t[:, None] * (seg_end - seg_start)