Blame for tests/test_closed_loop.py

"""closed-loop 평가 (nusim.sim.{hdscore,controller,policies,scenario}) — 서버·GPU 없이 합성 궤적으로."""

import json

import numpy as np
import pytest

from nusim.sim.controller import BicycleModel, ControllerConfig, WaypointTracker, dlqr, world_to_ego
from nusim.sim.geometry import (CONTACT_OBSTACLE, CONTACT_VEHICLE, ContactParams, Drivable, ObstacleGrid,
                                VehicleGeom, World)
from nusim.sim.hdscore import STEP_FIELDS, Route, comfort_signals, score_steps, ttc_ok
from nusim.sim.policies import ConstantVelocity, Observation, ReplayExpert
from nusim.sim.scenario import closed_loop_config

CFG = closed_loop_config()
HCFG = CFG["hdscore"]


def _world(obst_cells=(), road=((-50, -6), (300, -6), (300, 6), (-50, 6))):
    """y ∈ [-6, 6] 도로 (x -50~300), 장애물 칸 (x, y) 목록."""
    res, x0, y0, nx, ny = 0.2, -50.0, -20.0, 1750, 200
    g = np.zeros((ny, nx), np.uint8)
    for x, y in obst_cells:
        g[int((y - y0) / res), int((x - x0) / res)] = 255
    dr = Drivable([[np.array(road, float)]])
    return World(VehicleGeom(), ContactParams(), ObstacleGrid(g, x0, y0, res), dr)


def _steps(x, y, yaw, v, dt=0.05, decision_every=10, **cols):
    n = len(x)
    s = {k: np.zeros(n) for k in STEP_FIELDS}
    s.update(t=np.arange(n) * dt, x=np.asarray(x, float), y=np.asarray(y, float), yaw=np.asarray(yaw, float),
             v=np.asarray(v, float), on_road=np.ones(n))
    s["decision"][::decision_every] = 1
    s["decision"][-1] = 1
    for k, c in cols.items():
        s[k] = np.asarray(c, float)
    return s


def _straight(n=200, v=10.0, dt=0.05):
    x = np.arange(n) * v * dt
    return _steps(x, np.zeros(n), np.zeros(n), np.full(n, v), dt)


def test_perfect_straight_scores_one():
    st = _straight()
    route = Route(np.c_[np.linspace(0, st["x"][-1], 50), np.zeros(50)])
    r = score_steps(st, route, _world(), HCFG)
    assert r["hd_score"] == pytest.approx(1.0)
    assert r["nc"] == r["dac"] == r["ttc"] == r["com"] == 1.0
    assert r["n_decisions"] == 20


def test_route_completion_fraction():
    st = _straight(n=101)  # 0 → 50 m
    route = Route(np.c_[np.linspace(0, 100, 101), np.zeros(101)])
    r = score_steps(st, route, _world(), HCFG)
    assert r["rc"] == pytest.approx(0.5, abs=0.01)
    assert r["hd_score"] == pytest.approx(0.5, abs=0.01)
    # 완주 허용 거리 안이면 Rc = 1
    assert score_steps(st, route, _world(), HCFG, complete_m=51)["rc"] == 1.0


def test_collision_zeroes_nc_and_graze_does_not():
    st = _straight(n=100)
    st["contact"][45] = CONTACT_OBSTACLE
    st["contact_n"][45] = 1  # 스침 (graze 2 이하)
    route = Route(np.c_[[0, 100], [0, 0]])
    assert score_steps(st, route, _world(), HCFG)["nc"] == 1.0
    st["contact_n"][45] = 6
    r = score_steps(st, route, _world(), HCFG)
    assert r["per_step"][4]["nc"] == 0.0 and r["per_step"][4]["hd"] == 0.0  # 결정 구간 (40, 50]
    assert r["nc"] == pytest.approx(9 / 10)
    # 막힘(blocked)은 점 수와 무관하게 충돌. 차량 접촉도 0
    st["contact_n"][45] = 0
    st["blocked"][45] = 1
    st["contact"][45] = CONTACT_VEHICLE
    assert score_steps(st, route, _world(), HCFG)["per_step"][4]["nc"] == 0.0
    # NAVSIM 식 정적 물체 0.5
    st["contact"][45] = CONTACT_OBSTACLE
    cfg = {**HCFG, "nc": {**HCFG["nc"], "obstacle": 0.5}}
    assert score_steps(st, route, _world(), cfg)["per_step"][4]["nc"] == 0.5


def test_dac_from_on_road():
    st = _straight(n=100)
    st["on_road"][72] = 0
    r = score_steps(st, Route(np.c_[[0, 100], [0, 0]]), _world(), HCFG)
    assert [p["dac"] for p in r["per_step"]].count(0.0) == 1
    assert r["per_step"][7]["dac"] == 0.0


