"""Compare Intel's align on PC2 with an align on the host, on the SAME frameset. Read-only.

    python scripts/compare_host_align.py --url http://192.168.123.99:8878 [--framesets 10] [--tries 20]

For each frameset: GET /api/raw/latest (color + depth, not aligned) and
GET /api/capture (depth aligned on PC2) until both carry the same
frameset_seq; align the raw depth on the host with pyrealsense2 (a
software_device built from /api/intrinsics, then rs.align); report, pixel by
pixel, how the two aligned depths differ. Prints one JSON line per frameset
and a summary. Only GET requests; nothing on PC2 changes.

The host's pyrealsense2 may be newer than PC2's (2.58 here, 2.50 there): the
two aligns are then not even the same code version.
"""

import argparse
import base64
import json
import sys

import cv2
import httpx
import numpy as np
import pyrealsense2 as rs

SYNC_WAIT_MS = 200     # wait for the syncer's frameset after a push
SYNC_PUSHES_MAX = 5    # the syncer may emit the first pair split (color alone): push the same pair again


# Calibration ------------------------------------------------------------------

def rs_intrinsics(stream: dict) -> rs.intrinsics:
    """rs.intrinsics from one stream of /api/intrinsics ("ir" or "rgb")."""
    d = stream["distortion"] or {"k1": 0, "k2": 0, "k3": 0, "p1": 0, "p2": 0}
    intr = rs.intrinsics()
    intr.width, intr.height = stream["width"], stream["height"]
    intr.fx, intr.fy, intr.ppx, intr.ppy = stream["fx"], stream["fy"], stream["cx"], stream["cy"]
    intr.model = getattr(rs.distortion, stream["model"] or "none")
    intr.coeffs = [d["k1"], d["k2"], d["p1"], d["p2"], d["k3"]]   # the SDK's order
    return intr


def rs_extrinsics(ext: dict) -> rs.extrinsics:
    """rs.extrinsics depth->color from the API's (rotation row-major, translation mm)."""
    r = ext["rotation"]
    out = rs.extrinsics()
    out.rotation = [r[0], r[3], r[6], r[1], r[4], r[7], r[2], r[5], r[8]]   # back to column-major
    out.translation = [t / 1000.0 for t in ext["translation"]]
    return out


# Host align ---------------------------------------------------------------------

class HostAligner:
    """rs.align(color) on frames pushed into a software_device with PC2's calibration."""

    def __init__(self, calibration: dict):
        self._device = rs.software_device()
        self._depth = self._device.add_sensor("Depth")
        self._color = self._device.add_sensor("Color")
        self._depth.add_read_only_option(rs.option.depth_units, 0.001)   # we push millimeters
        self._depth_profile = self._depth.add_video_stream(
            _video_stream(rs.stream.depth, 0, rs.format.z16, 2, rs_intrinsics(calibration["ir"])))
        self._color_profile = self._color.add_video_stream(
            _video_stream(rs.stream.color, 1, rs.format.bgr8, 3, rs_intrinsics(calibration["rgb"])))
        self._depth_profile.register_extrinsics_to(self._color_profile, rs_extrinsics(calibration["extrinsics_ir_to_rgb"]))
        self._sync = rs.syncer()
        self._align = rs.align(rs.stream.color)
        self._number = 0
        self._started = []
        for sensor, profile in ((self._depth, self._depth_profile), (self._color, self._color_profile)):
            sensor.open(profile)
            sensor.start(self._sync)
            self._started.append(sensor)

    def close(self) -> None:
        for sensor in self._started:
            sensor.stop()
            sensor.close()
        self._started = []

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.close()

    def align(self, depth_mm: np.ndarray, color_bgr: np.ndarray) -> np.ndarray:
        """Depth (mm, depth grid) aligned onto the color grid by rs.align, as uint16 mm."""
        depth_mm, color_bgr = np.ascontiguousarray(depth_mm), np.ascontiguousarray(color_bgr)
        for _ in range(SYNC_PUSHES_MAX):
            self._number += 1
            stamp_ms = self._number * 1000.0
            _push(self._depth, self._depth_profile, depth_mm, 2, self._number, stamp_ms, depth=True)
            _push(self._color, self._color_profile, color_bgr, 3, self._number, stamp_ms)
            frames = self._complete_frameset(self._number)
            if frames is not None:
                aligned = self._align.process(frames)
                return np.asanyarray(aligned.get_depth_frame().get_data()).copy()
        raise RuntimeError(f"syncer gave no complete frameset in {SYNC_PUSHES_MAX} pushes")

    def _complete_frameset(self, number: int):
        """The syncer's frameset holding depth AND color of frame `number`, or None."""
        for _ in range(2):   # a split pair comes out as two framesets
            try:
                frames = self._sync.wait_for_frames(SYNC_WAIT_MS)
            except RuntimeError:   # "Frame did not arrive in time": try pushing again
                return None
            depth = frames.get_depth_frame()
            if frames.size() == 2 and depth and depth.get_frame_number() == number:
                return frames
        return None


