"""초기 계획노선 갈래의 후보 묶음 · 비중 탐색(PLAN 23-1 · 23-4).

`B05_Profile_Engine_RouteInitial` 이 부름.

후보 묶음(`build_pool`) = 같은 입력(지형 · 예상노선 · 계류 · 제약)이면 한 번만 짠다:
  · 제약 탐색 후보 — 목표 하나씩(토공 · 기울기 · 이격 · 연장) × 탐색 상한 단계(지킬 수 있을
    때) 또는 벌점 크기(지형상 불가능할 때 · 막은 기울기 안) — PLAN 23-1 그대로.
  · 토공 ↔ 기울기 쓸기 — 비중 계산(예상 0)을 토공 100 → 기울기 100 사이 `SWEEP_STEPS` 벌.
    갈래가 그 사이 어느 비중보다 제 목표에서 좋게(토공 최소 ≤ 쓸기 절성토 · 기울기 순한 ≤
    쓸기 최대 기울기).
  · 못 지킨 곳 수가 기준(엄격: 후보 최소 · 풂: 센 벌점 기준 길)보다 많은 후보는 뺀다.

갈래는 **비교값으로** 고른다(`pick`) — 토공 최소 = 절성토 합 · 기울기 순한 = 최대(같으면 평균) ·
짧은 노선 = 연장 · 흙 균형 = |절토 − 성토|(같으면 절성토 합).

비중 계산(`weighted`)은 고르기가 아니라 **탐색 비용의 가중합** 한 번(PLAN 23-4a):
  · 몫 = 예상노선에서 벗어난 거리 · 토공 가늠 · 기울기 — 각 몫을 끝 노선 셋(추적 · 토공 최소 ·
    기울기 순한) 사이 범위로 나눠(`scale`) 비중 1 이 비슷한 크기로 움직이게.
  · 제약은 막지 않는다 — 벌점 = (1 − 예상 몫) × 센 벌점(기울기 · 계류). 예상 100 이면 벌점 0
    (추적 노선도 넘는 지형) · 예상 0 이면 갈래와 같은 세기. 넘은 곳은 `violations` 에 뜬다.
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

import numpy as np
from shapely.geometry import LineString

from B05_Profile.B05_Profile_Engine_RouteInitial_Search import LENGTH_COST, NoRoute, RouteGraph
from common_util.common_util_route_polyline import build_planned_polyline

ZERO = {"expected": 0.0, "earthwork": 0.0, "grade": 0.0}
TARGETS = {
    "earthwork": {**ZERO, "earthwork": 100.0},
    "gentle": {**ZERO, "grade": 100.0},
    "follow": {**ZERO, "expected": 100.0},
    "short": dict(ZERO),  # 목표 몫 없음 — 거리 비용만(`SHORT_LENGTH_COST`)
}
# 탐색 상한 = 상한 × 이 값들 차례로 — 곡선 맞춤 뒤 측점 사이 지반이 넘으면 다음 값으로 다시 찾음
GRADE_CAP_STEPS = (1.0, 0.93, 0.85)
# 모든 제약을 지키는 길이 없을 때만 쓰는 벌점(1 m 당 × 넘은 기울기 / 상한) — 기준 길(가장 적게
# 넘긴 길)은 센 값, 갈래 후보는 목표와 겨룰 만한 값들(목표는 1 m 당 평균 1 로 정규화됨)
SOFT_GRADE_COST = 500.0
RELAXED_GRADE_COSTS = (50.0, 5.0)
# 짧은 노선 탐색의 1 m 당 거리 비용 — 꺾기 벌점(원호 길이)과 같은 단위라 「연장 최소」가 됨
SHORT_LENGTH_COST = 1.0
# 토공 100 → 기울기 100 쓸기 벌 수(양 끝 포함 · 10 % 걸음)
SWEEP_STEPS = 11
# 끝 노선 사이 범위가 링크 평균의 이 몫보다 작으면(끝 노선끼리 같음) 링크 평균으로 크기를 맞춤
MIN_RANGE_SHARE = 0.05
# 두 노선 사이 가장 먼 거리가 이보다 작으면 「같은 노선」으로 알림(`similar` · PLAN 23-4)
SIMILAR_M = 10.0
# 비중 계산의 제약 몫(기울기 · 계류 각각) = (1 − 예상 몫) × 이 값 — 크기 맞춘 목표 한 몫과 같은 세기
CONSTRAINT_SHARE = 1.0

Measure = Callable[..., dict[str, Any]]
Route = tuple[Any, dict[str, Any]]


@dataclass
class Pool:
    """같은 입력의 후보 묶음 — 갈래 고르기 · 비중 탐색이 같이 쓴다."""

    graph: RouteGraph
    candidates: list[Route]
    bound: float | None
    limit: int  # 후보로 남길 못 지킨 곳 수 상한
    # 비중 몫 크기(1 m 당) — 이격 · 토공 · 기울기 · 넘은 기울기 · 계류 나란히
    scale: np.ndarray = field(default_factory=lambda: np.ones(5))
    # 비중 탐색 상한 — 예상 0 이면 갈래 상한(`base_cap`) · 예상 100 이면 추적 노선이 지나는
    # 가장 센 링크 기울기(`follow_cap`) · 사이는 예상 몫만큼(추적 쪽으로 갈수록 덜 막음)
    base_cap: float = np.inf
    follow_cap: float = np.inf
    searched: dict[tuple[float, ...], Route] = field(default_factory=dict)
    traced: dict[tuple, Route] = field(default_factory=dict)  # 노드 → 잰 노선(`_trace`)
    follow: Route | None = None
    # 제약 안에서 갈 수 있는 띠 폭 어림 — 후보끼리 가장 먼 거리(m)
    band_width_m: float = 0.0


def _trace(
    nodes: list[tuple[float, float]],
    min_radius_m: float,
    measure: Measure,
    traced: dict[tuple, Route],
    **kw: Any,
) -> Route:
    # 노드마다 곡선(묶지 않음) — [확인] 재구성(`/route/replan`)과 같은 규칙 → 적용 뒤도 같은 선.
    # 같은 노드 · 같은 잼이면 다시 재지 않음(쓸기는 이웃 비중이 같은 길을 자주 냄 · 묶음마다 따로).
    key = (tuple(nodes), tuple(sorted(kw.items())))
    if key not in traced:
        outline = build_planned_polyline(
            nodes, min_radius_m=min_radius_m, simplify=False, curve_flags=[True] * len(nodes)
        )
        traced[key] = (outline, measure(outline, **kw))
    return traced[key]


def _length_cost(weights: dict[str, float]) -> float:
    return LENGTH_COST if any(weights.values()) else SHORT_LENGTH_COST


def _candidates(
    graph: RouteGraph,
    max_grade: float,
    min_radius_m: float,
    measure: Measure,
    traced: dict[tuple, Route],
) -> tuple[list[Route], float | None, int]:
    """제약 탐색 후보들 · 지형상 불가능 때 막은 기울기(아니면 None) · 후보로 남길 못 지킨 곳 상한.

    제약을 지키는 길이 있으면 목표마다 탐색 상한을 `GRADE_CAP_STEPS` 차례로 낮춰 찾은 노선들 ·
    없으면 막은 기울기 안에서 벌점 크기 · 목표를 바꿔 찾은 노선들.
    """
    found: list[Route] = []
    for weights in TARGETS.values():
        for step in GRADE_CAP_STEPS:
            try:
                nodes = graph.solve(
                    weights,
                    cap=max_grade * step,
                    strict_stream=True,
                    length_cost=_length_cost(weights),
                )
            except NoRoute:
                break
            found.append(_trace(nodes, min_radius_m, measure, traced, bound=None))
            if found[-1][1]["over_limit"]["total"] == 0:
                break
        if not found:  # 제약을 지키는 길 없음 — 목표와 무관하니 더 볼 것 없음
            break
    if found:
        least = min(item[1]["over_limit"]["total"] for item in found)
        return [item for item in found if item[1]["over_limit"]["total"] == least], None, least

    # 지형상 불가능 — 반지름 · 계류 규칙을 지키는 어떤 길도 넘어야 하는 가장 작은 기울기를
    # 구해(`terrain_bound`) 그 위는 막고, 상한과 그 사이는 벌점으로 넘긴다.
    bound = graph.terrain_bound()
    cap = bound + 1e-9
    seed = _trace(
        graph.solve(
            {"expected": 50.0, "earthwork": 25.0, "grade": 25.0},
            cap=cap,
            strict_stream=False,
            grade_cost=SOFT_GRADE_COST,
        ),
        min_radius_m,
        measure,
        traced,
        bound=bound,
    )
    found = [seed]
    for weights in TARGETS.values():
        for cost in RELAXED_GRADE_COSTS:
            nodes = graph.solve(
                weights,
                cap=cap,
                strict_stream=False,
                grade_cost=cost,
                length_cost=_length_cost(weights),
            )
            found.append(_trace(nodes, min_radius_m, measure, traced, bound=bound))
    most = seed[1]["over_limit"]["total"]
    return [item for item in found if item[1]["over_limit"]["total"] <= most], bound, most


def build_pool(
    graph: RouteGraph,
    follow: Route,
    max_grade: float,
    min_radius_m: float,
    measure: Measure,
) -> Pool:
    """후보 묶음 — 제약 탐색 후보 → 끝 노선 셋으로 비중 크기 맞춤 → 토공 ↔ 기울기 쓸기를 더함."""
    traced: dict[tuple, Route] = {}
    candidates, bound, limit = _candidates(graph, max_grade, min_radius_m, measure, traced)
    pool = Pool(graph, candidates, bound, limit, traced=traced, follow=follow)
    ends = [follow[0], pick("earthwork", candidates)[0], pick("gentle", candidates)[0]]
    pool.scale, steepest = _scale(graph, [outline.vertices for outline in ends], max_grade)
    pool.base_cap = (max_grade if bound is None else bound) + 1e-9
    pool.follow_cap = max(steepest + 1e-9, pool.base_cap)
    for k in np.linspace(0.0, 1.0, SWEEP_STEPS):
        weights = {"expected": 0.0, "earthwork": 100.0 * k, "grade": 100.0 * (1.0 - k)}
        route = weighted(pool, weights, min_radius_m, measure)
        if route[1]["over_limit"]["total"] <= limit:
            pool.candidates.append(route)
    lines = list(
        {id(outline): LineString(outline.vertices) for outline, _ in pool.candidates}.values()
    )
    pool.band_width_m = max(
        (a.hausdorff_distance(b) for i, a in enumerate(lines) for b in lines[i + 1 :]), default=0.0
    )
    return pool


def similar(pool: Pool, mode: str, outline: Any) -> list[dict[str, Any]]:
    """이 노선과 `SIMILAR_M` 안인 다른 갈래 — 갈래 이름 · 가장 먼 거리(m).

    화면이 「이 지형은 제약 안에서 갈 길이 한 띠뿐 · ○○ 과 거의 같은 노선」을 알리는 데 쓴다."""
    here = LineString(outline.vertices)
    others = {"follow": pool.follow, **{name: pick(name, pool.candidates) for name in GOALS}}
    found = []
    for name, (other, _) in others.items():
        if name == mode or other is None:
            continue
        gap = here.hausdorff_distance(LineString(other.vertices))
        if gap < SIMILAR_M:
            found.append({"logic": name, "max_gap_m": round(gap, 1)})
    return sorted(found, key=lambda item: item["max_gap_m"])


def _scale(
    graph: RouteGraph, lines: list[list[tuple[float, float]]], max_grade: float
) -> tuple[np.ndarray, float]:
    """비중 몫의 크기(1 m 당) — 끝 노선 셋 사이 몫 적분 범위 / 추적 연장 · 추적의 가장 센 기울기.

    비중 1 을 옮기면 몫마다 「끝 노선 사이 차이의 1 %」만큼 무게가 옮겨 가 몫끼리 영향이 비슷하다.
    끝 노선끼리 그 몫이 거의 같으면(평지 등) 링크 평균으로 맞춤.
    """
    measured = [graph.grid.objectives(line, limit=max_grade) for line in lines]
    values = np.array([value for value, _ in measured])
    reference = max(float(np.hypot(*np.diff(np.asarray(lines[0]), axis=0).T).sum()), 1.0)
    spread = (values.max(axis=0) - values.min(axis=0)) / reference
    per_link = np.vstack([graph.terms, graph.over / graph.limit, graph.along])
    mean = (per_link * graph.length).sum(axis=1) / max(float(graph.length.sum()), 1e-9)
    return np.where(spread > MIN_RANGE_SHARE * mean, spread, mean), measured[0][1]


def weighted(pool: Pool, weights: dict[str, float], min_radius_m: float, measure: Measure) -> Route:
    """비중 계산 — 탐색 비용의 가중합 한 번(같은 비중은 다시 안 찾음). 제약은 벌점만."""
    key = tuple(round(float(weights[name]), 6) for name in ZERO)
    if key not in pool.searched:
        share = weights["expected"] / max(sum(weights.values()), 1e-9)
        soft = CONSTRAINT_SHARE * (1.0 - share)
        nodes = pool.graph.solve(
            weights,
            cap=pool.base_cap + share * (pool.follow_cap - pool.base_cap),
            strict_stream=False,
            grade_cost=soft / max(pool.scale[3], 1e-12),
            stream_cost=soft / max(pool.scale[4], 1e-12),
            scale=pool.scale[:3],
        )
        pool.searched[key] = _trace(
            nodes, min_radius_m, measure, pool.traced, bound=pool.bound, weighted=True
        )
    return pool.searched[key]


def _earth(metrics: dict[str, Any]) -> float:
    return metrics["cut_m3"] + metrics["fill_m3"]


GOALS: dict[str, Callable[[dict[str, Any]], tuple[float, ...]]] = {
    "earthwork": lambda m: (_earth(m),),
    "gentle": lambda m: (m["max_grade_pct"], m["avg_grade_pct"]),
    "short": lambda m: (m["length_m"], _earth(m)),
    "balance": lambda m: (abs(m["cut_m3"] - m["fill_m3"]), _earth(m)),
}


def pick(mode: str, candidates: list[Route]) -> Route:
    """갈래 목표 비교값이 가장 좋은 후보. 같으면 앞 후보."""
    goal = GOALS[mode]
    return min(candidates, key=lambda item: goal(item[1]))