#!/usr/bin/env python3
"""Shared synthetic loading-area geometry for the paper's illustration figures.

Reproduces the pipeline of subsec:centerline / subsec:guide_route as faithfully
as the available toolchain allows:

  1. rasterize the site polygon, compute the Euclidean distance transform
     (clearance field) -- this is the quantity the extended Voronoi graph
     partitions space by;
  2. extract the clearance ridge (EVG skeleton approximation) by non-maximum
     suppression of the distance transform;
  3. connect entrance to loading bay by A* over the skeleton cells, with the
     edge cost penalizing low clearance;
  4. smooth the resulting polyline with the paper's QP model -- minimize the
     second-difference (smoothness) plus reference-deviation cost subject to a
     tangential/normal box corridor, solved as a projected linear system
     (the paper uses OSQP; here the same objective is solved by direct
     factorization with projection onto the corridor box, which is adequate
     for an illustration);
  5. offset by +-W/2 along the local normal to obtain the two guidance lines.

Everything here is a *construction illustration* on a synthetic site, not an
experimental result.
"""
import numpy as np
from scipy import ndimage
from matplotlib.path import Path

# ------------------------------------------------------------------ site plan
SITE = np.array([
    [0.0, 10.0], [16.0, 6.5], [33.0, 5.0], [50.0, 6.5], [63.0, 11.0],
    [70.0, 19.0], [66.0, 30.0], [52.0, 36.0], [36.0, 35.0], [21.0, 30.0],
    [8.0, 24.0], [0.0, 20.0],
])
ENTRANCE = np.array([0.6, 15.0])       # on the open left edge
LOAD_POSE = np.array([44.0, 24.0])     # loading position inside the bay
W = 9.0                                # channel width
RES = 0.25                             # raster resolution for the DT/skeleton


def rasterize(res=RES, pad=1.0):
    x0, x1 = SITE[:, 0].min() - pad, SITE[:, 0].max() + pad
    y0, y1 = SITE[:, 1].min() - pad, SITE[:, 1].max() + pad
    xs = np.arange(x0, x1, res)
    ys = np.arange(y0, y1, res)
    XX, YY = np.meshgrid(xs, ys)
    pts = np.column_stack([XX.ravel(), YY.ravel()])
    inside = Path(SITE).contains_points(pts).reshape(XX.shape)
    return xs, ys, XX, YY, inside


def clearance(inside, res=RES):
    """Euclidean distance to the nearest boundary/obstacle cell, in metres."""
    return ndimage.distance_transform_edt(inside) * res


def skeleton(dist, inside):
    """Clearance ridge: cells whose distance value is a local maximum along at
    least one of the four axes.  A cheap stand-in for EVG-thin thinning that
    needs no extra dependency."""
    d = dist
    ridge = np.zeros_like(inside, dtype=bool)
    for ax, sh in ((0, 1), (1, 1)):
        a = np.roll(d, sh, axis=ax)
        b = np.roll(d, -sh, axis=ax)
        ridge |= (d >= a) & (d >= b)
    # a genuine ridge needs some clearance; drop the noisy skin near the wall
    return ridge & inside & (d > 1.5 * RES)


def astar(dist, inside, start_rc, goal_rc, clear_w=6.0):
    """8-connected A* over cells.  Step cost = geometric length times a factor
    that grows as clearance drops, so the path hugs the clearance ridge."""
    import heapq
    nr, nc = dist.shape
    dmax = dist.max()
    inf = float("inf")
    g = np.full(dist.shape, inf)
    came = {}
    sr, sc = start_rc
    gr, gc = goal_rc
    g[sr, sc] = 0.0

    def h(r, c):
        return np.hypot(r - gr, c - gc) * RES

    openq = [(h(sr, sc), sr, sc)]
    nbrs = [(-1, 0, 1.0), (1, 0, 1.0), (0, -1, 1.0), (0, 1, 1.0),
            (-1, -1, 1.4142), (-1, 1, 1.4142), (1, -1, 1.4142), (1, 1, 1.4142)]
    seen = np.zeros(dist.shape, dtype=bool)
    while openq:
        _, r, c = heapq.heappop(openq)
        if seen[r, c]:
            continue
        seen[r, c] = True
        if (r, c) == (gr, gc):
            break
        for dr, dc, w in nbrs:
            rr, cc = r + dr, c + dc
            if not (0 <= rr < nr and 0 <= cc < nc) or not inside[rr, cc]:
                continue
            # penalty in [1, 1+clear_w]; lowest where clearance is largest
            pen = 1.0 + clear_w * (1.0 - dist[rr, cc] / dmax)
            ng = g[r, c] + w * RES * pen
            if ng < g[rr, cc]:
                g[rr, cc] = ng
                came[(rr, cc)] = (r, c)
                heapq.heappush(openq, (ng + h(rr, cc), rr, cc))
    if (gr, gc) not in came and (gr, gc) != (sr, sc):
        raise RuntimeError("A* failed to reach the goal")
    path = [(gr, gc)]
    while path[-1] != (sr, sc):
        path.append(came[path[-1]])
    return np.array(path[::-1])


def to_xy(rc, xs, ys):
    return np.column_stack([xs[rc[:, 1]], ys[rc[:, 0]]])


def resample(poly, step):
    seg = np.diff(poly, axis=0)
    L = np.hypot(seg[:, 0], seg[:, 1])
    cum = np.concatenate([[0.0], np.cumsum(L)])
    s = np.arange(0.0, cum[-1] + 1e-9, step)
    out = np.empty((s.size, 2))
    for i, si in enumerate(s):
        k = min(np.searchsorted(cum, si, side="right") - 1, len(seg) - 1)
        k = max(k, 0)
        t = (si - cum[k]) / L[k]
        out[i] = poly[k] + t * seg[k]
    return out


