#!/usr/bin/env python3
"""
UV-unroll finger tactile pipeline with 3D mesh visualization.

Offline: load Blender export -> build_or_load_domain -> unroll_map_uv.npz
Runtime:  camera frame -> UVUnrollFrameProcessor -> TactileProcessor -> forces/contact
          UV-space flow/HHD -> grey-white 3D finger mesh (Open3D)

For the full 2D dashboard demo see:
  tools/3d_patch_generation_blender/unroll_realtime.py
"""

import argparse
import os
import sys
from pathlib import Path

os.environ.setdefault("OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS", "0")

import cv2
from orisys.finger import DEFAULT_FINGER_EXPORT_DIR, DEFAULT_FINGER_MAP_CACHE, FingerRuntime

_EXAMPLES_DIR = Path(__file__).resolve().parent
if str(_EXAMPLES_DIR) not in sys.path:
    sys.path.insert(0, str(_EXAMPLES_DIR))

from viz import UVMeshViewer, build_uv_surface_mesh, show_finger_uv_panels


def open_mesh_viewer(runtime, args):
    domain = runtime.domain
    if domain.surface is None:
        raise RuntimeError(
            "UV domain has no world_xyz surface map. Rebuild cache with include_surface=True."
        )

    uv_mesh = build_uv_surface_mesh(domain.surface.world_xyz, domain.valid_mask)
    if args.no_3d:
        return None, uv_mesh

    viewer = UVMeshViewer(
        uv_mesh,
        max_magnitude=args.magnitude_max,
        threshold=args.depth_threshold if args.mesh_map == "depth" else args.flow_threshold,
        color_mode=args.mesh_map,
    )
    print(
        f"3D mesh viewer: {len(uv_mesh.rest_vertices)} vertices, "
        f"{len(uv_mesh.triangles)} triangles, {args.mesh_map} surface mode"
    )
    return viewer, uv_mesh


def main():
    parser = argparse.ArgumentParser(
        description="UV-unroll finger tactile pipeline with 3D mesh viewer"
    )
    parser.add_argument(
        "--export-dir",
        default=DEFAULT_FINGER_EXPORT_DIR,
        help="Blender export folder containing mesh.npz, camera_intrinsics.npz, world_to_cam.npy",
    )
    parser.add_argument("--map_cache", default=DEFAULT_FINGER_MAP_CACHE, help="Cached UV domain .npz")
    parser.add_argument("--camera_id", type=int, default=0)
    parser.add_argument("--video", help="Video file path (overrides camera)")
    parser.add_argument("--config", "-c", default="./config/dd02-ov.json", help="Orisys SDK config")
    parser.add_argument(
        "--source_scale",
        type=float,
        default=None,
        help="Resize live source before UV remap (default: camera.source_scale in config)",
    )
    parser.add_argument("--no_flip_source_vertical", action="store_true", help="Skip vertical flip before UV remap")
    parser.add_argument("--show", action="store_true", help="Show 2D unrolled image window with flow arrows")
    parser.add_argument("--verbose", action="store_true", help="SDK verbose logs")
    parser.add_argument(
        "--magnitude-max",
        type=float,
        default=None,
        help="Upper optical-flow magnitude for 3D surface heat-map scaling; default uses frame percentile.",
    )
    parser.add_argument(
        "--mesh-map",
        choices=["flow", "depth"],
        default="depth",
        help="Color the 3D mesh from optical-flow magnitude or the contact depth/divergence map.",
    )
    parser.add_argument(
        "--flow-threshold",
        type=float,
        default=0.4,
        help="Minimum optical-flow magnitude to color 3D surface / draw 2D arrows (matches draw_arrows).",
    )
    parser.add_argument(
        "--depth-threshold",
        type=float,
        default=0.0,
        help="Minimum depth/divergence value to color on the 3D mesh.",
    )
    parser.add_argument("--no-3d", action="store_true", help="Disable Open3D mesh viewer")
    args = parser.parse_args()

    video = args.camera_id
    if args.video:
        video = args.video

    try:
        runtime = FingerRuntime.open(
            export_dir=args.export_dir,
            config=args.config,
            video=video,
            verbose=args.verbose,
        )
    except RuntimeError as exc:
        print(str(exc), file=sys.stderr)
        sys.exit(1)

    if args.verbose:
        geometry = runtime.geometry
        print(
            f"UV geometry: native={geometry.native_width}x{geometry.native_height}, "
            f"source={geometry.source_width}x{geometry.source_height} "
            f"(scale={geometry.source_scale:g}), output={geometry.out_width}x{geometry.out_height}, "
            f"map_cache={geometry.map_cache}"
        )

    mesh_viewer, _uv_mesh = open_mesh_viewer(runtime, args)
    runtime.warmup(color=False)
    tactile = runtime.tactile

    print("Press 'q' to quit, 'r' to reset tracker. Close the 3D window to stop.")
    try:
        while True:
            ok, unrolled = runtime.grab_and_process(color=False)
            if not ok:
                break

            tactile.compute_deformation(decompose=True, check_motion=True, threshold=0)
            tactile.compute_contact()

            fps, flow, vnormal, img, contour, centroid, depth_map, fnormal, fshearx, fsheary = tactile.read_info(
                tactile.info.FPS,
                tactile.info.VRAW,
                tactile.info.VNORMAL,
                tactile.info.IMG,
                tactile.info.CONTOUR,
                tactile.info.CENTROID,
                tactile.info.DEPTH,
                tactile.info.FNORMAL,
                tactile.info.FSHEARX,
                tactile.info.FSHEARY,
            )
            flow_mag = float((flow[:, :, 0] ** 2 + flow[:, :, 1] ** 2).mean() ** 0.5) if flow.size else 0.0

            print(
                f"FPS={fps:.2f} flow_rms={flow_mag:.2f} "
                f"mask={tactile.mask_stats['coverage']:.2f} normal={fnormal:.2f} shear=({fshearx:.2f}, {fsheary:.2f})",
                end="\r",
            )

            if mesh_viewer is not None:
                if not mesh_viewer.update(flow=flow, scalar_map=depth_map):
                    print("\n3D viewer closed.")
                    break

            if args.show:
                show_finger_uv_panels(
                    unrolled,
                    vnormal,
                    depth_map,
                    flow_threshold=args.flow_threshold,
                )

            key = cv2.waitKey(1) & 0xFF
            if key == ord("q"):
                print("\nExiting.")
                break
            if key == ord("r"):
                tactile.reset()
                print("\nTracker reset.")
    finally:
        runtime.close()
        if mesh_viewer is not None:
            mesh_viewer.close()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()