#!/usr/bin/env python3
"""
RoMa + OrangePi WiFi 图传（FrameHeader分片）+ Unity 联调脚本（Roma 工程专用）。

功能：
- 从 OrangePi WiFi 图传端口接收 JPEG 帧（协议参考 orangepizero2w/UDP_PROTOCOL.md）
- 运行 RoMa 计算相机在参考图中的像素坐标
- 将结果通过 UDP 发送给 Unity（端口保持与旧 demo 一致：12348）
- 将接收的 JPEG 帧转发一份给 Unity 用于显示（默认 12366，FrameHeader 分片协议）
- 支持 Unity 控制口（12349）：n=当前帧设为参考图；s=GPU截屏设为参考图；r=下一帧设为参考图

注意：
- 该脚本不修改/不依赖 LightGlue_Deployment，仅复用相同端口与结果协议。
"""

from __future__ import annotations

import argparse
import os
import sys
from pathlib import Path
import threading
import time
import socket

_ROMA_ROOT = Path(__file__).resolve().parents[1]
if str(_ROMA_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROMA_ROOT))

import cv2
import numpy as np
import torch

from romatch import roma_outdoor, tiny_roma_v1_outdoor
from udp_jpeg_receiver import UDPJPEGReceiver
from udp_frame_forwarder import UDPFrameForwarder
from udp_result_sender import UDPResultSender

# DXGI Desktop Duplication 截屏（可选）
try:
    import dxcam

    try:
        _DXCAM_CAMERA = dxcam.create(output_color="RGB")
        DXCAM_AVAILABLE = True
    except Exception:
        _DXCAM_CAMERA = None
        DXCAM_AVAILABLE = False
except ImportError:
    dxcam = None  # type: ignore
    _DXCAM_CAMERA = None
    DXCAM_AVAILABLE = False

torch.set_grad_enabled(False)


def maybe_resize(frame_bgr: np.ndarray, resize_opt) -> np.ndarray:
    if frame_bgr is None:
        return frame_bgr
    if len(resize_opt) == 2:
        return cv2.resize(frame_bgr, tuple(resize_opt))
    if len(resize_opt) == 1 and resize_opt[0] > 0:
        h, w = frame_bgr.shape[:2]
        scale = resize_opt[0] / max(h, w)
        return cv2.resize(frame_bgr, (int(w * scale), int(h * scale)))
    return frame_bgr


def capture_screen_bgr(opt) -> np.ndarray | None:
    if not DXCAM_AVAILABLE or _DXCAM_CAMERA is None:
        return None
    try:
        rgb = _DXCAM_CAMERA.grab()
        if rgb is None:
            return None
        if len(rgb.shape) == 3 and rgb.shape[2] == 3:
            bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
        elif len(rgb.shape) == 3 and rgb.shape[2] == 4:
            bgr = cv2.cvtColor(rgb, cv2.COLOR_BGRA2BGR)
        else:
            bgr = rgb if len(rgb.shape) == 3 else cv2.cvtColor(rgb, cv2.COLOR_GRAY2BGR)
        bgr = maybe_resize(bgr, opt.resize)
        return bgr
    except Exception:
        return None