def test_dac_geometry_with_tolerance():
    w = _world()
    assert w.on_road(0, 0, 0)
    assert w.on_road(0, 6 - 0.87, 0)              # 차체 옆선이 경계에 딱 → 0.15 m 줄인 둘레는 안
    assert not w.on_road(0, 6 - 0.87 + 0.3, 0)     # 0.3 m 넘어감
    assert not w.on_road(0, 6 - 0.87, 0, tol=-0.1)  # 부풀린 둘레(음수 허용치)는 경계를 넘는다


def test_ttc_sees_obstacle_ahead():
    wall = [(30.0, y) for y in np.arange(-2, 2, 0.2)]
    w = _world(obst_cells=wall)
    # 앞끝 3.43 m. 1 s 등속 투영: 10 m/s 면 앞끝이 x + 13.4 까지 → 16.6 m 출발이면 벽(30)에 닿고, 10 m 출발이면 안 닿는다
    assert not ttc_ok(w, 17.0, 0.0, 0.0, 10.0, HCFG)
    assert ttc_ok(w, 10.0, 0.0, 0.0, 10.0, HCFG)
    assert ttc_ok(w, 25.0, 0.0, 0.0, 0.0, HCFG)    # 서 있으면 1
    assert ttc_ok(w, 17.0, 0.0, np.pi, 10.0, HCFG)  # 반대 방향
    # 장애물 옆을 지나가는 투영은 괜찮다
    assert ttc_ok(w, 17.0, 3.5, 0.0, 10.0, HCFG)


def test_comfort_limits():
    t = np.arange(200) * 0.05
    v = np.full(200, 8.0)
    sig = comfort_signals(t, v, np.zeros(200), 0.3)
    assert np.abs(sig["lon_accel"]).max() < 1e-9
    # 6 m/s² 급제동 → 종가속 한계 (-4.05) 밖
    v2 = np.clip(8.0 - 6.0 * np.clip(t - 5, 0, None), 0, None)
    st = _steps(np.cumsum(v2) * 0.05, np.zeros(200), np.zeros(200), v2)
    r = score_steps(st, Route(np.c_[[0, 100], [0, 0]]), _world(), HCFG)
    bad = [p for p in r["per_step"] if not p["com"]]
    assert bad and "lon_accel" in bad[0]["com_fail"]
    assert all(p["com"] for p in r["per_step"] if p["t"] < 4.5)  # 제동 전은 편하다
    # 급선회: 10 m/s 로 요레이트 1.2 rad/s (한계 0.95)
    yaw = np.where(t > 5, 1.2 * (t - 5), 0.0)
    st = _steps(np.zeros(200), np.zeros(200), yaw, np.full(200, 10.0))
    r = score_steps(st, Route(np.c_[[0, 100], [0, 0]]), _world(), HCFG)
    assert any("yaw_rate" in p["com_fail"] for p in r["per_step"])


def test_hd_score_formula():
    st = _straight(n=100)
    wall = [(46.0 + 3.43 + 0.5, y) for y in np.arange(-1, 1, 0.2)]  # 마지막 결정 (x≈49.5) 앞
    r = score_steps(st, Route(np.c_[[0, 49.5], [0, 0]]), _world(obst_cells=wall), HCFG)
    p = [q for q in r["per_step"] if q["ttc"] == 0]
    assert p, "TTC 가 0 인 결정이 있어야 한다"
    w_ttc, w_com = HCFG["w_ttc"], HCFG["w_com"]
    assert p[0]["hd"] == pytest.approx(w_com / (w_ttc + w_com))
    assert r["hd_score"] == pytest.approx(r["rc"] * np.mean([q["hd"] for q in r["per_step"]]))


def test_route_progress_monotone_and_local():
    # U 턴 경로: 되돌아오는 구간이 출발점 옆을 지나도 진행이 건너뛰지 않는다
    th = np.linspace(-np.pi / 2, np.pi / 2, 50)
    pts = np.r_[np.c_[np.linspace(0, 30, 30), np.full(30, -4.0)], np.c_[30 + 4 * np.cos(th), 4 * np.sin(th)],
                np.c_[np.linspace(30, 0, 30), np.full(30, 4.0)]]
    route = Route(pts)
    prog, dev = route.progress(np.array([0.0, 1.0, 2.0]), np.array([-3.0, -3.0, -3.0]))
    assert prog[-1] < 5.0 and dev.max() < 1.1
    assert route.command_at(20.0) == "left"
    assert route.command_at(0.0) == "straight"


# ------------------------------------------------------------------ 제어기
def _track(ref, T=20.0, y0=1.0, v0=5.0, yaw0=0.0):
    car = BicycleModel(x=0.0, y=y0, yaw=yaw0, v=v0)
    tr = WaypointTracker(ControllerConfig.from_dict(CFG["controller"]))
    out = []
    for _ in range(int(T / 0.5)):
        wp = ref(car)
        tr.set_plan(car.x, car.y, car.yaw, car.v, world_to_ego(wp, car.x, car.y, car.yaw), 0.5)
        for _ in range(10):
            a, d = tr.command(car.x, car.y, car.yaw, car.v, car.delta)
            car.step(a, d, 0.05)
            out.append((car.x, car.y, car.yaw, car.v))
    return np.array(out)


