code/tests/test_geometry.py

"""actionrope/geometry.py 的纯 CPU 单测：帧/cell 映射、known/new 掩码几何、日志采样、转场连续性。

    /opt/dlami/nvme/zhiyangdeng/ActionRoPE/.venv/bin/python -m pytest tests/test_geometry.py
"""
from __future__ import annotations

import json
import os
import sys

import numpy as np
import pytest

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from actionrope import geometry as G  # noqa: E402


# ------------------------------------------------------------ cell ↔ 帧 ----
def test_cell_frames_partition():
    """21 个 cell 恰好不重不漏地覆盖 81 帧，cell 0 只含帧 0。"""
    assert G.cell_frames(0) == [0]
    assert G.cell_frames(1) == [1, 2, 3, 4]
    assert G.cell_frames(20) == [77, 78, 79, 80]
    seen = [f for k in range(G.CELLS) for f in G.cell_frames(k)]
    assert seen == list(range(G.FRAMES))
    with pytest.raises(ValueError):
        G.cell_frames(21)


def test_frames_to_cells_mean():
    off = np.zeros((G.FRAMES, 2))
    off[:, 0] = np.arange(G.FRAMES) * 2.0          # dx 每帧 +2
    off[:, 1] = -np.arange(G.FRAMES) * 0.5         # dy 每帧 −0.5
    cells = G.frames_to_cells(off)
    assert cells.shape == (G.CELLS, 2)
    assert np.allclose(cells[0], 0.0)
    # cell 1 = 帧 1..4，均值帧号 2.5
    assert np.allclose(cells[1], [5.0, -1.25])
    assert np.allclose(cells[20], [np.mean([77, 78, 79, 80]) * 2.0, -np.mean([77, 78, 79, 80]) * 0.5])
    with pytest.raises(ValueError):
        G.frames_to_cells(np.zeros((80, 2)))


# ------------------------------------------------------------ known / new ----
def test_known_mask_zero_offset_all_known():
    off = np.zeros((3, 2))
    assert G.known_mask_tokens(off).all()
    assert G.known_mask_latent(off).all()
    assert G.known_mask_tokens(off).shape == (3, G.TOK_H, G.TOK_W)
    assert G.known_mask_latent(off).shape == (3, G.LAT_H, G.LAT_W)


def test_known_mask_dx_plus_64_right_two_columns_new():
    """玩家右走 64 px ⇒ 背景左移 ⇒ 屏幕最右 2 列 token 是首帧没见过的。"""
    m = G.known_mask_tokens(np.array([[64.0, 0.0]]))[0]
    assert not m[:, -2:].any()
    assert m[:, :-2].all()
    # latent 分辨率：64 px = 4 个 latent 像素
    ml = G.known_mask_latent(np.array([[64.0, 0.0]]))[0]
    assert not ml[:, -4:].any()
    assert ml[:, :-4].all()


def test_known_mask_dy_minus_32_top_row_new():
    """玩家上走 32 px（dy<0）⇒ 背景下移 ⇒ 顶部 1 行 token 是新的。"""
    m = G.known_mask_tokens(np.array([[0.0, -32.0]]))[0]
    assert not m[0].any()
    assert m[1:].all()
    ml = G.known_mask_latent(np.array([[0.0, -32.0]]))[0]
    assert not ml[:2].any()
    assert ml[2:].all()


def test_known_mask_fractional_and_boundary():
    """半格偏移落在闭区间边界上仍算 known；再多一点就是 new。"""
    m = G.known_mask_tokens(np.array([[16.0, 0.0]]))[0]      # 0.5 token
    assert m.all()
    m = G.known_mask_tokens(np.array([[16.5, 0.0]]))[0]
    assert not m[:, -1].any() and m[:, :-1].all()
    # 左上方向：dx=−40 ⇒ −1.25 token ⇒ 只有最左 1 列 new（j=0 → −1.25 < −0.5；j=1 → −0.25 ≥ −0.5 仍 known）
    m = G.known_mask_tokens(np.array([[-40.0, 0.0]]))[0]
    assert not m[:, 0].any() and m[:, 1:].all()
    m = G.known_mask_tokens(np.array([[-48.0, 0.0]]))[0]     # −1.5 ⇒ j=1 恰在边界 −0.5，known
    assert not m[:, 0].any() and m[:, 1:].all()
    m = G.known_mask_tokens(np.array([[-49.0, 0.0]]))[0]
    assert not m[:, :2].any() and m[:, 2:].all()


def test_known_mask_per_frame_independent():
    off = np.array([[0.0, 0.0], [64.0, 0.0], [0.0, -32.0], [1000.0, 0.0]])
    m = G.known_mask_tokens(off)
    assert m[0].all()
    assert not m[1][:, -2:].any()
    assert not m[2][0].any()
    assert not m[3].any()                                    # 位移超过整幅 ⇒ 全 new


