from pathlib import Path
import numpy as np
import cv2

DEPTH_DIR = Path("/data/autocon2/dataset/250827_P3/colmap/dense/stereo/depth_maps")
OUT_NPY_DIR = Path("/data/autocon2/dataset/250827_P3/colmap/dense/depth_npy")
OUT_PNG_DIR = Path("/data/autocon2/dataset/250827_P3/colmap/dense/depth_png")

OUT_NPY_DIR.mkdir(parents=True, exist_ok=True)
OUT_PNG_DIR.mkdir(parents=True, exist_ok=True)


def read_colmap_depth_map(path):
    with open(path, "rb") as f:
        header = b""
        while True:
            c = f.read(1)
            header += c
            if header.count(b"&") == 3:
                break

        width, height, channels = map(int, header.decode("ascii").split("&")[:3])
        data = np.fromfile(f, dtype=np.float32)

    depth = data.reshape((width, height, channels), order="F")
    depth = np.transpose(depth, (1, 0, 2))

    if channels == 1:
        depth = depth[:, :, 0]

    return depth


def save_png(depth, path):
    valid = (depth > 0) & np.isfinite(depth)

    if valid.sum() == 0:
        img = np.zeros_like(depth, dtype=np.uint8)
    else:
        dmin = np.percentile(depth[valid], 1)
        dmax = np.percentile(depth[valid], 99)

        depth = np.clip(depth, dmin, dmax)
        depth = (depth - dmin) / (dmax - dmin + 1e-8)
        img = (depth * 255).astype(np.uint8)

    cv2.imwrite(str(path), img)


files = sorted(DEPTH_DIR.glob("*.geometric.bin"))

print(f"Found {len(files)} depth maps")

for i, f in enumerate(files, 1):
    name = f.name.replace(".geometric.bin", "")
    stem = Path(name).stem

    depth = read_colmap_depth_map(f)

    np.save(OUT_NPY_DIR / f"{stem}.npy", depth)
    save_png(depth, OUT_PNG_DIR / f"{stem}.png")

    print(f"[{i}/{len(files)}] {stem}")

print("DONE")