#!/usr/bin/env python3
"""
Orbbec D2C Batch Converter (Software)

Usage:
  python d2c_batch.py            # Interactive: select profiles, fetch params, convert
  python d2c_batch.py --params   # Only fetch/update camera params (no conversion)
  python d2c_batch.py --convert  # Only convert using saved camera_params.json

Workflow:
  1. Select RGB profile (resolution / fps / format)
  2. Select Depth profile (resolution / fps / format)
  3. Camera intrinsics + extrinsics are fetched and saved to camera_params.json
  4. Enter depth image directory
  5. Batch-convert all PNG depth images via software D2C
     -> Output saved to <input_dir>_d2c/
"""

import os
import sys
import json
import time
import numpy as np
import cv2
from pathlib import Path

try:
    from PIL import Image as _PIL_Image
    _HAS_PIL = True
except ImportError:
    _HAS_PIL = False

PARAMS_FILE = "camera_params.json"


# ---------------------------------------------------------------------------
# Camera parameter acquisition
# ---------------------------------------------------------------------------

def _list_video_profiles(profile_list):
    """Return list of (index, VideoStreamProfile) for all video profiles."""
    results = []
    count = profile_list.get_count()
    for i in range(count):
        p = profile_list.get_stream_profile_by_index(i)
        vp = p.as_video_stream_profile()
        if vp is None:
            continue
        results.append((i, vp))
    return results


def _print_profiles(profiles, sensor_name):
    print(f"\nAvailable {sensor_name} profiles:")
    print(f"  {'#':<4} {'Resolution':<14} {'FPS':<6} {'Format'}")
    print(f"  {'-'*40}")
    for idx, vp in profiles:
        fmt = str(vp.get_format()).split(".")[-1]
        print(f"  [{idx:<2}] {vp.get_width()}x{vp.get_height():<8}  {vp.get_fps():<6} {fmt}")


def _select_profile(profile_list, sensor_name):
    """Interactive profile selection. Returns the selected VideoStreamProfile."""
    profiles = _list_video_profiles(profile_list)
    if not profiles:
        print(f"  No {sensor_name} profiles found!")
        return None

    _print_profiles(profiles, sensor_name)
    indices = [i for i, _ in profiles]

    while True:
        raw = input(f"Select {sensor_name} profile index: ").strip()
        try:
            choice = int(raw)
            if choice in indices:
                selected = next(vp for i, vp in profiles if i == choice)
                fmt = str(selected.get_format()).split(".")[-1]
                print(f"  -> Selected: {selected.get_width()}x{selected.get_height()} @ {selected.get_fps()}fps  {fmt}")
                return selected
        except ValueError:
            pass
        print(f"  Invalid. Choose from: {indices}")


def _intrinsic_to_dict(intr):
    return {
        "fx": float(intr.fx), "fy": float(intr.fy),
        "cx": float(intr.cx), "cy": float(intr.cy),
        "width": int(intr.width), "height": int(intr.height),
    }


def _distortion_to_dict(dist):
    return {
        "k1": float(dist.k1), "k2": float(dist.k2),
        "k3": float(dist.k3), "k4": float(dist.k4),
        "k5": float(dist.k5), "k6": float(dist.k6),
        "p1": float(dist.p1), "p2": float(dist.p2),
    }


