"""Example: Convert a single image to a metric point cloud with UniDepth.

Usage:
    python image_to_pointcloud.py room.jpg --output room.ply --checkpoint lpiccinelli/unidepth-v2-vits14
"""

import argparse
from pathlib import Path

from PIL import Image
from unidepth.inference import UniDepth, save_pointcloud_ply


def main():
    parser = argparse.ArgumentParser(description="Image → metric depth → point cloud")
    parser.add_argument("image", type=str, help="Input image path")
    parser.add_argument("--output", "-o", type=str, default="output.ply", help="Output PLY file")
    parser.add_argument("--checkpoint", type=str, default="lpiccinelli/unidepth-v2-vits14",
                        help="HF checkpoint: lpiccinelli/unidepth-v2-vits14 or lpiccinelli/unidepth-v2-vitl14")
    parser.add_argument("--device", type=str, default="cuda", choices=["cuda", "cpu"])
    parser.add_argument("--confidence-threshold", type=float, default=None,
                        help="Filter points with confidence > threshold (V2 only)")
    args = parser.parse_args()

    # ------------------------------------------------------------------ #
    # 1. Load image
    # ------------------------------------------------------------------ #
    image = Image.open(args.image).convert("RGB")
    print(f"Loaded image: {image.size[0]}×{image.size[1]}")

    # ------------------------------------------------------------------ #
    # 2. Load model (downloads weights on first run)
    # ------------------------------------------------------------------ #
    model = UniDepth.from_pretrained(args.checkpoint, device=args.device)

    # ------------------------------------------------------------------ #
    # 3. Inference
    # ------------------------------------------------------------------ #
    results = model(image, confidence_threshold=args.confidence_threshold)

    depth = results["depth"]
    points = results["points"]
    colors = results["colors"]
    intrinsics = results["intrinsics"]
    confidence = results["confidence"]

    print(f"Depth range: [{depth.min():.3f}, {depth.max():.3f}] meters")
    print(f"Predicted intrinsics K:\n{intrinsics}")
    if confidence is not None:
        print(f"Confidence range: [{confidence.min():.3f}, {confidence.max():.3f}]")
    print(f"Generated {len(points)} 3D points")

    # ------------------------------------------------------------------ #
    # 4. Save outputs
    # ------------------------------------------------------------------ #
    out_path = Path(args.output)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    # Save point cloud
    save_pointcloud_ply(str(out_path), points, colors)
    print(f"Saved point cloud to {out_path}")

    # Optionally save depth map as 16-bit PNG (mm precision)
    depth_png = out_path.with_suffix(".depth.png")
    import numpy as np
    from PIL import Image as PILImage
    depth_mm = (depth * 1000).astype(np.uint16)
    PILImage.fromarray(depth_mm).save(depth_png)
    print(f"Saved depth map to {depth_png}")


if __name__ == "__main__":
    main()