"""B05 종단 계획선 산출의 수치 계산부 (적분 가중치·PVI 최적화·종단곡선 기하).

[[B05_Profile_Engine_Grade]] 가 설계 기준을 확정한 뒤 이 모듈의 함수를 호출한다.
설계 기준 객체에 의존하지 않고 스칼라 인자만 받아, 기준 해석과 수치 계산을
분리한다(700줄 제한에 따른 분할).
"""

from typing import Any, Callable

import numpy as np
from scipy.optimize import LinearConstraint, minimize

# PVI 간격 산정: 최소 직선길이의 2배 이상을 두되 노선을 최대 이 개수로 나눈다.
_MAX_PVI_SEGMENTS = 40
# 종단곡선이 인접 곡선과 겹치지 않도록 확보하는 여유 비율
_CURVE_OVERLAP_RATIO = 0.9
# 구역 균형이 맞았다고 볼 평균 고저차 허용 오차 (m)
BALANCE_TOLERANCE_M = 1e-6
# 정확 균형이 기울기 기준과 양립하지 않을 때 쓰는 연화(soft) 균형 가중치
_SOFT_BALANCE_WEIGHT = 1.0e4
_TRUST_OPTIONS = {"maxiter": 1000, "gtol": 1e-10, "xtol": 1e-12}
_SLSQP_OPTIONS = {"maxiter": 500, "ftol": 1e-12}
# 균형 허용치(비율) 수렴용: 면적 상한을 다시 잡는 최대 횟수와 목표 여유 계수
_BALANCE_TIGHTEN_PASSES = 4
_BALANCE_TARGET_MARGIN = 0.8


def integration_weights(chainage: np.ndarray) -> np.ndarray:
    """사다리꼴 적분용 샘플 가중치(ds)."""
    weights = np.zeros_like(chainage)
    weights[1:-1] = (chainage[2:] - chainage[:-2]) / 2.0
    weights[0] = (chainage[1] - chainage[0]) / 2.0
    weights[-1] = (chainage[-1] - chainage[-2]) / 2.0
    return weights


def segment_weights(
    chainage: np.ndarray, lower: float, upper: float
) -> tuple[np.ndarray, np.ndarray]:
    """[lower, upper] 구간에 국한된 사다리꼴 적분 가중치를 만든다.

    전역 가중치를 그대로 쓰면 구역 경계 샘플의 반폭이 옆 구역까지 걸쳐 있어
    구역 면적이 과대 계상된다(가중치 합이 구간 길이와 어긋난다).
    """
    mask = (chainage >= lower - 1e-9) & (chainage <= upper + 1e-9)
    local = chainage[mask]
    weights = np.zeros(len(local), dtype=np.float64)
    if len(local) < 2:
        return mask, weights
    weights[1:-1] = (local[2:] - local[:-2]) / 2.0
    weights[0] = (local[1] - local[0]) / 2.0
    weights[-1] = (local[-1] - local[-2]) / 2.0
    return mask, weights


def balance_boundaries(total: float, segment_length: float | None) -> np.ndarray:
    """절·성토 균형을 맞출 구역 경계 chainage(시·종점 포함)를 만든다."""
    if not segment_length or segment_length >= total:
        return np.array([0.0, total], dtype=np.float64)
    count = max(1, int(round(total / segment_length)))
    return np.linspace(0.0, total, count + 1, dtype=np.float64)