def _video_stream(kind, uid: int, fmt, bpp: int, intrinsics: rs.intrinsics) -> rs.video_stream:
    vs = rs.video_stream()
    vs.type, vs.index, vs.uid, vs.fmt, vs.bpp = kind, 0, uid, fmt, bpp
    vs.width, vs.height, vs.fps, vs.intrinsics = intrinsics.width, intrinsics.height, 30, intrinsics
    return vs


def _push(sensor, profile, pixels: np.ndarray, bpp: int, number: int, stamp_ms: float, depth: bool = False) -> None:
    frame = rs.software_video_frame()
    frame.pixels = pixels
    frame.bpp = bpp
    frame.stride = pixels.shape[1] * bpp
    frame.timestamp = stamp_ms
    frame.domain = rs.timestamp_domain.hardware_clock
    frame.frame_number = number
    frame.profile = profile.as_video_stream_profile()
    if depth:
        frame.depth_units = 0.001
    sensor.on_video_frame(frame)


# PC2 side (GET only) ------------------------------------------------------------

def _data(response: httpx.Response) -> dict:
    response.raise_for_status()
    body = response.json()
    if not body.get("success"):
        raise RuntimeError(f"{response.request.url}: {body.get('error')}")
    return body["data"]


def _decode(b64: str, flags: int) -> np.ndarray:
    return cv2.imdecode(np.frombuffer(base64.b64decode(b64), np.uint8), flags)


def fetch_same_frameset(client: httpx.Client, tries: int) -> tuple[dict, dict, int]:
    """(raw, capture, attempts): /api/raw/latest and /api/capture of one frameset_seq."""
    for attempt in range(1, tries + 1):
        raw = _data(client.get("/api/raw/latest"))
        capture = _data(client.get("/api/capture"))
        if capture.get("frameset_seq") == raw["frameset_seq"]:
            return raw, capture, attempt
    raise RuntimeError(f"no matching frameset in {tries} attempts (another client may keep the align busy)")


# Comparison (pure) -------------------------------------------------------------

def compare_depths(pc2: np.ndarray, host: np.ndarray) -> dict:
    """Pixel-by-pixel difference of two aligned depth maps in mm (same shape)."""
    if pc2.shape != host.shape:
        raise ValueError(f"shapes differ: PC2 {pc2.shape}, host {host.shape}")
    both = (pc2 > 0) & (host > 0)
    diff = np.abs(pc2[both].astype(np.int32) - host[both].astype(np.int32))
    total = pc2.size
    return {
        "pixels": int(total),
        "valid_both": int(both.sum()),
        "only_pc2": int(((pc2 > 0) & (host == 0)).sum()),
        "only_host": int(((host > 0) & (pc2 == 0)).sum()),
        "equal_fraction_of_valid_both": float((diff == 0).mean()) if diff.size else None,
        "abs_diff_mm": None if not diff.size else {
            "mean": float(diff.mean()), "median": float(np.median(diff)),
            "p99": float(np.percentile(diff, 99)), "max": int(diff.max()),
        },
    }


def compare_one(client: httpx.Client, aligner: HostAligner, tries: int) -> dict:
    raw, capture, attempts = fetch_same_frameset(client, tries)
    host = aligner.align(_decode(raw["depth_png_b64"], cv2.IMREAD_UNCHANGED),
                         _decode(raw["color_jpeg_b64"], cv2.IMREAD_COLOR))
    pc2 = _decode(capture["depth_png_b64"], cv2.IMREAD_UNCHANGED)
    return {"frameset_seq": raw["frameset_seq"], "attempts": attempts, **compare_depths(pc2, host)}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--url", required=True, help="RealSense service, e.g. http://192.168.123.99:8878")
    ap.add_argument("--framesets", type=int, default=10)
    ap.add_argument("--tries", type=int, default=20, help="attempts to get raw and aligned of one frameset")
    ap.add_argument("--api-key", default="")
    args = ap.parse_args(argv)
    headers = {"X-Api-Key": args.api_key} if args.api_key else {}
    results = []
    with httpx.Client(base_url=args.url, headers=headers, timeout=10) as client:
        calibration = _data(client.get("/api/intrinsics"))
        with HostAligner(calibration) as aligner:
            for _ in range(args.framesets):
                result = compare_one(client, aligner, args.tries)
                results.append(result)
                print(json.dumps(result))
    means = [r["abs_diff_mm"]["mean"] for r in results if r["abs_diff_mm"]]
    print(json.dumps({"summary": {"framesets": len(results), "pyrealsense2_host": rs.__version__,
                                  "mean_abs_diff_mm_per_frameset": means}}))
    return 0


if __name__ == "__main__":
    sys.exit(main())