#!/usr/bin/env python3
"""Generate points3D.txt from TUM depth frames using scipy for pose conversion."""
import os, numpy as np
from PIL import Image
from pathlib import Path
from scipy.spatial.transform import Rotation

BASE    = "/home/szha0669/storage/blur_slam_exp"
TUM_DIR = f"{BASE}/data/TUM_RGBD/rgbd_dataset_freiburg1_desk"
OUT     = f"{BASE}/data/tum_fr1desk/scene/sparse/0/points3D.txt"

FX, FY, CX, CY = 517.3, 516.5, 318.6, 255.3
DEPTH_SCALE = 5000.0

def read_timestamps(path):
    data = {}
    with open(path) as f:
        for line in f:
            if line.startswith('#'): continue
            parts = line.strip().split()
            if len(parts) >= 2:
                try:
                    data[float(parts[0])] = parts[1:]
                except ValueError:
                    pass
    return data

def nearest_ts(query, ts_sorted, max_dt=0.05):
    idx = np.searchsorted(ts_sorted, query)
    best = None
    for i in [idx-1, idx]:
        if 0 <= i < len(ts_sorted):
            dt = abs(ts_sorted[i] - query)
            if dt < max_dt and (best is None or dt < best[0]):
                best = (dt, ts_sorted[i])
    return best[1] if best else None

rgb_data   = read_timestamps(f"{TUM_DIR}/rgb.txt")
depth_data = read_timestamps(f"{TUM_DIR}/depth.txt")
gt_data    = read_timestamps(f"{TUM_DIR}/groundtruth.txt")
depth_ts   = sorted(depth_data.keys())
gt_ts      = sorted(gt_data.keys())
rgb_ts     = sorted(rgb_data.keys())

# Sample every 30th frame
sampled = rgb_ts[::30]
print(f"Sampling {len(sampled)} frames")

all_pts, all_colors = [], []

for rgb_t in sampled:
    d_t = nearest_ts(rgb_t, depth_ts)
    g_t = nearest_ts(rgb_t, gt_ts)
    if d_t is None or g_t is None: continue

    tx, ty, tz, qx, qy, qz, qw = map(float, gt_data[g_t])
    R_c2w = Rotation.from_quat([qx, qy, qz, qw]).as_matrix()
    t_c2w = np.array([tx, ty, tz])

    depth_path = f"{TUM_DIR}/{depth_data[d_t][0]}"
    rgb_path   = f"{TUM_DIR}/{rgb_data[rgb_t][0]}"

    depth_img = np.array(Image.open(depth_path)).astype(np.float32) / DEPTH_SCALE
    rgb_img   = np.array(Image.open(rgb_path))

    H_d, W_d = depth_img.shape
    ys, xs = np.meshgrid(np.arange(0, H_d, 8), np.arange(0, W_d, 8), indexing='ij')
    ys, xs = ys.ravel(), xs.ravel()
    zs = depth_img[ys, xs]
    valid = (zs > 0.2) & (zs < 4.0)
    ys, xs, zs = ys[valid], xs[valid], zs[valid]

    xc = (xs - CX) / FX * zs
    yc = (ys - CY) / FY * zs
    pts_cam = np.stack([xc, yc, zs], axis=1)
    pts_world = (R_c2w @ pts_cam.T).T + t_c2w

    colors = rgb_img[ys, xs]
    all_pts.append(pts_world)
    all_colors.append(colors)

all_pts    = np.concatenate(all_pts, axis=0)
all_colors = np.concatenate(all_colors, axis=0)
print(f"Total: {len(all_pts)} points")

with open(OUT, 'w') as f:
    f.write("# 3D point list: POINT3D_ID X Y Z R G B ERROR TRACK[]\n")
    for i, (pt, col) in enumerate(zip(all_pts, all_colors)):
        f.write(f"{i+1} {pt[0]:.6f} {pt[1]:.6f} {pt[2]:.6f} "
                f"{int(col[0])} {int(col[1])} {int(col[2])} 0.5\n")

print(f"Written: {OUT}")