festifun/backend/tests/test_sfm.py

"""Tests for lane B: frames.py (sampling) + sfm.py (parsers, normalization,
interpolation, and the graceful-degradation reconstruct pipeline).

Everything here is verifiable without COLMAP: parsers run on hand-written TXT snippets, the
export path runs on a hand-built (synthetic-pose) COLMAP model injected in place of a real
COLMAP run, and the failure paths are exercised by monkeypatching. Spec M2 acceptance:
"on the synthetic fixture (which skips COLMAP and injects known poses) the whole export path
runs; if COLMAP is installed, the pipeline runs end-to-end without crashing."
"""

from __future__ import annotations

import struct
from pathlib import Path

import numpy as np
import pytest

from festival4d import config, db, frames, sfm, synthetic
from festival4d.geometry import quat_to_mat


# ===========================================================================
# COLMAP TXT parsers (hand-written snippets)
# ===========================================================================
IMAGES_TXT = """\
# Image list with two lines of data per image:
#   IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME
#   POINTS2D[] as (X, Y, POINT3D_ID)
# Number of images: 2, mean observations per image: 2
1 0.9998 0.0 0.02 0.0 -1.5 0.3 8.0 1 1/1_0.jpg
480.0 270.0 5 500.1 260.2 -1
2 0.7071 0.0 0.7071 0.0 2.0 0.1 7.5 2 2/2_30.jpg
100.0 120.0 -1 300.0 200.0 5
"""

CAMERAS_TXT = """\
# Camera list with one line of data per camera:
#   CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]
# Number of cameras: 2
1 PINHOLE 640 360 500.0 500.0 320.0 180.0
2 OPENCV 640 360 510.0 511.0 321.0 181.0 0.01 -0.02 0.0 0.0
"""

POINTS3D_TXT = """\
# 3D point list with one line of data per point:
#   POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)
# Number of points: 2, mean track length: 2
5 1.0 2.0 3.0 200 100 50 0.5 1 0 2 1
9 -1.0 0.0 4.0 10 20 30 0.8 1 1
"""


def test_parse_images_txt(tmp_path):
    p = tmp_path / "images.txt"
    p.write_text(IMAGES_TXT)
    recs = sfm.parse_images_txt(p)
    assert len(recs) == 2
    assert recs[0]["image_id"] == 1
    assert recs[0]["camera_id"] == 1
    assert recs[0]["name"] == "1/1_0.jpg"
    np.testing.assert_allclose(
        [recs[0]["qw"], recs[0]["qx"], recs[0]["qy"], recs[0]["qz"]],
        [0.9998, 0.0, 0.02, 0.0])
    np.testing.assert_allclose(
        [recs[1]["tx"], recs[1]["ty"], recs[1]["tz"]], [2.0, 0.1, 7.5])


def test_parse_cameras_txt(tmp_path):
    p = tmp_path / "cameras.txt"
    p.write_text(CAMERAS_TXT)
    cams = sfm.parse_cameras_txt(p)
    assert set(cams) == {1, 2}
    assert cams[1]["model"] == "PINHOLE"
    assert (cams[1]["fx"], cams[1]["fy"], cams[1]["cx"], cams[1]["cy"]) == \
        (500.0, 500.0, 320.0, 180.0)
    # OPENCV: first 4 params are fx, fy, cx, cy; distortion ignored
    assert (cams[2]["fx"], cams[2]["fy"], cams[2]["cx"], cams[2]["cy"]) == \
        (510.0, 511.0, 321.0, 181.0)


def test_parse_points3d_txt(tmp_path):
    p = tmp_path / "points3D.txt"
    p.write_text(POINTS3D_TXT)
    points, colors = sfm.parse_points3d_txt(p)
    assert points.shape == (2, 3) and colors.shape == (2, 3)
    np.testing.assert_allclose(points[0], [1.0, 2.0, 3.0])
    assert tuple(colors[0]) == (200, 100, 50)
    assert points.dtype == np.float32 and colors.dtype == np.uint8


def test_parse_points3d_empty(tmp_path):
    p = tmp_path / "points3D.txt"
    p.write_text("# header only\n")
    points, colors = sfm.parse_points3d_txt(p)
    assert points.shape == (0, 3) and colors.shape == (0, 3)