# ------------------------------------------------------------ 转场连续性 ----
def test_transition_offsets_continuity():
    rng = np.random.default_rng(0)
    a = np.cumsum(rng.normal(size=(G.FRAMES, 2)), 0); a -= a[0]
    b = np.cumsum(rng.normal(size=(G.FRAMES, 2)), 0); b -= b[0]
    n_pre = 17
    out = G.transition_frame_offsets(a, b, n_pre)
    assert out.shape == (G.FRAMES, 2)
    assert np.allclose(out[:n_pre], a[:n_pre])
    # 接缝处的增量来自 B：off[n_pre] − off[n_pre−1] = B[n_pre] − B[n_pre−1]
    assert np.allclose(out[n_pre] - out[n_pre - 1], b[n_pre] - b[n_pre - 1])
    assert np.allclose(np.diff(out[n_pre:], axis=0), np.diff(b[n_pre:], axis=0))
    assert np.allclose(out[0], 0.0)
    # 全程与 A 或 B 同源时退化为原序列
    assert np.allclose(G.transition_frame_offsets(a, a, n_pre), a)
    with pytest.raises(ValueError):
        G.transition_frame_offsets(a, b, 0)


# ------------------------------------------------------------ 日志采样 ----
def _synthetic_tracks(tmp_path, offset_segments, M_px1280, n=400, dt=0.05, gap=None):
    """造一个最小 session：直线匀速走 x，可选在中间挖一段日志断档。"""
    vt = np.arange(n) * dt
    x = 2.0 * vt                       # 2 世界单位/s
    y = np.zeros(n); z = np.zeros(n)
    if gap is not None:
        keep = (vt < gap[0]) | (vt > gap[1])
        vt, x, y, z = vt[keep], x[keep], y[keep], z[keep]
    # 故意加重复时间戳，load_session_tracks 必须去掉
    vt = np.concatenate([vt, vt[:3]]); x = np.concatenate([x, x[:3]])
    y = np.concatenate([y, y[:3]]); z = np.concatenate([z, z[:3]])
    d = tmp_path / "sess"
    d.mkdir()
    np.savez(d / "state.npz", vt=vt, x=x.astype(np.float32), y=y.astype(np.float32),
             z=z.astype(np.float32), qf=np.arange(vt.size), rg=np.zeros(vt.size, np.int32))
    (d / "offsets.json").write_text(json.dumps(dict(segments=offset_segments, M_px1280=M_px1280)))
    return G.load_session_tracks(str(tmp_path), "sess")


def test_load_session_tracks_dedupe_and_scale(tmp_path):
    M = [[-40.0, 0.0], [0.0, 18.0], [0.0, 27.0]]
    tr = _synthetic_tracks(tmp_path, [dict(vt_lo=0.0, vt_hi=1e12, offset=0.1)], M)
    assert tr["vt"].size == 400 and np.all(np.diff(tr["vt"]) > 0)
    assert np.allclose(tr["M_px832"][:, 0], np.array(M)[:, 0] * 832 / 1280)
    assert np.allclose(tr["M_px832"][:, 1], np.array(M)[:, 1] * 480 / 720)
    assert np.allclose(tr["vid_t"], tr["vt"] + 0.1)
    u, v = G.project_px832(tr, 1.0, 0.0, 0.0)
    assert np.isclose(u, -40.0 * 832 / 1280) and np.isclose(v, 0.0)


def test_offset_piecewise_constant(tmp_path):
    M = [[-40.0, 0.0], [0.0, 18.0], [0.0, 27.0]]
    segs = [dict(vt_lo=0.0, vt_hi=10.0, offset=0.1), dict(vt_lo=10.0, vt_hi=1e12, offset=-4.0)]
    tr = _synthetic_tracks(tmp_path, segs, M)
    assert G.offset_at(tr, 5.0) == 0.1
    assert G.offset_at(tr, 10.0) == -4.0          # 段以 vt_lo 为界，右侧闭
    assert G.offset_at(tr, 15.0) == -4.0


def test_sample_frame_offsets_linear_walk(tmp_path):
    """匀速沿世界 +x 走 2/s：M[0,0]=−40 ⇒ 1280 画幅下 −80 px/s ⇒ 832 画幅 −52 px/s。"""
    M = [[-40.0, 0.0], [0.0, 18.0], [0.0, 27.0]]
    tr = _synthetic_tracks(tmp_path, [dict(vt_lo=0.0, vt_hi=1e12, offset=0.1)], M)
    off, valid, reason = G.sample_frame_offsets(tr, video_t0=2.1, return_valid=True)
    assert valid and reason == ""
    assert off.shape == (G.FRAMES, 2)
    assert np.allclose(off[0], 0.0)
    # 最近邻采样：20 Hz 日志对 16 fps 的帧最多差 0.025 s ⇒ 1.3 px
    expect = -52.0 * np.arange(G.FRAMES) / G.FPS
    assert np.abs(off[:, 0] - expect).max() < 1.5
    assert np.allclose(off[:, 1], 0.0)
    assert np.isfinite(G.sample_frame_offsets(tr, 2.1)).all()


