#!/usr/bin/env python3
"""
scripts/preview_raw_map.py

快速预览工具：仅使用 simulated_gnss.json 中的 origin_xyz 对 PCD 进行简单平移拼接。
不进行任何旋转配准，也不进行 ICP 优化。
用于检查 GNSS 坐标是否大体正确，以及 PCD 加载是否正常。

Usage:
  python scripts/preview_raw_map.py --pcd_dir data/HK_PCD_DS --out results/raw_preview.pcd
"""

import argparse
import json
import re
import numpy as np
import open3d as o3d
from pathlib import Path

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--pcd_dir", type=str, default="data/HK_PCD")
    parser.add_argument("--json", type=str, default="data/HK_POSE/simulated_gnss.json")
    parser.add_argument("--out", type=str, default="results/raw_preview.pcd")
    parser.add_argument("--voxel", type=float, default=0.5, help="Downsample voxel size for preview")
    args = parser.parse_args()

    json_path = Path(args.json)
    if not json_path.exists():
        print(f"Error: JSON not found {json_path}")
        return

    with open(json_path, "r", encoding="utf-8") as f:
        info = json.load(f)

    # Sort blocks
    block_ids = sorted(info.keys(), key=lambda x: int(x.split("_")[-1]))
    print(f"Found blocks: {block_ids}")

    # Map PCDs
    pcd_dir = Path(args.pcd_dir)
    pcd_files = list(pcd_dir.glob("*.pcd"))
    
    merged = o3d.geometry.PointCloud()

    print(f"\nProcessing {len(block_ids)} blocks...")
    
    for bid in block_ids:
        # 1. Get Origin (Translation)
        origin = info[bid].get("origin_xyz", [0,0,0])
        t = np.array(origin, dtype=np.float64)
        
        # 2. Find PCD
        # Match ..._01.pcd -> block_1
        idx = int(bid.split("_")[-1])
        found_pcd = None
        for p in pcd_files:
            m = re.search(r"(\d+)\.pcd$", p.name)
            if m and int(m.group(1)) == idx:
                found_pcd = p
                break
        
        if not found_pcd:
            print(f"  [Warn] PCD not found for {bid} in {pcd_dir}")
            continue

        # 3. Load & Transform
        print(f"  [Load] {bid} -> {found_pcd.name} | Origin={t}")
        pcd = o3d.io.read_point_cloud(str(found_pcd))
        
        if pcd.is_empty():
            print("    -> Empty PCD, skip.")
            continue
            
        # Optional: pre-downsample to save memory
        if args.voxel > 0:
            pcd = pcd.voxel_down_sample(args.voxel)

        # Apply pure translation T = [I | t]
        # Assuming original PCDs are local (starts near 0,0,0 relative to their own start)
        pcd.translate(t)
        
        # Colorize (random color per block for visualization)
        # color = np.random.rand(3)
        # pcd.paint_uniform_color(color)
        
        merged += pcd

    # Save
    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    
    print(f"\nSaving merged preview to {out_path} ...")
    o3d.io.write_point_cloud(str(out_path), merged)
    print(f"Done. Total points: {len(merged.points)}")

if __name__ == "__main__":
    main()
