"""평면계획선 예정노선 맞춤(775 · 요구 615 · 617 · 분석 773 갈래 M) — 곡선 수 그대로 교각점 · R.

`..._Follow.merged_follow` 가 같은 쪽 합치기 뒤 부름(체인 처음 · `/route/plan` · `/route/replan` ·
`/route/follow` · 추천 「예상노선 추적」).
- `fit_to_expected` = 틀(곡선 수) 그대로 교각점 · R 를 예정노선 1 m 점 거리 제곱합 최소로(성긴
  자코비안 · 이웃 곡선만) → 가까운 일반측점 고름 → 「이웃 MC 사이 = 측점 사이」 · 「R ≥ 최소 R」
  벌점을 1 → 1000 으로 키우며 다시 맞춤 → 4 자리 · R 2 자리 반올림 → 끝 확인(MC 측점 어김 0 ·
  R ≥ 최소 R · 곡선 수 그대로) · 못 지나면 None(부르는 쪽이 지금 측점 맞춤으로 물러남 · 스냅 절대).
- `frame_for_count` = 수정 끝 틀 고르기 — 꺾임마다 곡선인 선에서 예정노선 거리가 가장 덜 늘어나는
  합침(교각점 지움 · 같은 쪽 둘 → 바깥 직선 교점)을 차례로 해 곡선 수 N · 굽이가 N 보다 적으면
  그대로.
"""

from __future__ import annotations

import numpy as np
from scipy.optimize import least_squares
from scipy.spatial import cKDTree
from shapely.geometry import LineString, Point

from B05_Profile.B05_Profile_Engine_RouteInitial_Radius import _flat
from B05_Profile.B05_Profile_Engine_RouteInitial_Station import (
    STATION_M,
    curve_mc_chainages,
    off_station_count,
)
from common_util.common_util_route_polyline import PlannedPolyline, build_planned_polyline

Xy = tuple[float, float]
PENALTY_STEPS = (1.0, 10.0, 100.0, 1000.0)


def _build(points: list[Xy], radii: list[float], min_radius_m: float) -> PlannedPolyline:
    n = len(points)
    return build_planned_polyline(
        points,
        min_radius_m=min_radius_m,
        simplify=False,
        curve_flags=[False] + [True] * (n - 2) + [False],
        radii=[None, *radii, None],
    )


def _along(coords: np.ndarray) -> np.ndarray:
    return np.concatenate([[0.0], np.cumsum(np.hypot(*np.diff(coords, axis=0).T))])


def _samples(vertices, step: float) -> np.ndarray:
    c = np.asarray(vertices, float)
    ch = _along(c)
    s = np.arange(0.0, ch[-1] + 1e-9, step)
    return np.c_[np.interp(s, ch, c[:, 0]), np.interp(s, ch, c[:, 1])]


def _plain(expected: list[Xy]) -> list[Xy]:
    """곧은 줄 위 덧점 지움 — 조밀화한 예정노선 · 원본이 같은 맞춤(체인 · `/route/follow`)."""
    return list(LineString(expected).simplify(1e-6).coords)


def _frame(line: PlannedPolyline, min_radius_m: float) -> tuple[list[Xy], list[float]]:
    points, _, radii = _flat(line)
    return points, [r if r else min_radius_m for r in radii[1:-1]]


def _fit(points0, radii0, expected, min_radius_m, targets=None, w_mc=0.0, r_pen=5.0, nfev=100):
    """교각점(끝 둘 고정) · R 최소제곱 — targets 주면 이웃 MC 사이 길이 = 측점 사이 벌점."""
    P = np.asarray(points0, float)
    n = len(P) - 2
    E = _samples(expected, 1.0)
    sE = np.arange(len(E), dtype=float)
    eline = LineString(expected)
    s_ip = np.array([eline.project(Point(p)) for p in P])
    m, n_t = len(E), (n if targets is not None else 0)
    sparsity = np.zeros((m + n + n_t, 3 * n), dtype=bool)
    for k in range(1, n + 1):
        lo, hi = s_ip[max(0, k - 2)] - 30, s_ip[min(n + 1, k + 2)] + 30
        rows = np.where((sE >= lo) & (sE <= hi))[0]
        for c in (2 * (k - 1), 2 * (k - 1) + 1, 2 * n + k - 1):
            sparsity[rows, c] = True
            sparsity[m + max(1, k - 1) - 1 : m + min(n, k + 1), c] = True
            if n_t:
                sparsity[m + n + max(1, k - 2) - 1 : m + n + min(n, k + 2), c] = True
    gaps = np.diff(np.concatenate([[0.0], targets])) if n_t else None

    def unpack(x):
        pts = [tuple(P[0]), *(tuple(p) for p in x[: 2 * n].reshape(n, 2)), tuple(P[-1])]
        return pts, list(x[2 * n :])

    def resid(x):
        pts, radii = unpack(x)
        line = _build(pts, radii, min_radius_m)
        d = cKDTree(_samples(line.vertices, 0.25)).query(E)[0]
        ok = len(line.curves) == n
        built = np.array([c.radius_m for c in line.curves]) if ok else np.zeros(n)
        out = [d, r_pen * np.maximum(0.0, min_radius_m + 0.01 - built)]
        if n_t:
            mc = np.array(curve_mc_chainages(line)) if ok else np.full(n, 1e4)
            out.append(w_mc * (np.diff(np.concatenate([[0.0], mc])) - gaps))
        return np.concatenate(out)

    x0 = np.concatenate([P[1:-1].ravel(), np.maximum(np.asarray(radii0, float), min_radius_m)])
    lower = np.concatenate([np.full(2 * n, -np.inf), np.full(n, min_radius_m)])
    scale = np.concatenate([np.full(2 * n, 5.0), np.full(n, 10.0)])
    found = least_squares(
        resid, x0, jac_sparsity=sparsity, bounds=(lower, np.inf), x_scale=scale, max_nfev=nfev
    )
    return unpack(found.x)