def pvi_chainages(boundaries: np.ndarray, total: float, min_tangent: float) -> np.ndarray:
    """균형 구역 경계를 반드시 포함하는 PVI 후보 chainage를 배치한다."""
    spacing = max(min_tangent * 2.0, total / _MAX_PVI_SEGMENTS)
    nodes: list[float] = [0.0]
    for start, end in zip(boundaries[:-1], boundaries[1:]):
        span = float(end - start)
        count = max(1, int(span // spacing))
        nodes.extend(float(value) for value in np.linspace(start, end, count + 1)[1:])
    values = np.unique(np.round(np.array(nodes, dtype=np.float64), 6))
    return values[(values >= -1e-9) & (values <= total + 1e-9)]


def _interp_matrix(pvi_s: np.ndarray, targets: np.ndarray) -> np.ndarray:
    """PVI 표고 벡터를 샘플 chainage의 직선 보간 표고로 옮기는 행렬."""
    matrix = np.zeros((len(targets), len(pvi_s)), dtype=np.float64)
    index = np.clip(np.searchsorted(pvi_s, targets, side="right") - 1, 0, len(pvi_s) - 2)
    span = pvi_s[index + 1] - pvi_s[index]
    ratio = np.where(span > 0, (targets - pvi_s[index]) / np.where(span > 0, span, 1.0), 0.0)
    rows = np.arange(len(targets))
    matrix[rows, index] = 1.0 - ratio
    matrix[rows, index + 1] = ratio
    return matrix


def _polyline_area_matrix(pvi_s: np.ndarray, lower: float, upper: float) -> np.ndarray:
    """[lower, upper] 구간에서 PVI 직선 폴리라인의 면적 적분 계수 벡터."""
    coefficients = np.zeros(len(pvi_s), dtype=np.float64)
    for index in range(len(pvi_s) - 1):
        start, end = float(pvi_s[index]), float(pvi_s[index + 1])
        if end <= lower + 1e-9 or start >= upper - 1e-9:
            continue
        length = end - start
        coefficients[index] += length / 2.0
        coefficients[index + 1] += length / 2.0
    return coefficients


def grade_limits(
    max_grade_pct: float, max_reverse_grade_pct: float, main_direction: str
) -> tuple[float, float]:
    """주 진행방향 기준 (상승 상한, 하강 상한)을 비율로 정한다.

    주 진행방향과 반대인 쪽이 역기울기이므로 더 엄한 상한을 받는다.
    `none`이면 주 진행방향을 특정할 수 없는 노선(V자·Λ자)이므로 양방향 모두
    순기울기 상한을 적용한다. 방향 판정은 [[_detect_main_direction]] 이 한다.
    """
    main = max_grade_pct / 100.0
    reverse = min(max_reverse_grade_pct, max_grade_pct) / 100.0
    if main_direction == "ascending":
        return main, reverse
    if main_direction == "descending":
        return reverse, main
    return main, main


def _best_solution(
    attempts: list[tuple[str, np.ndarray]],
    objective: Callable[[np.ndarray], float],
    objective_jac: Callable[[np.ndarray], np.ndarray],
    hessian: np.ndarray,
    constraints: list[LinearConstraint],
    violation: Callable[[np.ndarray], float],
) -> tuple[np.ndarray, float]:
    """여러 (해법, 초기값) 조합을 시도해 제약 위반이 가장 작은 해를 고른다.

    trust-constr는 활성 제약이 많아지면 Jacobian이 특이해져 엉뚱한 점에서 멈추는
    경우가 있다. 결과를 직접 검증하고 통과하면 즉시 채택, 아니면 다음 조합으로
    넘어가되 최소 위반 해를 남긴다.
    """
    best: tuple[np.ndarray, float] | None = None
    for method, initial in attempts:
        extra = {"hess": lambda _x: hessian} if method == "trust-constr" else {}
        settings = _TRUST_OPTIONS if method == "trust-constr" else _SLSQP_OPTIONS
        try:
            result = minimize(
                objective,
                initial,
                jac=objective_jac,
                constraints=constraints,
                method=method,
                options=settings,
                **extra,
            )
        except (ValueError, np.linalg.LinAlgError):
            continue
        score = violation(result.x)
        if best is None or score < best[1]:
            best = (np.asarray(result.x, dtype=np.float64), score)
        if score <= BALANCE_TOLERANCE_M:
            break
    if best is None:
        return np.asarray(attempts[0][1], dtype=np.float64), float("inf")
    return best


def optimize_pvi_elevations(
    pvi_s: np.ndarray,
    chainage: np.ndarray,
    ground: np.ndarray,
    weights: np.ndarray,
    boundaries: np.ndarray,
    targets: np.ndarray,
    fixed: tuple[float, float],
    up_limit: float,
    down_limit: float,
) -> tuple[np.ndarray, bool]:
    """구역별 절·성토 균형을 등식 제약으로 두고 PVI 표고를 최적화한다.

    목적함수는 지반 추종(∫(계획고−지반고)²ds), 등식 제약은 구역별
    ∫(계획고−지반고)ds = target, 부등식 제약은 구간별 종단기울기 상한이다.
    두 번째 반환값은 정확 균형 달성 여부다.
    """
    matrix = _interp_matrix(pvi_s, chainage)
    free = np.arange(1, len(pvi_s) - 1)
    base = np.zeros(len(pvi_s), dtype=np.float64)
    base[0], base[-1] = fixed
    if not len(free):
        return base, True

    matrix_free = matrix[:, free]
    offset = matrix @ base - ground
    # 노선 길이로 정규화해 목적함수와 제약의 스케일 차이를 없앤다.
    scale = float(weights.sum()) or 1.0
    normalized = weights / scale
    # 목적함수가 순수 이차식이므로 Hessian을 해석적으로 넘겨 QP로 정확히 푼다.
    hessian = 2.0 * (matrix_free.T @ (normalized[:, None] * matrix_free))

    def objective(x: np.ndarray) -> float:
        residual = matrix_free @ x + offset
        return float(residual @ (normalized * residual))

    def objective_jac(x: np.ndarray) -> np.ndarray:
        residual = matrix_free @ x + offset
        return 2.0 * (matrix_free.T @ (normalized * residual))

    equality: list[np.ndarray] = []
    equality_rhs: list[float] = []
    for index, (lower, upper) in enumerate(zip(boundaries[:-1], boundaries[1:])):
        coefficients = _polyline_area_matrix(pvi_s, float(lower), float(upper))
        mask, local_weights = segment_weights(chainage, float(lower), float(upper))
        ground_area = float(local_weights @ ground[mask])
        # 구역 길이로 나누어 면적(m²)이 아닌 평균 고저차(m) 단위로 다룬다.
        span = max(float(upper - lower), 1e-9)
        equality.append(coefficients[free] / span)
        equality_rhs.append(
            (ground_area + float(targets[index]) - float(coefficients @ base)) / span
        )
    equality_matrix = np.vstack(equality)
    equality_vector = np.array(equality_rhs, dtype=np.float64)

    spans = pvi_s[1:] - pvi_s[:-1]
    difference = np.zeros((len(spans), len(pvi_s)), dtype=np.float64)
    rows = np.arange(len(spans))
    difference[rows, rows] = -1.0 / spans
    difference[rows, rows + 1] = 1.0 / spans
    difference_free = difference[:, free]
    difference_base = difference @ base

    # 종단기울기 상한은 법정 기준이므로 항상 강제 제약으로 둔다.
    grade_constraint = LinearConstraint(
        difference_free, -down_limit - difference_base, up_limit - difference_base
    )
    balance_constraint = LinearConstraint(equality_matrix, equality_vector, equality_vector)

    def grade_violation(x: np.ndarray) -> float:
        grades = difference_free @ x + difference_base
        return float(
            max(np.max(grades - up_limit, initial=0.0), np.max(-down_limit - grades, initial=0.0))
        )

    def total_violation(x: np.ndarray) -> float:
        balance = float(np.max(np.abs(equality_matrix @ x - equality_vector), initial=0.0))
        return max(grade_violation(x), balance)

    # 시·종점 직선과 지반 추종, 두 출발점을 준비한다(전자는 기울기 제약을 자명하게 만족).
    straight = np.interp(
        pvi_s[free],
        [float(pvi_s[0]), float(pvi_s[-1])],
        [float(base[0]), float(base[-1])],
    )
    following = np.interp(pvi_s[free], chainage, ground)
    attempts = [
        ("trust-constr", straight),
        ("SLSQP", straight),
        ("trust-constr", following),
    ]
    elevations = base.copy()
    solution, violation = _best_solution(
        attempts,
        objective,
        objective_jac,
        hessian,
        [balance_constraint, grade_constraint],
        total_violation,
    )
    if violation <= BALANCE_TOLERANCE_M:
        elevations[free] = solution
        return elevations, True

    # 정확 균형이 기울기 기준과 양립하지 않는 경우: 균형을 벌점으로 완화하되
    # 법정 기울기 제약은 그대로 강제한 채 불균형이 최소인 해를 찾는다.
    soft_hessian = hessian + 2.0 * _SOFT_BALANCE_WEIGHT * (equality_matrix.T @ equality_matrix)

    def soft_objective(x: np.ndarray) -> float:
        gap = equality_matrix @ x - equality_vector
        return objective(x) + _SOFT_BALANCE_WEIGHT * float(gap @ gap)

    def soft_jac(x: np.ndarray) -> np.ndarray:
        gap = equality_matrix @ x - equality_vector
        return objective_jac(x) + 2.0 * _SOFT_BALANCE_WEIGHT * (equality_matrix.T @ gap)

    solution, _ = _best_solution(
        attempts, soft_objective, soft_jac, soft_hessian, [grade_constraint], grade_violation
    )
    elevations[free] = solution
    return elevations, False


def _regression_prefix(chainage: np.ndarray, ground: np.ndarray) -> dict[str, np.ndarray]:
    """구간 잔차제곱합을 O(1)로 구하기 위한 접두합 묶음."""
    zeros = np.zeros(1, dtype=np.float64)
    return {
        "n": np.concatenate([zeros, np.cumsum(np.ones_like(chainage))]),
        "x": np.concatenate([zeros, np.cumsum(chainage)]),
        "xx": np.concatenate([zeros, np.cumsum(chainage * chainage)]),
        "y": np.concatenate([zeros, np.cumsum(ground)]),
        "yy": np.concatenate([zeros, np.cumsum(ground * ground)]),
        "xy": np.concatenate([zeros, np.cumsum(chainage * ground)]),
    }


def _segment_sse(prefix: dict[str, np.ndarray], start: int, end: int) -> float:
    """샘플 [start, end] 구간을 직선으로 최소자승 근사했을 때의 잔차제곱합."""
    count = prefix["n"][end + 1] - prefix["n"][start]
    if count < 2:
        return 0.0
    sum_x = prefix["x"][end + 1] - prefix["x"][start]
    sum_y = prefix["y"][end + 1] - prefix["y"][start]
    centered_xx = (prefix["xx"][end + 1] - prefix["xx"][start]) - sum_x * sum_x / count
    centered_yy = (prefix["yy"][end + 1] - prefix["yy"][start]) - sum_y * sum_y / count
    centered_xy = (prefix["xy"][end + 1] - prefix["xy"][start]) - sum_x * sum_y / count
    if centered_xx <= 1e-12:
        return float(max(centered_yy, 0.0))
    return float(max(centered_yy - centered_xy * centered_xy / centered_xx, 0.0))


def station_breakpoints(
    chainage: np.ndarray,
    ground: np.ndarray,
    stations: np.ndarray,
    *,
    penalty_m2: float,
    min_segment_stations: int,
) -> np.ndarray:
    """지반 종단을 최소 개수의 직선으로 근사하는 변화점 chainage를 고른다.

    변화점 후보를 **기준 측점으로 한정**했기 때문에 동적계획법으로 정확해를 구할
    수 있다(구간 잔차제곱합이 접두합으로 O(1)이라 전체 O(n²)). 목적함수는
    ``Σ(지반고 − 직선)² + penalty_m2 × 구간 수`` 로, penalty를 올릴수록 직선이
    길고 적어진다. 반환값은 시·종점을 포함한 변화점 chainage 배열이다.
    """
    nodes = np.unique(np.round(stations.astype(np.float64), 6))
    nodes = nodes[(nodes >= chainage[0] - 1e-6) & (nodes <= chainage[-1] + 1e-6)]
    nodes = np.unique(np.concatenate([[float(chainage[0])], nodes, [float(chainage[-1])]]))
    if len(nodes) < 3:
        return nodes

    prefix = _regression_prefix(chainage, ground)
    sample_index = np.clip(np.searchsorted(chainage, nodes), 0, len(chainage) - 1)
    step = max(1, int(min_segment_stations))
    count = len(nodes)
    best = np.full(count, np.inf, dtype=np.float64)
    previous = np.zeros(count, dtype=np.int64)
    best[0] = 0.0
    for end in range(1, count):
        for start in range(0, end - step + 1):
            if not np.isfinite(best[start]):
                continue
            cost = (
                best[start]
                + _segment_sse(prefix, int(sample_index[start]), int(sample_index[end]))
                + penalty_m2
            )
            if cost < best[end]:
                best[end] = cost
                previous[end] = start
    if not np.isfinite(best[count - 1]):
        # 최소 구간 측점 수를 만족하는 분할이 없으면 시·종점 직선 하나로 둔다.
        return np.array([nodes[0], nodes[-1]], dtype=np.float64)
    picked = [count - 1]
    while picked[-1] != 0:
        picked.append(int(previous[picked[-1]]))
    return nodes[np.array(sorted(picked), dtype=np.int64)]


def solve_alignment_elevations(
    node_s: np.ndarray,
    chainage: np.ndarray,
    ground: np.ndarray,
    weights: np.ndarray,
    fixed: tuple[float, float],
    up_limit: float,
    down_limit: float,
    *,
    balance_tolerance_percent: float,
) -> tuple[np.ndarray, bool]:
    """변화점 표고를 "지반 추종 우선 + 균형은 허용 오차 이내" 로 결정한다.

    기존 [[optimize_pvi_elevations]] 는 구역별 절·성토 균형을 **등식 제약**으로
    강제해 계획선이 지반 형상에서 멀어지곤 했다. 여기서는 지반 추종을 목적으로
    두고, 균형은 `|절토−성토| / max(절토,성토) ≤ 허용치` 를 만족할 때까지만
    순 면적을 부등식으로 조인다. 종단기울기 상한은 법정 기준이라 항상 강제한다.
    """
    matrix = _interp_matrix(node_s, chainage)
    free = np.arange(1, len(node_s) - 1)
    base = np.zeros(len(node_s), dtype=np.float64)
    base[0], base[-1] = fixed
    if not len(free):
        return base, True

    matrix_free = matrix[:, free]
    offset = matrix @ base - ground
    scale = float(weights.sum()) or 1.0
    normalized = weights / scale

    def objective(x: np.ndarray) -> float:
        residual = matrix_free @ x + offset
        return float(residual @ (normalized * residual))

    def objective_jac(x: np.ndarray) -> np.ndarray:
        residual = matrix_free @ x + offset
        return 2.0 * (matrix_free.T @ (normalized * residual))

    spans = node_s[1:] - node_s[:-1]
    difference = np.zeros((len(spans), len(node_s)), dtype=np.float64)
    rows = np.arange(len(spans))
    difference[rows, rows] = -1.0 / spans
    difference[rows, rows + 1] = 1.0 / spans
    difference_free = difference[:, free]
    difference_base = difference @ base
    grade_constraint = LinearConstraint(
        difference_free, -down_limit - difference_base, up_limit - difference_base
    )

    def solve(constraints: list[LinearConstraint], initial: np.ndarray) -> np.ndarray:
        try:
            result = minimize(
                objective,
                initial,
                jac=objective_jac,
                constraints=constraints,
                method="SLSQP",
                options=_SLSQP_OPTIONS,
            )
        except (ValueError, np.linalg.LinAlgError):
            return initial
        return np.asarray(result.x, dtype=np.float64)

    def imbalance_of(x: np.ndarray) -> tuple[float, float, float]:
        elevations = base.copy()
        elevations[free] = x
        gap = matrix @ elevations - ground
        cut = float(weights[gap < 0] @ -gap[gap < 0])
        fill = float(weights[gap > 0] @ gap[gap > 0])
        reference = max(cut, fill)
        return cut, fill, (abs(cut - fill) / reference * 100.0 if reference > 1e-9 else 0.0)

    following = np.interp(node_s[free], chainage, ground)
    solution = solve([grade_constraint], following)

    # 허용치를 넘으면 순 면적(∫(계획고−지반고)ds)을 조여 다시 푼다.
    # 허용치는 |절토−성토| / max(절토,성토) 라는 **비율**이라 면적 상한을 한 번만
    # 잡으면 재계산 후 분모(max)가 줄어들며 비율이 다시 넘칠 수 있다. 분모를 갱신하며
    # 몇 번 조여 들어가고, 더 못 줄이면 그 시점의 최선을 채택한다.
    area = _polyline_area_matrix(node_s, float(node_s[0]), float(node_s[-1]))
    ground_area = float(weights @ ground)
    span = max(float(node_s[-1] - node_s[0]), 1e-9)
    _, _, imbalance = imbalance_of(solution)
    for _ in range(_BALANCE_TIGHTEN_PASSES):
        if imbalance <= balance_tolerance_percent + 1e-9:
            break
        cut, fill, _ = imbalance_of(solution)
        # 목표를 허용치보다 조금 더 조여, 분모가 줄어도 비율이 상한 안에 남게 한다.
        tolerance_area = balance_tolerance_percent / 100.0 * max(cut, fill) * _BALANCE_TARGET_MARGIN
        balance_constraint = LinearConstraint(
            area[free][None, :] / span,
            (ground_area - tolerance_area - float(area @ base)) / span,
            (ground_area + tolerance_area - float(area @ base)) / span,
        )
        candidate = solve([grade_constraint, balance_constraint], solution)
        _, _, candidate_imbalance = imbalance_of(candidate)
        if candidate_imbalance >= imbalance - 1e-9:
            break
        solution, imbalance = candidate, candidate_imbalance
    elevations = base.copy()
    elevations[free] = solution
    return elevations, imbalance <= balance_tolerance_percent + 1e-9


def build_vertical_curves(
    pvi_s: np.ndarray,
    pvi_z: np.ndarray,
    *,
    radius_m: float,
    min_curve_length_m: float,
    skip_delta_pct: float,
    paved: bool,
) -> tuple[list[dict[str, Any]], list[str]]:
    """기울기 변화점에 종단곡선(포물선)을 삽입하고 곡선 제원을 만든다."""
    spans = pvi_s[1:] - pvi_s[:-1]
    grades = (pvi_z[1:] - pvi_z[:-1]) / spans
    skip_delta = skip_delta_pct / 100.0
    curves: list[dict[str, Any]] = []
    warnings: list[str] = []
    for index in range(1, len(pvi_s) - 1):
        grade_in, grade_out = float(grades[index - 1]), float(grades[index])
        delta = grade_out - grade_in
        if abs(delta) < 1e-9:
            continue
        # 비포장 도로이면서 대수차가 기준 이하이면 법정 예외로 곡선을 두지 않는다.
        if not paved and abs(delta) <= skip_delta + 1e-12:
            continue
        length = radius_m * abs(delta)
        available = min(float(spans[index - 1]), float(spans[index])) * _CURVE_OVERLAP_RATIO
        if length > available:
            length = available
            warnings.append(
                f"chainage {pvi_s[index]:.1f}m: 인접 직선이 짧아 종단곡선 길이를 "
                f"{length:.1f}m로 줄였습니다."
            )
        if length < min_curve_length_m - 1e-9:
            warnings.append(
                f"chainage {pvi_s[index]:.1f}m: 종단곡선 길이 {length:.1f}m가 "
                f"기준 {min_curve_length_m:.1f}m에 미달합니다."
            )
        if length <= 1e-9:
            continue
        curves.append(
            {
                "index": index,
                "chainage_m": float(pvi_s[index]),
                "elevation_m": float(pvi_z[index]),
                "grade_in": grade_in,
                "grade_out": grade_out,
                "length_m": float(length),
                "radius_m": float(length / abs(delta)),
                # 포물선 삽입으로 직선 폴리라인 대비 발생하는 면적 차 (A·L²/24)
                "area_offset_m2": float(delta * length * length / 24.0),
            }
        )
    return curves, warnings


def evaluate_profile(
    pvi_s: np.ndarray, pvi_z: np.ndarray, curves: list[dict[str, Any]], targets: np.ndarray
) -> np.ndarray:
    """직선 + 종단곡선으로 구성된 계획선을 임의 chainage에서 평가한다."""
    values = np.interp(targets, pvi_s, pvi_z)
    for curve in curves:
        half = curve["length_m"] / 2.0
        start = curve["chainage_m"] - half
        end = curve["chainage_m"] + half
        mask = (targets >= start) & (targets <= end)
        if not np.any(mask):
            continue
        local = targets[mask] - start
        start_z = curve["elevation_m"] - curve["grade_in"] * half
        delta = curve["grade_out"] - curve["grade_in"]
        values[mask] = (
            start_z + curve["grade_in"] * local + delta / (2.0 * curve["length_m"]) * local * local
        )
    return values