def fetch_camera_params():
    """
    Interactively select color + depth profiles, start the pipeline briefly to
    capture one frameset, extract intrinsics/extrinsics, and return as a dict.
    """
    from pyorbbecsdk import Pipeline, Config, OBSensorType

    print("\n=== Phase 1: Camera Parameter Acquisition ===")
    pipeline = Pipeline()
    config = Config()

    # --- Color profile ---
    try:
        color_list = pipeline.get_stream_profile_list(OBSensorType.COLOR_SENSOR)
    except Exception as e:
        print(f"  Cannot get color profiles: {e}")
        return None

    color_profile = _select_profile(color_list, "RGB (Color)")
    if color_profile is None:
        return None

    # --- Depth profile ---
    try:
        depth_list = pipeline.get_stream_profile_list(OBSensorType.DEPTH_SENSOR)
    except Exception as e:
        print(f"  Cannot get depth profiles: {e}")
        return None

    depth_profile = _select_profile(depth_list, "Depth")
    if depth_profile is None:
        return None

    config.enable_stream(color_profile)
    config.enable_stream(depth_profile)

    print("\n  Starting pipeline...")
    try:
        pipeline.start(config)
    except Exception as e:
        print(f"  Failed to start pipeline: {e}")
        return None

    # --- Wait for first valid frameset ---
    frames = None
    deadline = time.time() + 10.0
    while time.time() < deadline:
        f = pipeline.wait_for_frames(200)
        if f and f.get_color_frame() and f.get_depth_frame():
            frames = f
            break

    if frames is None:
        print("  Timeout: no frameset received within 10 s.")
        pipeline.stop()
        return None

    color_frame = frames.get_color_frame()
    depth_frame = frames.get_depth_frame()

    # --- Extract profile-specific intrinsics from actual frames ---
    color_vp = color_frame.get_stream_profile().as_video_stream_profile()
    depth_vp = depth_frame.get_stream_profile().as_video_stream_profile()

    color_intr = color_vp.get_intrinsic()
    color_dist = color_vp.get_distortion()
    depth_intr = depth_vp.get_intrinsic()
    depth_dist = depth_vp.get_distortion()

    # Extrinsic: depth camera -> color camera
    extrinsic = depth_vp.get_extrinsic_to(color_vp)

    depth_scale = depth_frame.get_depth_scale()  # mm per raw unit

    pipeline.stop()
    print("  Pipeline stopped.")

    params = {
        "color": {
            "width":  color_frame.get_width(),
            "height": color_frame.get_height(),
            "fps":    color_profile.get_fps(),
            "format": str(color_profile.get_format()).split(".")[-1],
            "intrinsic":  _intrinsic_to_dict(color_intr),
            "distortion": _distortion_to_dict(color_dist),
        },
        "depth": {
            "width":  depth_frame.get_width(),
            "height": depth_frame.get_height(),
            "fps":    depth_profile.get_fps(),
            "format": str(depth_profile.get_format()).split(".")[-1],
            "scale":  float(depth_scale),
            "intrinsic":  _intrinsic_to_dict(depth_intr),
            "distortion": _distortion_to_dict(depth_dist),
        },
        # rot: 9-element flat array (row-major 3x3), transform: 3-element translation (mm)
        "extrinsic": {
            "rot":       extrinsic.rot.tolist(),
            "transform": extrinsic.transform.tolist(),
        },
    }
    return params


# ---------------------------------------------------------------------------
# Software D2C transformation
# ---------------------------------------------------------------------------

def _undistort_points(u, v, intr, dist):
    """Brown-Conrady lens undistortion (returns undistorted normalised coords)."""
    fx, fy = intr["fx"], intr["fy"]
    cx, cy = intr["cx"], intr["cy"]
    k1, k2, k3 = dist["k1"], dist["k2"], dist["k3"]
    p1, p2 = dist["p1"], dist["p2"]

    x = (u - cx) / fx
    y = (v - cy) / fy

    r2 = x * x + y * y
    radial = 1.0 + k1 * r2 + k2 * r2**2 + k3 * r2**3
    x_u = x * radial + 2.0 * p1 * x * y + p2 * (r2 + 2.0 * x * x)
    y_u = y * radial + p1 * (r2 + 2.0 * y * y) + 2.0 * p2 * x * y
    return x_u, y_u


