#!/usr/bin/env python3
"""
Minimum UV-unroll finger tactile pipeline.

Offline: load Blender export -> build_or_load_domain -> unroll_map_uv.npz
Runtime:  camera frame -> UVUnrollFrameProcessor -> Sensor -> forces/contact

For the full demo (UV debug panels, dashboard, benchmarks) see:
  tools/3d_patch_generation_blender/unroll_realtime.py
"""

import argparse
import os
import sys

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

import cv2
import orisys
from orisys.configs import load_config
from orisys.frame_processors import UVUnrollFrameProcessor
from orisys.unroll import (
    build_or_load_domain,
    compute_uv_distortion_metrics,
    load_calib_json,
    load_exported_data,
    recommend_edge_expand,
)


def open_source(video_path, camera_id):
    return cv2.VideoCapture(video_path) if video_path else cv2.VideoCapture(camera_id)


def resolve_export_paths(export_dir):
    """Return Blender export paths expected inside one folder."""
    export_dir = os.path.abspath(export_dir)
    paths = {
        "mesh": os.path.join(export_dir, "mesh.npz"),
        "intrinsics": os.path.join(export_dir, "camera_intrinsics.npz"),
        "pose": os.path.join(export_dir, "world_to_cam.npy"),
    }
    missing = [name for name, path in paths.items() if not os.path.isfile(path)]
    if missing:
        raise FileNotFoundError(
            f"Missing export file(s) in {export_dir}: {', '.join(missing)}"
        )
    return paths["mesh"], paths["intrinsics"], paths["pose"]


def main():
    parser = argparse.ArgumentParser(description="Minimum UV-unroll finger tactile pipeline")
    parser.add_argument(
        "--export-dir",
        required=True,
        help="Blender export folder containing mesh.npz, camera_intrinsics.npz, world_to_cam.npy",
    )
    parser.add_argument("--calib", default=None, help="OpenCV calibration JSON for the live camera")
    parser.add_argument("--map_cache", default="unroll_map_uv.npz", 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("--no_flip_source_vertical", action="store_true", help="Skip vertical flip before UV remap")
    parser.add_argument("--show", action="store_true", help="Show unrolled image window")
    parser.add_argument("--verbose", action="store_true", help="SDK verbose logs")
    args = parser.parse_args()

    sdk_config = load_config(args.config)
    out_width, out_height = sdk_config["camera"]["resolution_stitch"]

    mesh_path, intrinsics_path, pose_path = resolve_export_paths(args.export_dir)

    data = load_exported_data(mesh_path, intrinsics_path, pose_path)
    distortion = compute_uv_distortion_metrics(data)
    edge_expand = recommend_edge_expand(distortion)

    if args.calib:
        source_K, dist_coeffs, img_size = load_calib_json(args.calib)
        source_width, source_height = img_size
    else:
        source_K = data["K"]
        dist_coeffs = None
        source_width = data["width"]
        source_height = data["height"]

    domain = build_or_load_domain(
        mesh_path,
        intrinsics_path,
        pose_path,
        args.map_cache,
        source_K,
        source_width,
        source_height,
        out_width,
        out_height,
        edge_expand=edge_expand,
        include_surface=True,
        verbose=args.verbose,
    )
    domain.remap.ensure_fast_maps()

    cap = open_source(args.video, args.camera_id)
    if not cap.isOpened():
        print("Cannot open camera/video source", file=sys.stderr)
        sys.exit(1)
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, source_width)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, source_height)

    sensor = orisys.Sensor("external", config_name=args.config, verbose=args.verbose, mirror=False)
    processor = UVUnrollFrameProcessor(
        domain.remap,
        camera_matrix=source_K if dist_coeffs is not None else None,
        dist_coeffs=dist_coeffs,
        source_width=source_width,
        source_height=source_height,
        flip_source_vertical=not args.no_flip_source_vertical,
    )

    ok, frame = cap.read()
    if ok:
        if args.show:
            prepared = processor.prepare_source(frame)
            sensor.get_img(
                frame=processor.flatten_prepared(prepared),
                already_processed=True,
                valid_mask=processor.valid_mask,
            )
        else:
            prepared_gray = processor.prepare_source_gray(frame)
            sensor.get_img(
                frame=processor.flatten_prepared_gray(prepared_gray),
                already_processed=True,
                already_gray=True,
                valid_mask=processor.valid_mask,
            )
        sensor.warmup_deformation()

    print("Press 'q' to quit, 'r' to reset tracker.")
    try:
        while True:
            ok, frame = cap.read()
            if not ok:
                break

            if args.show:
                prepared = processor.prepare_source(frame)
                unrolled = processor.flatten_prepared(prepared)
                sensor.get_img(
                    frame=unrolled,
                    already_processed=True,
                    valid_mask=processor.valid_mask,
                )
            else:
                prepared_gray = processor.prepare_source_gray(frame)
                unrolled = processor.flatten_prepared_gray(prepared_gray)
                sensor.get_img(
                    frame=unrolled,
                    already_processed=True,
                    already_gray=True,
                    valid_mask=processor.valid_mask,
                )
            sensor.compute_deformation(check_motion=True, threshold=0)
            sensor.compute_contact()

            fps, fnormal, fshearx, fsheary = sensor.read_info(
                sensor.info.FPS,
                sensor.info.FNORMAL,
                sensor.info.FSHEARX,
                sensor.info.FSHEARY,
            )

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

            if args.show:
                cv2.imshow("UV Unrolled", unrolled)

            key = cv2.waitKey(1) & 0xFF
            if key == ord("q"):
                print("\nExiting.")
                break
            if key == ord("r"):
                sensor.reset()
                print("\nTracker reset.")
    finally:
        cap.release()
        sensor.disconnect()
        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()