import argparse
import os
import sys
import time

import cv2
import numpy as np


SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
REPO_ROOT = os.path.abspath(os.path.join(SCRIPT_DIR, "..", ".."))
SRC_DIR = os.path.join(REPO_ROOT, "src")
if SRC_DIR not in sys.path:
    sys.path.insert(0, SRC_DIR)

import orisys


def load_flatten_remap(path):
    remap = np.load(path)
    result = {
        "map_x": remap["map_x"].astype(np.float32),
        "map_y": remap["map_y"].astype(np.float32),
        "valid_mask": remap["valid_mask"].astype(bool),
    }
    if "flat_xyz" in remap:
        result["flat_xyz"] = remap["flat_xyz"].astype(np.float32)
    return result


def flip_sensor_frames_vertical(sensor):
    """Flip captured BGR/gray frames to match Blender-exported map orientation."""
    if sensor._img_bgr_raw is not None:
        sensor._img_bgr_raw = cv2.flip(sensor._img_bgr_raw, 0)
    if sensor.img is not None:
        cv2.flip(sensor.img, 0, dst=sensor.img)


def flatten_frame(frame, remap, border_value=0):
    flat = cv2.remap(
        frame,
        remap["map_x"],
        remap["map_y"],
        interpolation=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=border_value,
    )
    flat = flat.copy()
    flat[~remap["valid_mask"]] = border_value
    return flat


def load_exported_data(data_dir):
    mesh = np.load(os.path.join(data_dir, "mesh.npz"))
    cam_data = np.load(os.path.join(data_dir, "camera_intrinsics.npz"))
    xyz_map_path = os.path.join(data_dir, "pixel_to_xyz.npz")
    if not os.path.exists(xyz_map_path):
        raise FileNotFoundError(f"Missing precomputed 3D map: {xyz_map_path}")

    xyz_map = np.load(xyz_map_path)["xyz_map"]
    width = int(cam_data["width"])
    height = int(cam_data["height"])
    if xyz_map.shape[:2] != (height, width):
        raise ValueError(
            f"pixel_to_xyz shape {xyz_map.shape[:2]} does not match camera size {(height, width)}"
        )

    return {
        "vertices": mesh["vertices"],
        "faces": mesh["faces"],
        "uv_data": mesh["uv_data"],
        "loop_to_vertex": mesh["loop_to_vertex"],
        "K": cam_data["K"],
        "width": width,
        "height": height,
        "world_to_cam": np.load(os.path.join(data_dir, "world_to_cam.npy")),
        "xyz_map": xyz_map,
    }


def parse_source(value):
    value = str(value).strip()
    for prefix in ("camera:", "cam:", "webcam:"):
        if value.lower().startswith(prefix):
            value = value[len(prefix):].strip()
            break

    try:
        return int(value)
    except ValueError:
        return value


def prepare_static_points(data, display_stride):
    xyz_map = data["xyz_map"]
    valid_mask = ~np.isnan(xyz_map[:, :, 0])
    valid_points = xyz_map[valid_mask].reshape(-1, 3)
    if len(valid_points) == 0:
        raise ValueError("pixel_to_xyz.npz contains no valid mesh hit points.")

    display_indices = np.arange(len(valid_points))[::max(1, int(display_stride))]
    return valid_mask, display_indices, valid_points[display_indices]


def align_image_to_map(img_bgr, data, valid_mask, display_indices, flip_vertical=True, flip_horizontal=True):
    target_w = int(data["width"])
    target_h = int(data["height"])
    h, w = img_bgr.shape[:2]
    if (w, h) != (target_w, target_h):
        img_bgr = cv2.resize(img_bgr, (target_w, target_h), interpolation=cv2.INTER_LINEAR)

    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    if flip_vertical:
        img_rgb = np.flipud(img_rgb).copy()
    if flip_horizontal:
        img_rgb = np.fliplr(img_rgb).copy()

    colors = img_rgb.astype(np.float64) / 255.0
    return colors[valid_mask].reshape(-1, 3)[display_indices]


def arrow_map_reference_colors(count, color=(0.04, 0.04, 0.04)):
    return np.tile(np.asarray(color, dtype=np.float64), (count, 1))


def magnitude_map_colors(flow, valid_mask, display_indices, max_magnitude=None):
    mag = np.linalg.norm(flow, axis=2)
    valid_mag = mag[valid_mask].reshape(-1)
    display_mag = valid_mag[display_indices]
    if max_magnitude is None or max_magnitude <= 0:
        max_magnitude = max(float(np.percentile(display_mag, 95)), 1e-6)

    value = np.clip(display_mag / max_magnitude, 0.0, 1.0)
    return np.column_stack(
        (
            value,
            np.clip(1.0 - np.abs(value - 0.5) * 2.0, 0.0, 1.0),
            1.0 - value,
        )
    ).astype(np.float64)


