"""긴 씬의 구간 분할 (nusim.recon.batch.plan_segments) 과 증강의 프레임 배정 (nusim.aug.map_aug.assign_front_frames)."""

from types import SimpleNamespace

import numpy as np
import pytest

from nusim.aug.map_aug import assign_front_frames
from nusim.recon.batch import plan_segments

T0 = 1_533_000_000_000_000


def _scene(speeds, hz=12.0):
    """CAM_FRONT 만 있는 가짜 씬: 프레임 사이 속도 speeds[i] (m/s) 로 x 축을 따라 달린다."""
    x, frames = 0.0, []
    for i, v in enumerate(speeds):
        frames.append(SimpleNamespace(timestamp=T0 + round(i / hz * 1e6), ego_translation=(x, 0.0, 0.0)))
        x += v / hz
    return SimpleNamespace(frames={"CAM_FRONT": frames})


def test_short_scene_single_segment():
    segs = plan_segments(_scene([1.5] * 240), 60, 10)  # 30 m
    assert len(segs) == 1 and segs[0]["t0"] is None


@pytest.mark.parametrize("speed", [6.0, 12.0])
def test_segments_cover_with_overlap(speed):
    sc = _scene([speed] * 240)  # 20 초
    segs = plan_segments(sc, 60, 10)
    n = len(segs)
    total = speed * 239 / 12
    assert n == int(np.ceil((total - 10) / 50))
    for a, b in zip(segs, segs[1:]):
        assert b["t0"] < a["t1"]  # 학습 구간은 겹친다
        assert a["own1"] == b["own0"]  # 맡는 구간은 맞닿는다
        assert a["t0"] <= a["own0"] < a["own1"] <= a["t1"]
    for s in segs:  # 구간 길이 ≤ 60 m (프레임 양자화 여유)
        assert (s["t1"] - s["t0"]) * speed <= 60 + speed / 12 + 1e-6
    assert segs[0]["t0"] == 0.0 and segs[-1]["t1"] == pytest.approx(239 / 12, abs=1e-3)


def test_stop_stays_inside_segment():
    # 30 m 달리고 5 초 서 있다가 다시 60 m — 서 있던 프레임이 어느 구간에서도 빠지지 않는다
    sc = _scene([10.0] * 36 + [0.0] * 60 + [10.0] * 72)
    segs = plan_segments(sc, 60, 10)
    owned = assign_front_frames(sc, segs)
    got = [f.timestamp for fr in owned for f in fr]
    assert sorted(got) == [f.timestamp for f in sc.frames["CAM_FRONT"]]  # 프레임마다 정확히 한 구간


def test_assign_each_frame_once():
    sc = _scene([11.9] * 240)
    segs = plan_segments(sc, 60, 10)
    owned = assign_front_frames(sc, segs)
    got = [f.timestamp for fr in owned for f in fr]
    assert len(got) == len(set(got)) == len(sc.frames["CAM_FRONT"])
    start = sc.frames["CAM_FRONT"][0].timestamp
    for s, fr in zip(segs, owned):  # 맡은 프레임은 그 구간의 학습 구간 안
        ts = [(f.timestamp - start) / 1e6 for f in fr]
        assert s["t0"] - 1e-3 <= min(ts) and max(ts) <= s["t1"] + 1e-3