def parse_args():
    p = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    # WiFi 图传输入
    p.add_argument("--wifi_listen_ip", type=str, default="0.0.0.0")
    p.add_argument("--wifi_listen_port", type=int, default=9000)

    # Unity 图像显示转发
    p.add_argument("--forward_ip", type=str, default="127.0.0.1")
    p.add_argument("--forward_port", type=int, default=12366)
    p.add_argument("--forward_fps", type=float, default=30.0, help="0=每帧转发")

    # Unity 结果回传
    p.add_argument("--result_ip", type=str, default="127.0.0.1")
    p.add_argument("--result_port", type=int, default=12348)

    # Unity 控制口
    p.add_argument("--control_port", type=int, default=12349)

    # RoMa 参数（继承原 demo 的主要参数）
    p.add_argument("--resize", type=int, nargs="+", default=[640, 480])
    p.add_argument("--model", type=str, default="tiny", choices=["tiny", "roma"])
    p.add_argument("--sample_num", type=int, default=2500)
    p.add_argument("--sample_thresh", type=float, default=0.05)
    p.add_argument("--ransac_reproj_threshold", type=float, default=3.0)
    p.add_argument("--min_matches", type=int, default=30)
    p.add_argument("--min_inlier_ratio", type=float, default=0.25)
    p.add_argument("--smooth_alpha", type=float, default=0.8)
    p.add_argument("--roma_interval", type=int, default=10)

    # 兼容旧 Roma demo 的常用参数（WiFi版可选使用/仅用于避免 argparse 报错）
    p.add_argument("--show_fps", action="store_true", help="在 console/画面中显示FPS（可选）")
    p.add_argument("--max_fps", type=float, default=0.0, help="限制主循环处理FPS（<=0 不限制）")
    p.add_argument("--max_display_fps", type=float, default=0.0, help="兼容参数（WiFi版不单独限显示帧率）")
    p.add_argument("--idle_sleep_ms", type=float, default=0.0, help="兼容参数（WiFi版无帧时可睡眠）")
    p.add_argument("--timer_print_interval", type=float, default=0.0, help="console 打印间隔（秒，0=不打印）")

    p.add_argument("--force_cpu", action="store_true")
    p.add_argument("--no_display", action="store_true")
    p.add_argument("--no_ui", action="store_true")

    # 启动健壮性
    p.add_argument("--first_frame_timeout", type=float, default=0.0,
                   help="等待首帧的超时秒数（<=0 表示无限等待，不因无帧退出）")

    return p.parse_args()


def start_control_listener(port: int, on_n, on_s, on_r):
    if port <= 0:
        return None, None

    stop_event = threading.Event()

    def _worker():
        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.bind(("0.0.0.0", port))
        sock.settimeout(0.5)
        print(f"[Control] Listening on 0.0.0.0:{port} (n/s/r)", flush=True)
        while not stop_event.is_set():
            try:
                try:
                    data, _addr = sock.recvfrom(1024)
                except socket.timeout:
                    continue
                if not data:
                    continue
                c = data[0]
                if c in (ord("n"), ord("N"), 1):
                    on_n()
                elif c == ord("s"):
                    on_s()
                elif c == ord("r"):
                    on_r()
            except OSError:
                break
            except Exception as exc:
                print(f"[Control] error: {exc}", flush=True)
        try:
            sock.close()
        except Exception:
            pass

    th = threading.Thread(target=_worker, daemon=True, name="RomaControlListener")
    th.start()
    return stop_event, th