def align_scalar_field_to_map(field, data, flip_vertical=True):
    if field.ndim == 3:
        field = np.linalg.norm(field, axis=2)

    target_w = int(data["width"])
    target_h = int(data["height"])
    h, w = field.shape[:2]
    if (w, h) != (target_w, target_h):
        field = cv2.resize(field, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
    else:
        field = field.copy()

    if flip_vertical:
        field = np.flipud(field).copy()
    return field


def scalar_map_colors(field, valid_mask, display_indices, value_min=None, value_max=None):
    valid_values = field[valid_mask].reshape(-1)
    display_values = valid_values[display_indices]
    finite_values = display_values[np.isfinite(display_values)]
    if len(finite_values) == 0:
        return arrow_map_reference_colors(len(display_indices))

    if value_min is None:
        value_min = float(np.percentile(finite_values, 5))
    if value_max is None:
        value_max = float(np.percentile(finite_values, 95))
    if value_max <= value_min:
        value_max = value_min + 1e-6

    value = np.clip((display_values - value_min) / (value_max - value_min), 0.0, 1.0)
    return np.column_stack(
        (
            value,
            np.clip(1.0 - np.abs(value - 0.5) * 2.0, 0.0, 1.0),
            1.0 - value,
        )
    ).astype(np.float64)


def align_vector_field_to_map(field, data, flip_vertical=True, scale_as_flow=True):
    target_w = int(data["width"])
    target_h = int(data["height"])
    h, w = field.shape[:2]
    if (w, h) != (target_w, target_h):
        sx = target_w / max(w, 1)
        sy = target_h / max(h, 1)
        field = cv2.resize(field, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
        if scale_as_flow:
            field[:, :, 0] *= sx
            field[:, :, 1] *= sy
    else:
        field = field.copy()

    if flip_vertical:
        field = np.flipud(field).copy()
        field[:, :, 1] *= -1.0
    return field


def estimate_surface_normals(xyz_map):
    valid = ~np.isnan(xyz_map[:, :, 0])
    filled = xyz_map.copy()
    filled[~valid] = 0
    d_y = np.gradient(filled, axis=0)
    d_x = np.gradient(filled, axis=1)
    normals = np.cross(d_x, d_y)
    norm = np.linalg.norm(normals, axis=2, keepdims=True)
    normals = np.divide(normals, norm, out=np.zeros_like(normals), where=norm > 1e-8)
    normals[~valid] = 0
    return normals


def line_colors(mag, arrow_map_style):
    if len(mag) == 0:
        return np.empty((0, 3), dtype=np.float64)
    if arrow_map_style:
        return np.tile(np.array([0.0, 0.9, 1.0], dtype=np.float64), (len(mag), 1))
    norm_mag = mag / max(float(np.percentile(mag, 95)), 1e-6)
    norm_mag = np.clip(norm_mag, 0.0, 1.0)
    return np.column_stack((norm_mag, 1.0 - norm_mag, np.zeros_like(norm_mag))).astype(np.float64)


def lines_from_vectors(starts, vectors, mag, flow_scale, arrow_map_style):
    vec_norm = np.linalg.norm(vectors, axis=1)
    good = np.isfinite(vec_norm) & (vec_norm > 1e-8)
    starts, vectors, mag = starts[good], vectors[good], mag[good]
    if len(starts) == 0:
        return np.empty((0, 3)), np.empty((0, 2), dtype=np.int32), np.empty((0, 3))

    points = np.empty((len(starts) * 2, 3), dtype=np.float64)
    points[0::2] = starts
    points[1::2] = starts + vectors * flow_scale
    lines = np.column_stack((np.arange(0, len(points), 2), np.arange(1, len(points), 2))).astype(np.int32)
    return points, lines, line_colors(mag, arrow_map_style)


def build_flow_lines(
    xyz_map,
    valid_mask,
    flow,
    arrow_stride=32,
    flow_scale=1.0,
    min_magnitude=0.5,
    arrow_map_style=False,
):
    h, w = flow.shape[:2]
    ys, xs = np.mgrid[0:h:arrow_stride, 0:w:arrow_stride]
    ys = ys.reshape(-1)
    xs = xs.reshape(-1)
    sample_flow = flow[ys, xs]
    mag = np.linalg.norm(sample_flow, axis=1)
    keep = valid_mask[ys, xs] & (mag >= min_magnitude)
    xs, ys, sample_flow, mag = xs[keep], ys[keep], sample_flow[keep], mag[keep]
    if len(xs) == 0:
        return np.empty((0, 3)), np.empty((0, 2), dtype=np.int32), np.empty((0, 3))

    tx = np.rint(xs + sample_flow[:, 0]).astype(np.int32)
    ty = np.rint(ys + sample_flow[:, 1]).astype(np.int32)
    inside = (0 <= tx) & (tx < w) & (0 <= ty) & (ty < h)
    xs, ys, tx, ty, mag = xs[inside], ys[inside], tx[inside], ty[inside], mag[inside]
    if len(xs) == 0:
        return np.empty((0, 3)), np.empty((0, 2), dtype=np.int32), np.empty((0, 3))

    valid_target = valid_mask[ty, tx]
    xs, ys, tx, ty, mag = xs[valid_target], ys[valid_target], tx[valid_target], ty[valid_target], mag[valid_target]
    starts = xyz_map[ys, xs]
    vectors = xyz_map[ty, tx] - starts
    return lines_from_vectors(starts, vectors, mag, flow_scale, arrow_map_style)


def build_shear_normal_lines(
    xyz_map,
    valid_mask,
    vshear,
    vnormal,
    normals,
    arrow_stride=32,
    flow_scale=1.0,
    min_magnitude=0.01,
    normal_gain=1.0,
    arrow_map_style=False,
):
    h, w = vshear.shape[:2]
    ys, xs = np.mgrid[0:h:arrow_stride, 0:w:arrow_stride]
    ys = ys.reshape(-1)
    xs = xs.reshape(-1)

    shear = vshear[ys, xs]
    normal_field = vnormal[ys, xs]
    shear_mag = np.linalg.norm(shear, axis=1)
    normal_mag = np.linalg.norm(normal_field, axis=1) * normal_gain
    combined_mag = shear_mag + normal_mag
    keep = valid_mask[ys, xs] & (combined_mag >= min_magnitude)
    xs, ys, shear, shear_mag, normal_mag, combined_mag = (
        xs[keep],
        ys[keep],
        shear[keep],
        shear_mag[keep],
        normal_mag[keep],
        combined_mag[keep],
    )
    if len(xs) == 0:
        return np.empty((0, 3)), np.empty((0, 2), dtype=np.int32), np.empty((0, 3))

    tx = np.rint(xs + shear[:, 0]).astype(np.int32)
    ty = np.rint(ys + shear[:, 1]).astype(np.int32)
    inside = (0 <= tx) & (tx < w) & (0 <= ty) & (ty < h)
    xs, ys, tx, ty, shear_mag, normal_mag, combined_mag = (
        xs[inside],
        ys[inside],
        tx[inside],
        ty[inside],
        shear_mag[inside],
        normal_mag[inside],
        combined_mag[inside],
    )
    if len(xs) == 0:
        return np.empty((0, 3)), np.empty((0, 2), dtype=np.int32), np.empty((0, 3))

    valid_target = valid_mask[ty, tx]
    xs, ys, tx, ty, shear_mag, normal_mag, combined_mag = (
        xs[valid_target],
        ys[valid_target],
        tx[valid_target],
        ty[valid_target],
        shear_mag[valid_target],
        normal_mag[valid_target],
        combined_mag[valid_target],
    )

    starts = xyz_map[ys, xs]
    tangent = xyz_map[ty, tx] - starts
    tangent_norm = np.linalg.norm(tangent, axis=1, keepdims=True)
    tangent_unit = np.divide(tangent, tangent_norm, out=np.zeros_like(tangent), where=tangent_norm > 1e-8)
    normal_unit = normals[ys, xs]

    shear_ratio = shear_mag / np.maximum(combined_mag, 1e-8)
    normal_ratio = normal_mag / np.maximum(combined_mag, 1e-8)
    vectors = tangent_unit * shear_ratio[:, None] + normal_unit * normal_ratio[:, None]
    vector_norm = np.linalg.norm(vectors, axis=1, keepdims=True)
    vectors = np.divide(vectors, vector_norm, out=np.zeros_like(vectors), where=vector_norm > 1e-8)
    return lines_from_vectors(starts, vectors, combined_mag, flow_scale, arrow_map_style)


def init_viewer(cloud_points, cloud_colors, flow_points, flow_lines, flow_colors, arrow_map_style=False):
    import open3d as o3d

    vis = o3d.visualization.Visualizer()
    vis.create_window(window_name="SDK 3D Flow", width=1280, height=720)

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(cloud_points.astype(np.float64))
    pcd.colors = o3d.utility.Vector3dVector(cloud_colors.astype(np.float64))
    vis.add_geometry(pcd)

    line_set = o3d.geometry.LineSet()
    line_set.points = o3d.utility.Vector3dVector(flow_points.astype(np.float64))
    line_set.lines = o3d.utility.Vector2iVector(flow_lines)
    line_set.colors = o3d.utility.Vector3dVector(flow_colors.astype(np.float64))
    vis.add_geometry(line_set)

    render_option = vis.get_render_option()
    render_option.point_size = 1.0 if arrow_map_style else 2.0
    render_option.line_width = 2.0
    render_option.background_color = np.array([0.0, 0.0, 0.0])
    vis.poll_events()
    vis.update_renderer()
    return vis, pcd, line_set, o3d


def run_sdk_3d_flow(args):
    data = load_exported_data(args.data_dir)
    flatten_remap = load_flatten_remap(args.flatten_remap) if args.flatten_remap else None
    view_data = data
    if flatten_remap is not None:
        if "flat_xyz" not in flatten_remap:
            raise ValueError("Flatten remap must contain flat_xyz for 3D visualization.")
        view_data = dict(data)
        view_data["xyz_map"] = flatten_remap["flat_xyz"]
        view_data["height"], view_data["width"] = view_data["xyz_map"].shape[:2]

    valid_mask, display_indices, cloud_points = prepare_static_points(view_data, args.display_stride)
    surface_normals = estimate_surface_normals(view_data["xyz_map"])
    flip_vertical = not args.no_flip_vertical

    sensor = orisys.Sensor(
        parse_source(args.source),
        config_name=args.config,
        isstitch=False,
        cuda=args.cuda,
        verbose=args.verbose,
        mirror=not args.no_mirror,
    )

    vis = pcd = line_set = o3d = None
    frame_count = 0
    t0 = time.perf_counter()

    try:
        while True:
            if vis is not None and not vis.poll_events():
                break

            if sensor.get_img() is None:
                break
            if flip_vertical:
                flip_sensor_frames_vertical(sensor)
            flat_bgr = None
            if flatten_remap is not None:
                flat_bgr = flatten_frame(sensor.read_info(sensor.info.IMG)[0], flatten_remap)
                flat_gray = cv2.cvtColor(flat_bgr, cv2.COLOR_BGR2GRAY)
                sensor.img = flat_gray
                sensor._img_bgr_raw = flat_bgr
                sensor.img_size_raw = flat_bgr.shape
            needs_decomposition = args.vector_mode in ("shear-normal", "depth-map")
            sensor.compute_deformation(decompose=needs_decomposition, check_motion=False)
            if args.vector_mode == "depth-map":
                sensor.compute_contact()
            flow, img_bgr, vshear, vnormal, depth_map = sensor.read_info(
                sensor.info.VRAW,
                sensor.info.IMG,
                sensor.info.VSHEAR,
                sensor.info.VNORMAL,
                sensor.info.DEPTH,
            )
            if flat_bgr is not None:
                img_bgr = flat_bgr

            active_data = view_data if flatten_remap is not None else data
            # Input is already flipped above; flattened domain uses flat_xyz directly.
            flip_for_map = False if (flatten_remap is not None or flip_vertical) else True
            flow_map = align_vector_field_to_map(flow, active_data, flip_vertical=flip_for_map)
            if args.vector_mode in ("vraw-magnitude", "vraw-magnitude-arrows"):
                cloud_colors = magnitude_map_colors(flow_map, valid_mask, display_indices, args.magnitude_max)
            elif args.vector_mode == "depth-map":
                depth_map = align_scalar_field_to_map(depth_map, active_data, flip_vertical=flip_for_map)
                cloud_colors = scalar_map_colors(
                    depth_map,
                    valid_mask,
                    display_indices,
                    value_min=args.depth_min,
                    value_max=args.depth_max,
                )
            elif args.arrow_map_style:
                cloud_colors = arrow_map_reference_colors(len(cloud_points))
            else:
                cloud_colors = align_image_to_map(
                    img_bgr,
                    active_data,
                    valid_mask,
                    display_indices,
                    flip_vertical=flip_for_map,
                )

            if args.vector_mode == "shear-normal":
                vshear_map = align_vector_field_to_map(vshear, active_data, flip_vertical=flip_for_map)
                vnormal_map = align_vector_field_to_map(vnormal, active_data, flip_vertical=flip_for_map)
                flow_points, flow_lines, flow_colors = build_shear_normal_lines(
                    active_data["xyz_map"],
                    valid_mask,
                    vshear_map,
                    vnormal_map,
                    surface_normals,
                    arrow_stride=args.arrow_stride,
                    flow_scale=args.flow_scale,
                    min_magnitude=args.min_flow,
                    normal_gain=args.normal_gain,
                    arrow_map_style=args.arrow_map_style,
                )
            elif args.vector_mode in ("vraw-magnitude", "depth-map"):
                flow_points = np.empty((0, 3), dtype=np.float64)
                flow_lines = np.empty((0, 2), dtype=np.int32)
                flow_colors = np.empty((0, 3), dtype=np.float64)
            else:
                flow_points, flow_lines, flow_colors = build_flow_lines(
                    active_data["xyz_map"],
                    valid_mask,
                    flow_map,
                    arrow_stride=args.arrow_stride,
                    flow_scale=args.flow_scale,
                    min_magnitude=args.min_flow,
                    arrow_map_style=args.arrow_map_style,
                )

            if vis is None:
                vis, pcd, line_set, o3d = init_viewer(
                    cloud_points,
                    cloud_colors,
                    flow_points,
                    flow_lines,
                    flow_colors,
                    arrow_map_style=args.arrow_map_style,
                )
                print("SDK 3D flow viewer started. Close the Open3D window to stop.")
            else:
                pcd.colors = o3d.utility.Vector3dVector(cloud_colors)
                line_set.points = o3d.utility.Vector3dVector(flow_points.astype(np.float64))
                line_set.lines = o3d.utility.Vector2iVector(flow_lines)
                line_set.colors = o3d.utility.Vector3dVector(flow_colors.astype(np.float64))
                vis.update_geometry(pcd)
                vis.update_geometry(line_set)
                vis.update_renderer()

            frame_count += 1
            if frame_count % 30 == 0:
                elapsed = max(time.perf_counter() - t0, 1e-6)
                print(f"frames={frame_count}, viewer FPS={frame_count / elapsed:.1f}, arrows={len(flow_lines)}")
    finally:
        sensor.disconnect()
        if vis is not None:
            vis.destroy_window()


def parse_args():
    parser = argparse.ArgumentParser(description="Map SDK flow fields to a 3D point cloud.")
    parser.add_argument("--data-dir", default=os.path.join(SCRIPT_DIR, "exported_data"))
    parser.add_argument(
        "--source",
        default=os.path.join(SCRIPT_DIR, "20260602-094916.mp4"),
        help="Video path or live camera ID, for example 0, 1, camera:0, or cam:1.",
    )
    parser.add_argument("--config", default=os.path.join(REPO_ROOT, "config", "dd02-ov.json"))
    parser.add_argument("--display-stride", type=int, default=4)
    parser.add_argument("--arrow-stride", type=int, default=32)
    parser.add_argument("--flow-scale", type=float, default=3.0)
    parser.add_argument("--min-flow", type=float, default=0.5)
    parser.add_argument("--normal-gain", type=float, default=1.0)
    parser.add_argument("--magnitude-max", type=float, default=None,
                        help="Upper VRAW magnitude for heat-map color scaling; default uses frame percentile.")
    parser.add_argument("--depth-min", type=float, default=None,
                        help="Minimum depth value for heat-map color scaling; default uses frame percentile.")
    parser.add_argument("--depth-max", type=float, default=None,
                        help="Maximum depth value for heat-map color scaling; default uses frame percentile.")
    parser.add_argument("--flatten-remap", default=None,
                        help="Optional flatten_remap.npz to run optical flow in UV-flattened image space.")
    parser.add_argument(
        "--vector-mode",
        choices=["vraw", "vraw-magnitude", "vraw-magnitude-arrows", "shear-normal", "depth-map"],
        default="vraw",
    )
    parser.add_argument(
        "--arrow-map-style",
        action="store_true",
        help="Hide real texture and use a dark Qt-viewer-like arrow map style.",
    )
    parser.add_argument(
        "--no-flip-vertical",
        action="store_true",
        help="Do not flip captured video vertically before processing.",
    )
    parser.add_argument("--no-mirror", action="store_true")
    parser.add_argument("--cuda", action="store_true")
    parser.add_argument("--verbose", action="store_true")
    return parser.parse_args()


if __name__ == "__main__":
    run_sdk_3d_flow(parse_args())