def software_d2c(depth_img, params, splat=True):
    """
    Transform a uint16 depth image (depth camera space) to an aligned uint16
    depth image in color camera space.

    Parameters
    ----------
    depth_img : np.ndarray  (H_d x W_d, uint16)  raw depth in sensor units
    params    : dict  camera_params.json content

    Returns
    -------
    aligned : np.ndarray  (H_c x W_c, uint16)  aligned depth in sensor units
    """
    depth_intr = params["depth"]["intrinsic"]
    depth_dist = params["depth"]["distortion"]
    color_intr = params["color"]["intrinsic"]
    ext = params["extrinsic"]

    color_w = color_intr["width"]
    color_h = color_intr["height"]
    fx_c, fy_c = color_intr["fx"], color_intr["fy"]
    cx_c, cy_c = color_intr["cx"], color_intr["cy"]

    dh, dw = depth_img.shape
    depth_scale = params["depth"]["scale"]

    # Pixel grid
    u_d = np.arange(dw, dtype=np.float32)
    v_d = np.arange(dh, dtype=np.float32)
    u_d, v_d = np.meshgrid(u_d, v_d)  # (dh, dw)

    # Depth in mm (float)
    Z = depth_img.astype(np.float32) * depth_scale
    valid = Z > 0

    # Undistort + unproject to 3-D (depth camera space, mm)
    x_u, y_u = _undistort_points(u_d, v_d, depth_intr, depth_dist)
    X = x_u * Z   # (dh, dw)
    Y = y_u * Z

    # Extrinsic: R (3x3) and t (3,) in mm
    R = np.asarray(ext["rot"], dtype=np.float64).reshape(3, 3)
    t = np.asarray(ext["transform"], dtype=np.float64)

    # Vectorised transform  [X_c, Y_c, Z_c] = R @ [X, Y, Z]^T + t
    pts = np.stack([X.ravel(), Y.ravel(), Z.ravel()], axis=0).astype(np.float64)  # (3, N)
    pts_c = R @ pts + t[:, np.newaxis]                                              # (3, N)

    Xc = pts_c[0].reshape(dh, dw).astype(np.float32)
    Yc = pts_c[1].reshape(dh, dw).astype(np.float32)
    Zc = pts_c[2].reshape(dh, dw).astype(np.float32)

    # Project onto color image plane
    valid_c = valid & (Zc > 0)
    u_c = np.where(valid_c, fx_c * Xc / Zc + cx_c, -1.0).astype(np.float32)
    v_c = np.where(valid_c, fy_c * Yc / Zc + cy_c, -1.0).astype(np.float32)

    u_ci = np.round(u_c).astype(np.int32)
    v_ci = np.round(v_c).astype(np.int32)

    in_bounds = (u_ci >= 0) & (u_ci < color_w) & (v_ci >= 0) & (v_ci < color_h)
    mask = valid_c & in_bounds

    u_vals = u_ci[mask]
    v_vals = v_ci[mask]
    z_vals = depth_img[mask]   # keep original uint16 sensor units

    # Write far pixels first so nearer pixels overwrite (z small = close)
    order = np.argsort(z_vals)[::-1]
    u_vals = u_vals[order]
    v_vals = v_vals[order]
    z_vals = z_vals[order]

    aligned = np.zeros((color_h, color_w), dtype=np.uint16)

    if splat:
        # Splat each depth pixel to a 3×3 neighbourhood to fill sub-pixel gaps
        # that arise from forward projection.  Far→near order is preserved so
        # the nearest surface always wins when splats overlap.
        for dv in range(-1, 2):
            for du in range(-1, 2):
                us = np.clip(u_vals + du, 0, color_w - 1)
                vs = np.clip(v_vals + dv, 0, color_h - 1)
                aligned[vs, us] = z_vals
    else:
        aligned[v_vals, u_vals] = z_vals

    return aligned


# ---------------------------------------------------------------------------
# Batch conversion
# ---------------------------------------------------------------------------

def depth_to_colormap(aligned, min_depth_mm=200, max_depth_mm=5000, depth_scale=1.0):
    """Convert uint16 aligned depth to a JET pseudo-color BGR image."""
    depth_mm = aligned.astype(np.float32) * depth_scale
    valid = (depth_mm > min_depth_mm) & (depth_mm < max_depth_mm)
    norm = np.zeros_like(depth_mm)
    norm[valid] = (depth_mm[valid] - min_depth_mm) / (max_depth_mm - min_depth_mm)
    norm = np.clip(norm, 0.0, 1.0)
    gray = (norm * 255).astype(np.uint8)
    colormap = cv2.applyColorMap(gray, cv2.COLORMAP_JET)
    colormap[~valid] = 0  # black for invalid pixels
    return colormap