# ===========================================================================
# Scene normalization
# ===========================================================================
def _all_synth_poses():
    poses = []
    for i in range(3):
        poses.extend(synthetic.camera_track(i, 20.0, 30, 640, 360))
    return poses


def test_normalize_scene_invariants():
    points, _ = synthetic.generate_point_cloud()
    poses = _all_synth_poses()
    npoints, nposes = sfm.normalize_scene(points.astype(float), poses)

    # centroid at origin
    assert np.linalg.norm(npoints.mean(axis=0)) < 1e-3
    # camera bounding sphere (from origin = point centroid) has radius 10
    centers = np.array([sfm._camera_center(p) for p in nposes])
    assert abs(float(np.max(np.linalg.norm(centers, axis=1))) - 10.0) < 1e-6
    # average camera-up aligns with +Y
    ups = np.array([sfm._camera_up(p) for p in nposes])
    up_avg = ups.mean(axis=0)
    up_avg /= np.linalg.norm(up_avg)
    np.testing.assert_allclose(up_avg, [0.0, 1.0, 0.0], atol=1e-6)


def test_normalize_scene_preserves_projection():
    """A world point projects to the same pixel before and after normalization
    (a similarity transform of the world must not change any image)."""
    points, _ = synthetic.generate_point_cloud()
    poses = _all_synth_poses()
    intr = synthetic.intrinsics(640, 360)
    for p in poses:
        p.update(intr)
    X = np.array([0.7, 0.9, 0.3])  # arbitrary world point

    def project(pose, Xw):
        R = quat_to_mat([pose["qw"], pose["qx"], pose["qy"], pose["qz"]])
        t = np.array([pose["tx"], pose["ty"], pose["tz"]])
        xc = R @ Xw + t
        return np.array([pose["fx"] * xc[0] / xc[2] + pose["cx"],
                         pose["fy"] * xc[1] / xc[2] + pose["cy"]])

    npoints, nposes = sfm.normalize_scene(points.astype(float), poses)
    # transform the same world point by X' = s*Rn@(X-c): recover s,Rn,c from the point map.
    # Easier: the invariant is projection equality, so map X through the same transform used
    # on the cloud by fitting it — instead, just check pixel equality using each pose pair.
    # Reconstruct transform from three non-collinear cloud points is overkill; use the fact
    # that projection is invariant, so compare original point's pixel to the transformed
    # point's pixel where the transform is inferred from the point cloud centroid+scale.
    c = points.astype(float).mean(axis=0)
    centers = np.array([sfm._camera_center(p) for p in poses])
    scale = 10.0 / float(np.max(np.linalg.norm(centers - c, axis=1)))
    ups = np.array([sfm._camera_up(p) for p in poses])
    up_avg = ups.mean(axis=0)
    Rn = sfm._rotation_aligning(up_avg, np.array([0.0, 1.0, 0.0]))
    Xn = scale * Rn @ (X - c)

    for p_old, p_new in zip(poses, nposes):
        np.testing.assert_allclose(project(p_old, X), project(p_new, Xn), atol=1e-6)


# ===========================================================================
# Pose interpolation
# ===========================================================================
def _pose(frame_idx, t, q, t3, registered=True):
    return {"frame_idx": frame_idx, "t_video_s": t,
            "qw": q[0], "qx": q[1], "qy": q[2], "qz": q[3],
            "tx": t3[0], "ty": t3[1], "tz": t3[2],
            "fx": 500.0, "fy": 500.0, "cx": 320.0, "cy": 180.0,
            "registered": registered, "video_id": 1}


