"""Generate sparse depth supervision PNGs from a COLMAP sparse model, by
projecting each image's observed 3D points (points3D.txt, via the 2D-3D
correspondences already stored in images.txt) onto its own pixel grid using
that image's COLMAP pose. Output format matches what BAGS/TriSplat's
load_gt_depth_cache expects: <source_path>/depth/<image_name>.png, uint16,
depth_units * 1000 (same convention as ScanNet sensor depth -> meters via
/1000, but here "units" = whatever scale the COLMAP reconstruction itself
is in, which is the SAME scale the renderer's surf_depth/depth will be in
since it's trained on these same poses)."""
import os, sys, numpy as np, cv2

REPO_TRI = "/srv2/szha0669/blur_slam_exp/repos/triangle-splatting"
sys.path.insert(0, REPO_TRI)
from scene.colmap_loader import read_extrinsics_text, qvec2rotmat

SPARSE = "/home/szha0669/storage/blur_slam_exp/data/i2slam_trigsplat/tum_fr2_xyz_abl1/sparse/0"
OUT_DIR = "/home/szha0669/storage/blur_slam_exp/data/i2slam_trigsplat/tum_fr2_xyz_abl1/depth"
W, H = 640, 480
MAX_DEPTH = 65.0  # uint16/1000 storage cap; drop (not clamp) farther points as
                  # likely-noisy long-tail COLMAP triangulations (p95~64.5, p99~140)


def read_points3D_xyz(path):
    pts = {}
    with open(path) as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            elems = line.split()
            pid = int(elems[0])
            pts[pid] = np.array(list(map(float, elems[1:4])))
    return pts


def main():
    os.makedirs(OUT_DIR, exist_ok=True)
    images = read_extrinsics_text(os.path.join(SPARSE, "images.txt"))
    points3D = read_points3D_xyz(os.path.join(SPARSE, "points3D.txt"))

    all_depths = []
    n_pts_per_img = []
    n_dropped = 0
    n_total = 0
    for img in images.values():
        R = qvec2rotmat(img.qvec)
        t = img.tvec
        depth_map = np.zeros((H, W), dtype=np.float32)
        n = 0
        for (x, y), pid in zip(img.xys, img.point3D_ids):
            if pid == -1 or pid not in points3D:
                continue
            xyz_cam = R @ points3D[pid] + t
            d = xyz_cam[2]
            if d <= 0:
                continue
            n_total += 1
            if d > MAX_DEPTH:
                n_dropped += 1
                continue
            px, py = int(round(x)), int(round(y))
            if 0 <= px < W and 0 <= py < H:
                if depth_map[py, px] == 0 or d < depth_map[py, px]:
                    depth_map[py, px] = d
                n += 1
                all_depths.append(d)
        n_pts_per_img.append(n)

        raw = np.clip(depth_map * 1000.0, 0, 65535).astype(np.uint16)
        out_path = os.path.join(OUT_DIR, img.name)  # name already e.g. "000000.png"
        cv2.imwrite(out_path, raw)

    all_depths = np.array(all_depths)
    print(f"images: {len(images)}, depth maps written to {OUT_DIR}")
    print(f"points/image: min={min(n_pts_per_img)} max={max(n_pts_per_img)} "
          f"mean={np.mean(n_pts_per_img):.1f}")
    print(f"depth value range: min={all_depths.min():.4f} max={all_depths.max():.4f} "
          f"mean={all_depths.mean():.4f}  (units = COLMAP pose-scale, same as "
          f"renderer's surf_depth/depth output)")
    print(f"dropped (depth > {MAX_DEPTH}): {n_dropped}/{n_total} "
          f"({100 * n_dropped / n_total:.2f}%)")


if __name__ == "__main__":
    main()