def test_dlqr_stabilizes():
    A = np.array([[1, 0.25], [0, 1.0]])
    B = np.array([[0.006], [0.1]])
    K = dlqr(A, B, np.eye(2), np.array([[1.0]]))
    assert np.abs(np.linalg.eigvals(A - B @ K)).max() < 1.0


def test_lqr_straight_line_converges():
    tr = _track(lambda c: np.c_[c.x + np.arange(1, 7) * 4.0, np.zeros(6)], y0=1.0, v0=5.0, yaw0=0.1)
    assert abs(tr[-1, 1]) < 0.05 and abs(tr[-1, 2]) < 0.01
    assert np.abs(tr[len(tr) // 2:, 1]).max() < 0.05
    assert tr[-1, 3] == pytest.approx(8.0, abs=0.2)  # 웨이포인트 간격 4 m / 0.5 s


@pytest.mark.parametrize("R", [15.0, 40.0])
def test_lqr_curve_converges(R):
    def circle(c):  # 중심 (0, R), 반시계, 6 m/s
        th = np.arctan2(c.x, R - c.y) + np.arange(1, 7) * 0.5 * 6.0 / R
        return np.c_[R * np.sin(th), R - R * np.cos(th)]

    tr = _track(circle, y0=0.5, v0=6.0)
    radial = np.hypot(tr[:, 0], tr[:, 1] - R) - R
    assert np.abs(radial[len(radial) // 2:]).max() < 0.05
    assert tr[-1, 3] == pytest.approx(6.0, abs=0.2)


def test_controller_stops_on_stationary_plan():
    car = BicycleModel(v=5.0)
    tr = WaypointTracker(ControllerConfig.from_dict(CFG["controller"]))
    for _ in range(12):
        tr.set_plan(car.x, car.y, car.yaw, car.v, np.zeros((6, 2)), 0.5)
        for _ in range(10):
            car.step(*tr.command(car.x, car.y, car.yaw, car.v, car.delta), 0.05)
    assert abs(car.v) < 0.05 and car.v >= 0  # 서되 후진하지 않는다


# ------------------------------------------------------------------ 정책
def test_replay_expert_follows_log_and_restarts_after_stop():
    t = np.arange(0, 20, 0.1)
    # 0~5 s 직진 10 m/s, 5~8 s 정차, 이후 다시 10 m/s
    v = np.where((t >= 5) & (t < 8), 0.0, 10.0)
    x = np.r_[0, np.cumsum(v[:-1] * 0.1)]
    log = {"t": t, "xy": np.c_[x, np.zeros_like(x)], "yaw": np.zeros_like(x)}
    p = ReplayExpert()
    p.reset({"log": log, "planner": CFG["planner"]})
    ego = {"x": 10.0, "y": 0.5, "yaw": 0.0, "v": 10.0, "t": 1.0}
    out = p.act(Observation(None, [], ego))
    assert out.waypoints.shape == (6, 2)
    assert out.waypoints[0, 0] == pytest.approx(5.0, abs=0.1) and out.waypoints[0, 1] == pytest.approx(-0.5, abs=1e-6)
    # 정차 지점에 서 있는 동안 시간이 흐르면 다시 출발하는 웨이포인트가 나온다
    stop_x = x[t < 5][-1] + 1.0
    for k in range(8):
        out = p.act(Observation(None, [], {"x": stop_x, "y": 0.0, "yaw": 0.0, "v": 0.0, "t": 5.0 + 0.5 * k}))
    assert out.waypoints[-1, 0] > 5.0


def test_constant_policy_goes_straight():
    p = ConstantVelocity()
    p.reset({"planner": CFG["planner"], "start_speed": 7.0})
    out = p.act(Observation(None, [], {"x": 0, "y": 0, "yaw": 1.0, "v": 3.0, "t": 0}))
    assert np.allclose(out.waypoints[:, 1], 0) and out.waypoints[-1, 0] == pytest.approx(7.0 * 3.0)


def test_scenario_files_parse():
    from nusim.sim.scenario import Scenario, scenarios_dir

    files = sorted(scenarios_dir().glob("*.json"))
    assert files, "configs/scenarios 가 비었다"
    for f in files[:10]:
        s = Scenario.load(str(f))
        assert s.scene.startswith("scene-") and s.max_duration_s > 0
        json.dumps(s.to_dict())


def test_obstacle_pgm_roundtrip(tmp_path):
    g = np.zeros((10, 20), np.uint8)
    g[3, 7] = 255
    p = tmp_path / "o.pgm"
    p.write_bytes(b"P5\n# nusim obstacles -2.000 -1.000 0.2\n20 10\n255\n" + g.tobytes())
    og = ObstacleGrid.load(p)
    assert og.occupied(np.array([[-2 + 7 * 0.2 + 0.1, -1 + 3 * 0.2 + 0.1], [0.0, 0.0], [100.0, 0.0]])).tolist() == \
        [True, False, False]