def test_interpolate_poses_fills_gaps_no_extrapolation():
    q0 = [1.0, 0.0, 0.0, 0.0]
    ang = np.radians(60.0)
    q1 = [np.cos(ang / 2), 0.0, np.sin(ang / 2), 0.0]  # 60 deg yaw about Y
    registered = [_pose(0, 0.0, q0, [0, 0, 0]), _pose(60, 2.0, q1, [6, 0, 0])]
    sampled = [{"frame_idx": f, "t_video_s": f / 30.0} for f in (0, 30, 60, 90)]

    out = sfm.interpolate_poses(registered, sampled)
    by_frame = {p["frame_idx"]: p for p in out}
    assert set(by_frame) == {0, 30, 60}          # frame 90 dropped (beyond last registered)
    assert by_frame[30]["registered"] is False    # interpolated
    assert by_frame[0]["registered"] is True and by_frame[60]["registered"] is True
    # midpoint translation is the lerp; rotation is ~30 deg
    np.testing.assert_allclose(
        [by_frame[30]["tx"], by_frame[30]["ty"], by_frame[30]["tz"]], [3, 0, 0], atol=1e-9)
    mid_angle = 2.0 * np.degrees(np.arccos(min(1.0, abs(by_frame[30]["qw"]))))
    assert abs(mid_angle - 30.0) < 1e-6


def test_interpolate_poses_empty_registered():
    assert sfm.interpolate_poses([], [{"frame_idx": 0, "t_video_s": 0.0}]) == []


# ===========================================================================
# frames.py — sharpness + windowed sampling
# ===========================================================================
def test_sharpness_sharp_beats_blurred():
    import cv2

    rng = np.random.default_rng(0)
    noise = (rng.integers(0, 256, size=(120, 120), dtype=np.uint8))   # high-frequency
    flat = np.full((120, 120), 128, dtype=np.uint8)                    # no edges
    blurred = cv2.GaussianBlur(noise, (0, 0), sigmaX=5)
    assert frames.sharpness(noise) > frames.sharpness(blurred) > frames.sharpness(flat)
    # accepts BGR too
    bgr = cv2.cvtColor(noise, cv2.COLOR_GRAY2BGR)
    assert frames.sharpness(bgr) > 0