def qp_smooth(ref, half_width, w_smooth=12.0, w_ref=1.0, iters=400,
              r_min=12.0):
    """Minimize  w_smooth * sum |p_{i-1} - 2 p_i + p_{i+1}|^2
              +  w_ref    * sum |p_i - p_i^ref|^2
       s.t. |p_i - p_i^ref| <= half_width   (isotropic proxy for the
                                             tangential/normal box corridor)
            |p_{i-1} - 2 p_i + p_{i+1}| <= ell^2 / r_min
                                            (curvature bound, subsec:centerline:
                                             ||d_i|| = ell^2 kappa_i + O(ell^4))
       with the endpoints fixed.

    Solved by projected gradient iterations; the paper solves the same objective
    with OSQP under linear constraints.  The curvature projection is applied as
    a Gauss-Seidel sweep: where the second difference exceeds its bound, the
    point is pulled toward the midpoint of its neighbours, which reduces |d_i|
    monotonically.
    """
    P = ref.copy()
    # Per-point step length: the second-difference bound ||d_i|| <= ell_i^2/r_min
    # is local, so use the local spacing rather than the mean (the mean would
    # under-constrain the segments that are longer than average).

    def cap_of(Q):
        """Local second-difference cap ell_i^2 / r_min, evaluated on Q's own
        spacing so that the enforced bound matches the curvature actually
        measured on the returned polyline."""
        seg = np.hypot(*np.diff(Q, axis=0).T)
        ell_i = np.minimum(seg[:-1], seg[1:])
        return ell_i ** 2 / r_min
    for _ in range(iters):
        lap = np.zeros_like(P)
        lap[1:-1] = P[:-2] - 2.0 * P[1:-1] + P[2:]
        grad = np.zeros_like(P)
        grad[1:-1] += 2.0 * w_smooth * (-2.0) * lap[1:-1]
        grad[2:-1] += 2.0 * w_smooth * lap[1:-2]
        grad[1:-2] += 2.0 * w_smooth * lap[2:-1]
        grad += 2.0 * w_ref * (P - ref)
        step = 0.02 / (w_smooth + w_ref)
        P[1:-1] = P[1:-1] - step * grad[1:-1]

        # --- corridor projection -------------------------------------------
        off = P - ref
        r = np.hypot(off[:, 0], off[:, 1])
        bad = r > half_width
        if bad.any():
            P[bad] = ref[bad] + off[bad] * (half_width / r[bad])[:, None]
        P[0], P[-1] = ref[0], ref[-1]

        # --- curvature projection, applied last so that it is the binding
        #     constraint on the returned polyline ---------------------------
        for _ in range(60):
            d_cap = cap_of(P)
            d = P[:-2] - 2.0 * P[1:-1] + P[2:]
            mag = np.hypot(d[:, 0], d[:, 1])
            over = mag > d_cap
            if not over.any():
                break
            # shift p_i along +d to shrink |d_i| toward the cap
            shrink = np.zeros_like(mag)
            shrink[over] = 0.5 * (mag[over] - d_cap[over]) / mag[over]
            P[1:-1] += d * shrink[:, None]
            P[0], P[-1] = ref[0], ref[-1]
    return P


def tangent_normal(P):
    t = np.gradient(P, axis=0)
    t /= np.linalg.norm(t, axis=1)[:, None]
    n = np.stack([-t[:, 1], t[:, 0]], axis=1)
    return t, n


def build_centerline(verbose=False):
    """Full stage-1 geometry.  Returns a dict of everything the figures need."""
    xs, ys, XX, YY, inside = rasterize()
    dist = clearance(inside)
    skel = skeleton(dist, inside)

    def nearest_cell(pt, mask):
        rr, cc = np.nonzero(mask)
        d = np.hypot(xs[cc] - pt[0], ys[rr] - pt[1])
        i = int(np.argmin(d))
        return (rr[i], cc[i])

    start = nearest_cell(ENTRANCE, inside)
    goal = nearest_cell(LOAD_POSE, inside)
    raw_rc = astar(dist, inside, start, goal)
    raw = to_xy(raw_rc, xs, ys)

    # the centerline covers the entrance->bay main channel; resample then smooth
    coarse = resample(raw, 1.0)
    center = qp_smooth(coarse, half_width=1.6)

    t, n = tangent_normal(center)
    park_line = center + (W / 2.0) * n
    exit_line = center - (W / 2.0) * n

    if verbose:
        print("grid %d x %d, res %.2f m" % (XX.shape[1], XX.shape[0], RES))
        print("max clearance %.2f m" % dist.max())
        print("A* raw length %.2f m, %d cells" %
              (np.hypot(*np.diff(raw, axis=0).T).sum(), len(raw)))
        print("centerline length %.2f m, %d points" %
              (np.hypot(*np.diff(center, axis=0).T).sum(), len(center)))
        # curvature of the smoothed centerline via second differences
        d2 = center[:-2] - 2 * center[1:-1] + center[2:]
        ell = np.hypot(*np.diff(center, axis=0).T).mean()
        kap = np.hypot(d2[:, 0], d2[:, 1]) / ell ** 2
        print("centerline |kappa| max %.4f (1/m)  -> R_min %.2f m"
              % (kap.max(), 1.0 / max(kap.max(), 1e-9)))

    return dict(xs=xs, ys=ys, XX=XX, YY=YY, inside=inside, dist=dist,
                skel=skel, raw=raw, center=center,
                park_line=park_line, exit_line=exit_line, tang=t, nrm=n)


if __name__ == "__main__":
    build_centerline(verbose=True)