def batch_convert(params, depth_dir_str, splat=True):
    depth_dir = Path(depth_dir_str.strip().strip('"\''))
    if not depth_dir.exists():
        print(f"  Directory not found: {depth_dir}")
        return

    png_files = sorted(depth_dir.glob("*.png")) + sorted(depth_dir.glob("*.PNG"))
    png_files = sorted(set(png_files))
    if not png_files:
        print(f"  No PNG files found in: {depth_dir}")
        return

    out_dir = depth_dir.parent / (depth_dir.name + "_d2c")
    vis_dir = depth_dir.parent / (depth_dir.name + "_d2c_vis")
    out_dir.mkdir(exist_ok=True)
    vis_dir.mkdir(exist_ok=True)

    color_w = params["color"]["intrinsic"]["width"]
    color_h = params["color"]["intrinsic"]["height"]
    depth_scale = params["depth"]["scale"]

    print(f"\n=== Phase 2: Batch D2C Conversion ===")
    print(f"  Input dir    : {depth_dir}")
    print(f"  Aligned depth: {out_dir}")
    print(f"  Pseudo-color : {vis_dir}")
    print(f"  Files        : {len(png_files)}")
    print(f"  Depth scale  : {depth_scale:.6f} mm/unit")
    print(f"  Output size  : {color_w}x{color_h}  (color resolution)")
    print(f"  Splat (3×3)  : {'on' if splat else 'off'}")
    print()

    t0 = time.time()
    ok = 0
    for i, fpath in enumerate(png_files, 1):
        depth_img = None
        if _HAS_PIL:
            try:
                depth_img = np.array(_PIL_Image.open(str(fpath)))
            except Exception:
                depth_img = None
        if depth_img is None:
            depth_img = cv2.imread(str(fpath), cv2.IMREAD_UNCHANGED)
        if depth_img is None:
            print(f"  [{i:>4}/{len(png_files)}] SKIP (unreadable): {fpath.name}")
            continue
        if depth_img.ndim != 2:
            depth_img = depth_img[:, :, 0]
        depth_img = depth_img.astype(np.uint16)

        aligned = software_d2c(depth_img, params, splat=splat)

        # Save aligned uint16 depth
        cv2.imwrite(str(out_dir / fpath.name), aligned)

        # Save pseudo-color visualization
        colormap = depth_to_colormap(aligned, depth_scale=depth_scale)
        vis_name = fpath.stem + "_vis.png"
        cv2.imwrite(str(vis_dir / vis_name), colormap)

        ok += 1
        print(f"  [{i:>4}/{len(png_files)}] {fpath.name}")

    elapsed = time.time() - t0
    print(f"\n  Done. {ok}/{len(png_files)} files converted in {elapsed:.1f}s")
    print(f"  Aligned depth : {out_dir}")
    print(f"  Pseudo-color  : {vis_dir}")


# ---------------------------------------------------------------------------
# Parameter display
# ---------------------------------------------------------------------------

def print_params_summary(params):
    c = params["color"]
    d = params["depth"]
    ci = c["intrinsic"]
    di = d["intrinsic"]
    ext = params["extrinsic"]
    R = np.array(ext["rot"]).reshape(3, 3)
    t = np.array(ext["transform"])

    print("\n--- Camera Parameters Summary ---")
    print(f"  RGB   : {c['width']}x{c['height']} @ {c['fps']}fps  {c['format']}")
    print(f"          fx={ci['fx']:.4f}  fy={ci['fy']:.4f}  cx={ci['cx']:.4f}  cy={ci['cy']:.4f}")
    print(f"  Depth : {d['width']}x{d['height']} @ {d['fps']}fps  {d['format']}")
    print(f"          fx={di['fx']:.4f}  fy={di['fy']:.4f}  cx={di['cx']:.4f}  cy={di['cy']:.4f}")
    print(f"          scale={d['scale']:.6f} mm/unit")
    print(f"  Extrinsic R :\n    {R[0]}\n    {R[1]}\n    {R[2]}")
    print(f"  Extrinsic t : {t}  (mm)")


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main():
    only_params = "--params" in sys.argv
    only_convert = "--convert" in sys.argv
    no_splat = "--no-splat" in sys.argv

    params = None

    # ---- Load or fetch camera parameters ----
    if not only_convert:
        if os.path.exists(PARAMS_FILE) and not only_params:
            print(f"Found saved camera params: {PARAMS_FILE}")
            ans = input("Use existing params? [y=use existing / n=re-fetch from camera]: ").strip().lower()
            if ans == "y":
                with open(PARAMS_FILE) as f:
                    params = json.load(f)
                print("  Loaded existing params.")

        if params is None:
            params = fetch_camera_params()
            if params is None:
                print("Failed to fetch camera parameters. Exiting.")
                sys.exit(1)
            with open(PARAMS_FILE, "w") as f:
                json.dump(params, f, indent=2)
            print(f"\n  Camera params saved to: {PARAMS_FILE}")

    else:
        # --convert: must have existing params file
        if not os.path.exists(PARAMS_FILE):
            print(f"No {PARAMS_FILE} found. Run without --convert first to fetch camera params.")
            sys.exit(1)
        with open(PARAMS_FILE) as f:
            params = json.load(f)
        print(f"Loaded camera params from: {PARAMS_FILE}")

    print_params_summary(params)

    if only_params:
        print("\n--params mode: done.")
        return

    # ---- Batch conversion ----
    depth_dir = input("\nEnter depth images directory: ").strip()
    if not depth_dir:
        print("No directory entered. Exiting.")
        return

    batch_convert(params, depth_dir, splat=not no_splat)


if __name__ == "__main__":
    main()