"""Generic version of gen_sparse_colmap_depth.py: 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 * 1000.
"""
import os, sys, argparse, 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

MAX_DEPTH = 65.0  # uint16/1000 storage cap; drop far long-tail COLMAP points


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():
    ap = argparse.ArgumentParser()
    ap.add_argument("--sparse", required=True, help="path to sparse/0 dir")
    ap.add_argument("--out_dir", required=True)
    ap.add_argument("--w", type=int, default=640)
    ap.add_argument("--h", type=int, default=480)
    args = ap.parse_args()

    W, H = args.w, args.h
    os.makedirs(args.out_dir, exist_ok=True)
    images = read_extrinsics_text(os.path.join(args.sparse, "images.txt"))
    points3D = read_points3D_xyz(os.path.join(args.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(args.out_dir, img.name)
        cv2.imwrite(out_path, raw)

    all_depths = np.array(all_depths)
    print(f"images: {len(images)}, depth maps written to {args.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}")
    if len(all_depths):
        print(f"depth value range: min={all_depths.min():.4f} max={all_depths.max():.4f} "
              f"mean={all_depths.mean():.4f}")
    print(f"dropped (depth > {MAX_DEPTH}): {n_dropped}/{max(1,n_total)} "
          f"({100 * n_dropped / max(1,n_total):.2f}%)")


if __name__ == "__main__":
    main()