def test_sample_frame_offsets_gap_invalid(tmp_path):
    M = [[-40.0, 0.0], [0.0, 18.0], [0.0, 27.0]]
    tr = _synthetic_tracks(tmp_path, [dict(vt_lo=0.0, vt_hi=1e12, offset=0.0)], M, gap=(5.0, 5.6))
    off, valid, reason = G.sample_frame_offsets(tr, video_t0=3.0, return_valid=True)
    assert not valid and reason.startswith("log_gap")
    assert np.isnan(G.sample_frame_offsets(tr, 3.0)).all()
    # 阈值放宽到 0.5 s 就有效（最近邻误差最大 0.3 s）
    _, valid2, _ = G.sample_frame_offsets(tr, 3.0, return_valid=True, max_sample_dt=0.5)
    assert valid2
    # 完全落在日志范围之外 ⇒ 断档
    _, valid3, reason3 = G.sample_frame_offsets(tr, video_t0=100.0, return_valid=True)
    assert not valid3 and reason3.startswith("log_gap")


def test_sample_frame_offsets_teleport_invalid(tmp_path):
    M = [[-40.0, 0.0], [0.0, 18.0], [0.0, 27.0]]
    tr = _synthetic_tracks(tmp_path, [dict(vt_lo=0.0, vt_hi=1e12, offset=0.0)], M)
    tr["x"] = tr["x"].copy(); tr["x"][tr["vt"] > 4.0] += 50.0          # 传送 50 世界单位 = 1300 px
    _, valid, reason = G.sample_frame_offsets(tr, video_t0=3.0, return_valid=True)
    assert not valid and reason.startswith("teleport")


def test_sample_frame_offsets_log_jump_invalid(tmp_path):
    """offset 台阶：vt=10 处 offset 从 0 跳到 −4 ⇒ vid_t 局部非单调，跨过台阶的 clip 最近邻会跳到
    4 s 之后的日志样本（时间差仍 < 0.15 s），必须按 log_jump 判 invalid。"""
    M = [[-40.0, 0.0], [0.0, 18.0], [0.0, 27.0]]
    segs = [dict(vt_lo=0.0, vt_hi=10.0, offset=0.0), dict(vt_lo=10.0, vt_hi=1e12, offset=-4.0)]
    tr = _synthetic_tracks(tmp_path, segs, M)
    assert np.diff(tr["vid_t"]).min() < 0
    _, valid, reason = G.sample_frame_offsets(tr, video_t0=4.0, return_valid=True)
    assert not valid and reason.startswith("log_jump")
    # 台阶之前 / 之后整段落在单调区间里的 clip 不受影响
    _, valid_a, _ = G.sample_frame_offsets(tr, video_t0=0.5, return_valid=True)
    _, valid_b, _ = G.sample_frame_offsets(tr, video_t0=8.0, return_valid=True)
    assert valid_a and valid_b


def test_action_dirs_unit():
    n = np.hypot(G.ACTION_DIRS[:, 0], G.ACTION_DIRS[:, 1])
    assert np.isclose(n[0], 0.0) and np.allclose(n[1:], 1.0)
    assert np.allclose(G.ACTION_DIRS[7], [1.0, 0.0])        # moving right
    assert np.allclose(G.ACTION_DIRS[1], [0.0, -1.0])       # moving up（屏幕 v 向下为正）


def test_shaped_loss_weight():
    """dist 形状：known=1，紧贴足迹边界 ≈1，远处到 new_weight；σ 调制只作用于额外部分。"""
    import torch
    from actionrope.arope import loss_weight_map
    off = torch.tensor([[[0.0, 0.0], [96.0, 0.0]]])  # 帧 1 右移 96 px = 6 latent 格 ⇒ 最右 6 列 new
    flat = loss_weight_map(off, 30, 52, new_weight=3.0)[0, 0]
    assert flat[1, :, :46].eq(1).all() and flat[1, :, 46:].eq(3).all() and flat[0].eq(1).all()
    w = loss_weight_map(off, 30, 52, new_weight=3.0, ramp_latent=4.0)[0, 0]
    assert w[0].eq(1).all() and w[1, :, :46].eq(1).all()
    col = w[1, 0, 46:].tolist()          # 距边界 0.5,1.5,…,5.5 格 ⇒ 1+2·min(d/4,1)
    assert abs(col[0] - 1.25) < 1e-5 and abs(col[1] - 1.75) < 1e-5 and abs(col[3] - 2.75) < 1e-5 and col[4] == 3.0 and col[5] == 3.0
    ws = loss_weight_map(off, 30, 52, new_weight=3.0, ramp_latent=4.0, sigma=torch.tensor([1.0]))[0, 0]
    assert abs(ws[1, 0, 51] - 4.0) < 1e-5     # 1 + 2·1.5
    w0 = loss_weight_map(off, 30, 52, new_weight=3.0, ramp_latent=4.0, sigma=torch.tensor([0.0]))[0, 0]
    assert abs(w0[1, 0, 51] - 2.0) < 1e-5     # 1 + 2·0.5
    assert ws[1, :, :46].eq(1).all()