def main():
    opt = parse_args()

    if opt.no_ui:
        sys.stdout = open(os.devnull, "w")
        sys.stderr = open(os.devnull, "w")

    device = "cuda" if torch.cuda.is_available() and not opt.force_cpu else "cpu"
    print(f'Running inference on device "{device}"', flush=True)

    roma_model = tiny_roma_v1_outdoor(device=torch.device(device)) if opt.model == "tiny" else roma_outdoor(device=torch.device(device))
    roma_model.sample_thresh = opt.sample_thresh

    receiver = UDPJPEGReceiver(host=opt.wifi_listen_ip, port=opt.wifi_listen_port, timeout=0.5, max_frame_buffers=8)
    receiver.start()

    forwarder = UDPFrameForwarder(target_ip=opt.forward_ip, target_port=opt.forward_port)
    result_sender = UDPResultSender(unity_ip=opt.result_ip, unity_port=opt.result_port)

    # 控制事件
    set_ref_now_event = threading.Event()
    screen_ref_event = threading.Event()
    next_frame_ref_event = threading.Event()

    def _on_n():
        set_ref_now_event.set()

    def _on_s():
        screen_ref_event.set()

    def _on_r():
        next_frame_ref_event.set()

    control_stop, control_thread = start_control_listener(opt.control_port, _on_n, _on_s, _on_r)

    ref_frame = None
    last_good_H_slow = None
    H_ref2cur_fast = None
    last_camera_center_ref = None
    frame_id = 0

    last_forward_time = 0.0
    last_timer_print_time = time.time()
    fps_ema = 0.0

    try:
        # 等第一帧（默认无限等待；避免 Unity 启动顺序/图传配置尚未生效时直接退出）
        t0 = time.time()
        last_wait_log = 0.0
        first = None
        while first is None:
            if opt.first_frame_timeout and opt.first_frame_timeout > 0 and (time.time() - t0) >= opt.first_frame_timeout:
                raise RuntimeError("No WiFi frames received within timeout. Check port/firewall/device config.")
            first = receiver.get_image(timeout=0.2)
            if first is None:
                if (time.time() - last_wait_log) >= 2.0:
                    print(
                        f"[RomaWiFi] Waiting for WiFi frames on {opt.wifi_listen_ip}:{opt.wifi_listen_port}... "
                        "Make sure OrangePi target_ip is your PC LAN IP (not 127.0.0.1), and target_port matches.",
                        flush=True,
                    )
                    last_wait_log = time.time()
                if opt.idle_sleep_ms and opt.idle_sleep_ms > 0:
                    time.sleep(opt.idle_sleep_ms / 1000.0)

        first = maybe_resize(first, opt.resize)
        ref_frame = first.copy()
        print("[RomaWiFi] First frame received, using as reference", flush=True)

        while True:
            loop_start_time = time.time()
            frame_id += 1

            # s：截屏设为参考图
            if screen_ref_event.is_set():
                screen_ref_event.clear()
                sc = capture_screen_bgr(opt)
                if sc is not None:
                    ref_frame = sc.copy()
                    last_good_H_slow = None
                    H_ref2cur_fast = None
                    last_camera_center_ref = None
                    print("[Control] Reference updated from screen capture", flush=True)
                elif not DXCAM_AVAILABLE:
                    print("[Control] Screen capture skipped: install dxcam (pip install dxcam)", flush=True)

            frame = receiver.get_image(timeout=0.2)
            if frame is None:
                if opt.idle_sleep_ms and opt.idle_sleep_ms > 0:
                    time.sleep(opt.idle_sleep_ms / 1000.0)
                continue
            frame = maybe_resize(frame, opt.resize)

            # r：下一帧设为参考图（Roma 版先用 WiFi 下一帧作为参考图）
            if next_frame_ref_event.is_set():
                next_frame_ref_event.clear()
                ref_frame = frame.copy()
                last_good_H_slow = None
                H_ref2cur_fast = None
                last_camera_center_ref = None
                print("[Control] Reference updated from next WiFi frame (r)", flush=True)
                continue

            # n：当前帧设为参考图
            if set_ref_now_event.is_set():
                set_ref_now_event.clear()
                ref_frame = frame.copy()
                last_good_H_slow = None
                H_ref2cur_fast = None
                last_camera_center_ref = None
                print("[Control] Reference updated from current frame (n)", flush=True)

            # 转发图像给 Unity（尽量转发 JPEG 原始字节；如果拿不到则回退为编码）
            now = time.time()
            do_forward = opt.forward_fps <= 0 or (now - last_forward_time) >= (1.0 / max(opt.forward_fps, 1e-6))
            if do_forward:
                jpeg_bytes = receiver.get_jpeg(timeout=0.0)
                if jpeg_bytes is None:
                    ok, enc = cv2.imencode(".jpg", frame, [int(cv2.IMWRITE_JPEG_QUALITY), 80])
                    jpeg_bytes = enc.tobytes() if ok else None
                if jpeg_bytes:
                    forwarder.send_jpeg(jpeg_bytes)
                    last_forward_time = now

            # ---- RoMa 计算 H_ref_to_cur ----
            h_ref, w_ref = ref_frame.shape[:2]
            h_cur, w_cur = frame.shape[:2]
            camera_center_current = (w_cur // 2, h_cur // 2)

            H_for_pose = None
            inliers_ratio = 0.0
            num_matches = 0

            run_roma_this_frame = (frame_id % opt.roma_interval == 0) or (last_good_H_slow is None)
            if run_roma_this_frame:
                from PIL import Image

                ref_pil = Image.fromarray(cv2.cvtColor(ref_frame, cv2.COLOR_BGR2RGB))
                cur_pil = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))

                warp, certainty = roma_model.match(ref_pil, cur_pil) if opt.model == "tiny" else roma_model.match(ref_pil, cur_pil, device=torch.device(device))
                matches, _ = roma_model.sample(warp, certainty, num=opt.sample_num)

                k_ref, k_cur = roma_model.to_pixel_coordinates(matches, h_ref, w_ref, h_cur, w_cur)
                pts_ref = k_ref.detach().cpu().numpy().astype(np.float32)
                pts_cur = k_cur.detach().cpu().numpy().astype(np.float32)
                num_matches = len(pts_ref)

                if num_matches >= opt.min_matches:
                    H_tmp, mask = cv2.findHomography(pts_ref, pts_cur, cv2.RANSAC, opt.ransac_reproj_threshold)
                    if H_tmp is not None and mask is not None:
                        inliers_ratio = float(mask.mean())
                        if inliers_ratio >= opt.min_inlier_ratio:
                            last_good_H_slow = H_tmp
                            H_ref2cur_fast = H_tmp.copy()

            H_for_pose = H_ref2cur_fast if H_ref2cur_fast is not None else last_good_H_slow

            camera_center_ref = None
            is_valid = False
            if H_for_pose is not None and num_matches >= opt.min_matches:
                try:
                    H_cur_to_ref = np.linalg.inv(H_for_pose)
                    center_cur = np.array([[camera_center_current]], dtype=np.float32)
                    center_ref_now = cv2.perspectiveTransform(center_cur, H_cur_to_ref)[0, 0]
                    center_ref_now = np.clip(center_ref_now, [0, 0], [w_ref - 1, h_ref - 1])
                    if last_camera_center_ref is None:
                        camera_center_ref = center_ref_now
                    else:
                        camera_center_ref = opt.smooth_alpha * last_camera_center_ref + (1.0 - opt.smooth_alpha) * center_ref_now
                    last_camera_center_ref = camera_center_ref
                    is_valid = True
                except np.linalg.LinAlgError:
                    is_valid = False

            # 结果回传给 Unity
            if is_valid and camera_center_ref is not None:
                result_sender.send_result(True, num_matches, inliers_ratio, float(camera_center_ref[0]), float(camera_center_ref[1]))
            else:
                result_sender.send_result(False, num_matches, 0.0, 0.0, 0.0)

            if not opt.no_display:
                disp = frame.copy()
                cv2.putText(disp, f"matches={num_matches} inliers={inliers_ratio:.1%}", (10, 30),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0) if is_valid else (0, 0, 255), 2)
                if opt.show_fps:
                    cv2.putText(disp, f"FPS: {fps_ema:.1f}", (10, 60),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                cv2.imshow("RoMa WiFi", disp)
                key = cv2.waitKey(1) & 0xFF
                if key == ord("q"):
                    break
                if key == ord("n"):
                    set_ref_now_event.set()
                if key == ord("s"):
                    screen_ref_event.set()

            if opt.timer_print_interval and opt.timer_print_interval > 0:
                now = time.time()
                if (now - last_timer_print_time) >= opt.timer_print_interval:
                    print(f"[RomaWiFi] FPS={fps_ema:.1f} matches={num_matches} inliers={inliers_ratio:.1%}", flush=True)
                    last_timer_print_time = now

            # 限速（把 sleep 计入 loop_elapsed，避免出现“几十万FPS”的假值）
            if opt.max_fps and opt.max_fps > 0:
                loop_elapsed = time.time() - loop_start_time
                target_dt = 1.0 / float(opt.max_fps)
                if loop_elapsed < target_dt:
                    time.sleep(target_dt - loop_elapsed)
            loop_elapsed = time.time() - loop_start_time
            dt = max(loop_elapsed, 1e-6)
            fps_ema = 0.9 * fps_ema + 0.1 * (1.0 / dt)


    finally:
        try:
            receiver.stop()
        except Exception:
            pass
        try:
            forwarder.close()
        except Exception:
            pass
        try:
            result_sender.close()
        except Exception:
            pass
        if control_stop is not None:
            control_stop.set()
        try:
            cv2.destroyAllWindows()
        except Exception:
            pass


if __name__ == "__main__":
    main()