def _targets(line: PlannedPolyline) -> np.ndarray:
    """MC 마다 가까운 일반측점(No.1 이상 · 앞 측점보다 뒤)."""
    out, last = [], 0.0
    for mc in curve_mc_chainages(line):
        last = max(last + STATION_M, round(mc / STATION_M) * STATION_M)
        out.append(last)
    return np.array(out)


def passes(line: PlannedPolyline, count: int, min_radius_m: float) -> bool:
    """끝 확인 — 곡선 수 그대로 · MC 측점 어김 0 · R ≥ 최소 R."""
    return (
        len(line.curves) == count
        and off_station_count(line) == 0
        and all(c.radius_m >= min_radius_m - 1e-6 for c in line.curves)
    )


def fit_to_expected(
    frame: PlannedPolyline, expected: list[Xy], min_radius_m: float
) -> PlannedPolyline | None:
    """틀 곡선 수 그대로 예정노선에 맞추고 MC 를 일반측점에 놓은 선 · 끝 확인 못 지나면 None."""
    count = len(frame.curves)
    if count == 0:
        return None
    expected = _plain(expected)
    try:
        points, radii = _fit(*_frame(frame, min_radius_m), expected, min_radius_m)
        targets = _targets(_build(points, radii, min_radius_m))
        for w in PENALTY_STEPS:
            points, radii = _fit(
                points, radii, expected, min_radius_m, targets, w, max(5.0, w), nfev=60
            )
    except (ValueError, np.linalg.LinAlgError):
        return None
    points = [(round(x, 4), round(y, 4)) for x, y in points]
    line = _build(points, [round(r, 2) for r in radii], min_radius_m)
    return line if passes(line, count, min_radius_m) else None


def _gap(line: PlannedPolyline, tree: cKDTree) -> float:
    return float(np.square(tree.query(_samples(line.vertices, 1.0))[0]).sum())


def _crossing(a: Xy, b: Xy, c: Xy, d: Xy) -> Xy | None:
    """a→b 반직선 · d→c 반직선(b · c 넘어) 교점 — 둘 다 앞쪽일 때만."""
    u, v = np.subtract(b, a), np.subtract(c, d)
    den = u[0] * v[1] - u[1] * v[0]
    if abs(den) < 1e-9:
        return None
    w = np.subtract(d, a)
    t, s = (w[0] * v[1] - w[1] * v[0]) / den, (w[0] * u[1] - w[1] * u[0]) / den
    if t <= 1.0 or s <= 1.0:
        return None
    return (float(a[0] + u[0] * t), float(a[1] + u[1] * t))


def frame_for_count(
    kinked: PlannedPolyline, expected: list[Xy], min_radius_m: float, count: int
) -> PlannedPolyline:
    """꺾임마다 곡선인 선 → 예정노선 거리가 가장 덜 늘어나는 합침을 차례로 해 곡선 수 `count`."""
    tree = cKDTree(_samples(_plain(expected), 1.0))
    points, radii = _frame(kinked, min_radius_m)
    line = _build(points, radii, min_radius_m)
    # ponytail: 한 걸음마다 후보 전부 다시 짬(곡선 30 개 · 수 초) · 느리면 바뀐 곡선 둘레만 잼
    while len(points) - 2 > max(count, 0):
        best = None
        for k in range(1, len(points) - 1):
            trials = [(points[:k] + points[k + 1 :], radii[: k - 1] + radii[k:])]
            if k + 1 < len(points) - 1:
                at = _crossing(points[k - 1], points[k], points[k + 1], points[k + 2])
                if at is not None:
                    r = max(radii[k - 1], radii[k])
                    trials.append(
                        (points[:k] + [at] + points[k + 2 :], radii[: k - 1] + [r] + radii[k + 1 :])
                    )
            for pts, rr in trials:
                trial = _build(pts, rr, min_radius_m)
                if len(trial.curves) != len(pts) - 2:
                    continue
                cost = _gap(trial, tree)
                if best is None or cost < best[0]:
                    best = (cost, pts, rr, trial)
        if best is None:
            break
        _, points, radii, line = best
    return line if len(line.curves) == len(points) - 2 else kinked