def test_sample_frames_picks_sharpest_per_window(tmp_path):
    import cv2

    w = h = 128
    fps = 30.0
    n_frames = 30                       # 1.0 s -> two 0.5 s windows
    vid_path = tmp_path / "clip.mp4"
    writer = cv2.VideoWriter(str(vid_path), cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
    if not writer.isOpened():
        pytest.skip("no usable VideoWriter codec in this environment")
    checker = np.indices((h, w)).sum(axis=0) % 2   # fine checkerboard -> sharp
    sharp = (checker * 255).astype(np.uint8)
    flat = np.full((h, w), 128, dtype=np.uint8)
    sharp_frames = {7, 22}
    for i in range(n_frames):
        g = sharp if i in sharp_frames else flat
        writer.write(cv2.cvtColor(g, cv2.COLOR_GRAY2BGR))
    writer.release()

    out = frames.sample_frames(vid_path, video_id=5, out_dir=tmp_path / "frames")
    got_frames = sorted(int(Path(p).stem.split("_")[1]) for p in out)

    # ground-truth: decode the written stream and take the per-window argmax ourselves
    cap = cv2.VideoCapture(str(vid_path))
    sh, idx = [], 0
    while True:
        ok, fr = cap.read()
        if not ok:
            break
        sh.append(frames.sharpness(fr)); idx += 1
    cap.release()
    bucket = {}
    for i, s in enumerate(sh):
        win = int((i / fps) / 0.5)
        if win not in bucket or s > bucket[win][0]:
            bucket[win] = (s, i)
    expected = sorted(v[1] for v in bucket.values())

    assert got_frames == expected
    assert all(Path(p).name.startswith("5_") for p in out)


# ===========================================================================
# run_reconstruct — graceful degradation + full export path (no real COLMAP)
# ===========================================================================
def _snapshot_poses():
    return {
        v.id: [(p.frame_idx, round(p.qw, 9), round(p.tx, 9), p.registered)
               for p in db.get_poses(v.id)]
        for v in db.get_videos()
    }


def _write_colmap_model(model_dir: Path, records: list[dict], cameras: dict,
                        points, colors) -> None:
    """Write a COLMAP TXT model (images/cameras/points3D) from pose records."""
    model_dir.mkdir(parents=True, exist_ok=True)
    with open(model_dir / "images.txt", "w") as f:
        f.write("# image list\n")
        for i, r in enumerate(records, start=1):
            f.write(f"{i} {r['qw']} {r['qx']} {r['qy']} {r['qz']} "
                    f"{r['tx']} {r['ty']} {r['tz']} {r['camera_id']} {r['name']}\n")
            f.write("100.0 100.0 -1\n")  # dummy points2D line
    with open(model_dir / "cameras.txt", "w") as f:
        f.write("# camera list\n")
        for cid, cam in cameras.items():
            f.write(f"{cid} PINHOLE {cam['width']} {cam['height']} "
                    f"{cam['fx']} {cam['fy']} {cam['cx']} {cam['cy']}\n")
    with open(model_dir / "points3D.txt", "w") as f:
        f.write("# point list\n")
        for j, (p, c) in enumerate(zip(points, colors), start=1):
            f.write(f"{j} {p[0]} {p[1]} {p[2]} {int(c[0])} {int(c[1])} {int(c[2])} 0.5 1 0\n")


def test_reconstruct_no_videos():
    db.init_engine()
    db.reset_db()
    result = sfm.run_reconstruct()
    assert result["status"] == "no_videos"


def test_reconstruct_skips_without_colmap(monkeypatch):
    synthetic.build(run_ffmpeg=False)
    monkeypatch.setattr(sfm, "colmap_available", lambda: False)
    before = _snapshot_poses()
    result = sfm.run_reconstruct()
    assert result["status"] == "skipped_no_colmap"
    assert _snapshot_poses() == before          # poses untouched


def _fake_sampling(monkeypatch, frames_per_video):
    """Make sampling hermetic: placeholder raw files + fake sample_frames output."""
    for v in db.get_videos():
        (config.RAW_DIR / v.filename).parent.mkdir(parents=True, exist_ok=True)
        (config.RAW_DIR / v.filename).write_bytes(b"stub")

    def fake_sample(video_path, video_id, out_dir, **kw):
        return [Path(f"{video_id}_{f}.jpg") for f in frames_per_video]

    monkeypatch.setattr(sfm.frames, "sample_frames", fake_sample)
    monkeypatch.setattr(sfm, "colmap_available", lambda: True)


def test_reconstruct_failed_no_model(monkeypatch):
    synthetic.build(run_ffmpeg=False)
    _fake_sampling(monkeypatch, [0, 30, 60])
    monkeypatch.setattr(sfm, "run_colmap", lambda frames_dir, workspace: None)
    before = _snapshot_poses()
    result = sfm.run_reconstruct()
    assert result["status"] == "failed_no_model"
    assert _snapshot_poses() == before          # poses untouched


def test_reconstruct_failed_weak(monkeypatch, tmp_path):
    synthetic.build(run_ffmpeg=False)
    _fake_sampling(monkeypatch, [0, 30, 60])     # 3 videos x 3 = 9 sampled
    # model registers only one image of one video -> ~11%, < 60% and < 2 videos
    track = synthetic.camera_track(0, 20.0, 30, 640, 360)
    rec = dict(track[0]); rec["camera_id"] = 1; rec["name"] = "1_0.jpg"
    cams = {1: {"width": 640, "height": 360, **synthetic.intrinsics(640, 360)}}
    pts, cols = synthetic.generate_point_cloud()
    model_dir = tmp_path / "weak"
    _write_colmap_model(model_dir, [rec], cams, pts[:50], cols[:50])
    monkeypatch.setattr(sfm, "run_colmap", lambda frames_dir, workspace: model_dir)
    before = _snapshot_poses()
    result = sfm.run_reconstruct()
    assert result["status"] == "failed_weak"
    assert _snapshot_poses() == before          # poses untouched


def test_reconstruct_full_export_path(monkeypatch, tmp_path):
    """The whole export path runs on injected known poses (spec M2 acceptance)."""
    synthetic.build(run_ffmpeg=False)
    videos = db.get_videos()
    id_by_cam = {i: videos[i].id for i in range(len(videos))}   # cam index -> db id
    sampled_frames = [0, 15, 30, 45, 60]
    _fake_sampling(monkeypatch, sampled_frames)

    # Build a model: videos for cam0/cam1 fully registered; cam2 registers a subset so
    # interpolation must fill its gaps. All frame indices come from the synthetic tracks.
    records, cameras = [], {}
    reg_plan = {0: sampled_frames, 1: sampled_frames, 2: [0, 30, 60]}
    for cam_idx, frame_list in reg_plan.items():
        vid = id_by_cam[cam_idx]
        cameras[vid] = {"width": 640, "height": 360, **synthetic.intrinsics(640, 360)}
        track = {p["frame_idx"]: p for p in synthetic.camera_track(cam_idx, 20.0, 30, 640, 360)}
        for f in frame_list:
            p = track[f]
            records.append({**p, "camera_id": vid, "name": f"{vid}_{f}.jpg"})
    pts, cols = synthetic.generate_point_cloud()
    model_dir = tmp_path / "good"
    _write_colmap_model(model_dir, records, cameras, pts, cols)
    monkeypatch.setattr(sfm, "run_colmap", lambda frames_dir, workspace: model_dir)

    result = sfm.run_reconstruct()
    assert result["status"] == "ok"
    assert result["registered"] == 13           # 5 + 5 + 3
    assert result["videos_in_model"] == 3
    assert result["interpolated"] == 2           # cam2 frames 15 and 45

    # cam2 poses: 5 total (3 registered + 2 interpolated), sorted, with interp flags
    cam2_poses = db.get_poses(id_by_cam[2])
    assert len(cam2_poses) == 5
    assert sum(1 for p in cam2_poses if not p.registered) == 2
    assert [p.frame_idx for p in cam2_poses] == [0, 15, 30, 45, 60]

    # normalized point cloud written in the frozen PLY format; camera sphere radius ~ 10
    assert config.POINTS_PLY.exists()
    npoints, _ = synthetic.read_ply(config.POINTS_PLY)
    assert len(npoints) == len(pts)
    all_centers = []
    for v in videos:
        for p in db.get_poses(v.id):
            R = quat_to_mat([p.qw, p.qx, p.qy, p.qz])
            all_centers.append(-R.T @ np.array([p.tx, p.ty, p.tz]))
    max_r = float(np.max(np.linalg.norm(np.array(all_centers), axis=1)))
    assert abs(max_r - 10.0) < 0.5              # registered cams normalized to radius 10


def test_colmap_model_parser_roundtrip(tmp_path):
    """Writing then parsing a model recovers the poses/intrinsics/points."""
    track = synthetic.camera_track(0, 20.0, 30, 640, 360)[:3]
    records = [{**p, "camera_id": 1, "name": f"1_{p['frame_idx']}.jpg"} for p in track]
    cams = {1: {"width": 640, "height": 360, **synthetic.intrinsics(640, 360)}}
    pts, cols = synthetic.generate_point_cloud()
    model_dir = tmp_path / "rt"
    _write_colmap_model(model_dir, records, cams, pts[:20], cols[:20])

    images = sfm.parse_images_txt(model_dir / "images.txt")
    cameras = sfm.parse_cameras_txt(model_dir / "cameras.txt")
    ppoints, pcolors = sfm.parse_points3d_txt(model_dir / "points3D.txt")
    assert len(images) == 3
    np.testing.assert_allclose(images[0]["qw"], track[0]["qw"], atol=1e-6)
    assert cameras[1]["fx"] == pytest.approx(synthetic.intrinsics(640, 360)["fx"])
    assert len(ppoints) == 20


def test_parse_cameras_txt_truncated_params(tmp_path):
    """A camera line with too few PARAMS gives a clean ValueError, not a raw IndexError."""
    p = tmp_path / "cameras.txt"
    p.write_text("# header\n7 PINHOLE 640 360 500.0 500.0 320.0\n")  # PINHOLE needs 4 params
    with pytest.raises(ValueError, match="too few PARAMS"):
        sfm.parse_cameras_txt(p)


def test_largest_model_dir_picks_by_registered_count(tmp_path):
    """Binary models are ranked by the images.bin header count, not file size."""
    sparse = tmp_path / "sparse"
    # model 0: 10 registered images but a LARGE images.bin (many keypoint observations)
    (sparse / "0").mkdir(parents=True)
    (sparse / "0" / "images.bin").write_bytes(struct.pack("<Q", 10) + b"\x00" * 5000)
    # model 1: 15 registered images but a SMALL images.bin (few keypoints)
    (sparse / "1").mkdir(parents=True)
    (sparse / "1" / "images.bin").write_bytes(struct.pack("<Q", 15) + b"\x00" * 100)
    assert sfm._registered_image_count(sparse / "1") == 15
    assert sfm._largest_model_dir(sparse).name == "1"   # more images wins despite smaller file


def test_interpolate_poses_drops_before_first_registered():
    """Lower no-extrapolation guard: a sampled frame before the first registered one is dropped."""
    q = [1.0, 0.0, 0.0, 0.0]
    registered = [_pose(30, 1.0, q, [3, 0, 0]), _pose(60, 2.0, q, [6, 0, 0])]
    sampled = [{"frame_idx": f, "t_video_s": f / 30.0} for f in (0, 30, 45, 60)]
    out = sfm.interpolate_poses(registered, sampled)
    frames_out = {p["frame_idx"] for p in out}
    assert 0 not in frames_out                # tv=0 < t_first=1.0 -> dropped, no extrapolation
    assert frames_out == {30, 45, 60}
    by = {p["frame_idx"]: p for p in out}
    assert by[45]["registered"] is False      # interior gap still interpolated


def test_reconstruct_subset_leaves_others_untouched(monkeypatch, tmp_path):
    """When only a subset of videos registers, the rest keep their existing poses (DB safety)."""
    synthetic.build(run_ffmpeg=False)
    videos = db.get_videos()
    id_by_cam = {i: videos[i].id for i in range(len(videos))}
    sampled_frames = [0, 15, 30, 45, 60]
    _fake_sampling(monkeypatch, sampled_frames)

    # Register cam0 + cam1 fully (10/15 = 67% >= 60%, 2 videos -> ok); cam2 NOT in the model.
    records, cameras = [], {}
    for cam_idx in (0, 1):
        vid = id_by_cam[cam_idx]
        cameras[vid] = {"width": 640, "height": 360, **synthetic.intrinsics(640, 360)}
        track = {p["frame_idx"]: p for p in synthetic.camera_track(cam_idx, 20.0, 30, 640, 360)}
        for f in sampled_frames:
            records.append({**track[f], "camera_id": vid, "name": f"{vid}_{f}.jpg"})
    pts, cols = synthetic.generate_point_cloud()
    model_dir = tmp_path / "subset"
    _write_colmap_model(model_dir, records, cameras, pts, cols)
    monkeypatch.setattr(sfm, "run_colmap", lambda frames_dir, workspace: model_dir)

    cam2_before = [(p.frame_idx, p.qw, p.tx, p.registered) for p in db.get_poses(id_by_cam[2])]
    result = sfm.run_reconstruct()
    assert result["status"] == "ok"
    assert result["videos_in_model"] == 2
    # cam2 absent from the model: its rows must be byte-identical afterward
    cam2_after = [(p.frame_idx, p.qw, p.tx, p.registered) for p in db.get_poses(id_by_cam[2])]
    assert cam2_after == cam2_before
    assert len(cam2_before) == 41                       # original synthetic poses intact
    assert len(db.get_poses(id_by_cam[0])) == 5         # cam0 replaced with sampled frames


def test_reconstruct_no_frames_when_videos_missing(monkeypatch):
    """All raw files absent -> videos skipped -> no_frames, poses untouched."""
    synthetic.build(run_ffmpeg=False)
    for f in config.RAW_DIR.glob("*.mp4"):              # clear stubs left by other tests
        f.unlink()
    monkeypatch.setattr(sfm, "colmap_available", lambda: True)
    before = _snapshot_poses()
    result = sfm.run_reconstruct()
    assert result["status"] == "no_frames"
    assert _snapshot_poses() == before


def test_reconstruct_no_frames_empty_sampling(monkeypatch):
    """Raw files exist but sampling yields nothing -> no_frames, poses untouched."""
    synthetic.build(run_ffmpeg=False)
    _fake_sampling(monkeypatch, [])                     # sample_frames returns []
    before = _snapshot_poses()
    result = sfm.run_reconstruct()
    assert result["status"] == "no_frames"
    assert _snapshot_poses() == before