"""
Prototype tactile 2.5D motion analysis on top of the existing dense flow output.

This script compares several motion representations derived from optical flow:
1. Raw flow arrows / magnitude
2. Radial and tangential flow around an estimated contact center
3. Differential flow features: divergence, curl, strain magnitude
4. Pseudo-height reconstructed from divergence (qualitative only)
"""
import os
import sys
import argparse
from pathlib import Path
import time

os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"

import cv2
import numpy as np

WINDOW_TITLE = "Tactile 2.5D Motion Prototype"

_local_src = os.path.join(os.path.dirname(__file__), "src")
if _local_src not in sys.path:
    sys.path.insert(0, _local_src)

import orisys


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--video", "-v", type=str, default="0", help="Camera ID / video file path")
    parser.add_argument("--config", "-c", type=str, default="./config/ddjx01.json", help="Configuration file path")
    parser.add_argument("--cal", "-cal", type=str, default="./config/ddjx01.npy", help="Calibration file path")
    parser.add_argument("--verbose", action="store_true", help="Print detailed sensor logs")
    parser.add_argument("--no-display", action="store_true", help="Disable OpenCV windows")
    parser.add_argument("--save", type=str, default="", help="Optional output video path for the comparison panel")
    parser.add_argument("--max-frames", type=int, default=0, help="Stop after N frames (0 = unlimited)")
    parser.add_argument("--grid-spacing", type=int, default=20, help="Arrow visualization grid spacing")
    parser.add_argument("--flow-threshold", type=float, default=2.0, help="Arrow visualization magnitude threshold")
    parser.add_argument("--print-every", type=int, default=1, help="Print scalar metrics every N frames")
    parser.add_argument("--panel-size", type=int, default=320, help="Width/height of each panel tile")
    parser.add_argument("--plane-width-mm", type=float, default=20.0, help="Physical tactile imaging plane width in mm")
    parser.add_argument("--plane-height-mm", type=float, default=20.0, help="Physical tactile imaging plane height in mm")
    parser.add_argument("--camera-distance-mm", type=float, default=10.0, help="Distance from camera plane to imaging plane in mm")
    parser.add_argument("--display-max-width", type=int, default=1800, help="Maximum preview window width in pixels")
    parser.add_argument("--display-max-height", type=int, default=1000, help="Maximum preview window height in pixels")
    return parser.parse_args()


def create_sensor(args):
    try:
        video_input = int(args.video)
    except ValueError:
        video_input = args.video

    return orisys.Sensor(
        video_input,
        config_name=args.config,
        cal_path=args.cal,
        verbose=args.verbose,
    )


def ensure_bgr(image):
    if image.ndim == 2:
        return cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
    return image


def resize_panel(image, panel_size):
    return cv2.resize(ensure_bgr(image), (panel_size, panel_size), interpolation=cv2.INTER_LINEAR)


def fit_display_image(image, max_width, max_height):
    h, w = image.shape[:2]
    scale = min(max_width / max(float(w), 1.0), max_height / max(float(h), 1.0), 1.0)
    if scale >= 0.999:
        return image
    new_size = (max(1, int(w * scale)), max(1, int(h * scale)))
    return cv2.resize(image, new_size, interpolation=cv2.INTER_AREA)


def signed_heatmap(field, colormap=cv2.COLORMAP_TURBO):
    max_abs = float(np.max(np.abs(field)))
    if max_abs <= 1e-9:
        norm = np.full(field.shape, 127, dtype=np.uint8)
    else:
        norm = np.clip((field / max_abs + 1.0) * 127.5, 0, 255).astype(np.uint8)
    return cv2.applyColorMap(norm, colormap)


def magnitude_heatmap(field, colormap=cv2.COLORMAP_TURBO):
    max_val = float(np.max(field))
    if max_val <= 1e-9:
        norm = np.zeros(field.shape, dtype=np.uint8)
    else:
        norm = np.clip(field / max_val * 255.0, 0, 255).astype(np.uint8)
    return cv2.applyColorMap(norm, colormap)


def add_title(image, title):
    panel = image.copy()
    cv2.rectangle(panel, (0, 0), (panel.shape[1], 28), (0, 0, 0), -1)
    cv2.putText(panel, title, (8, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2, cv2.LINE_AA)
    return panel


def make_panel(image, title, panel_size):
    return add_title(resize_panel(image, panel_size), title)


def stack_panels(panels, columns=4):
    filler = np.zeros_like(panels[0])
    rows = []
    for idx in range(0, len(panels), columns):
        row = panels[idx:idx + columns]
        if len(row) < columns:
            row = row + [filler] * (columns - len(row))
        rows.append(cv2.hconcat(row))
    return cv2.vconcat(rows)


def estimate_contact_center(divergence, flow_magnitude):
    positive = np.maximum(divergence, 0.0)
    weights = positive * (flow_magnitude + 1e-6)
    total = float(weights.sum())
    h, w = divergence.shape

    if total <= 1e-9:
        return (w // 2, h // 2)

    ys, xs = np.indices(divergence.shape, dtype=np.float32)
    cx = int(np.clip((weights * xs).sum() / total, 0, w - 1))
    cy = int(np.clip((weights * ys).sum() / total, 0, h - 1))
    return (cx, cy)


def prepare_support_mask(support_map):
    if support_map is None:
        return None

    support = np.abs(support_map).astype(np.float32)
    max_support = float(np.max(support))
    if max_support > 0.0:
        support /= max_support
    support = (support > 0.05).astype(np.float32)
    if float(np.mean(support)) < 0.002:
        return None
    return support


def decompose_radial_tangential(flow, center):
    h, w = flow.shape[:2]
    ys, xs = np.indices((h, w), dtype=np.float32)
    cx, cy = center

    rx = xs - float(cx)
    ry = ys - float(cy)
    r_norm = np.sqrt(rx * rx + ry * ry)
    r_norm = np.where(r_norm < 1e-6, 1.0, r_norm)

    rhat_x = rx / r_norm
    rhat_y = ry / r_norm
    that_x = -rhat_y
    that_y = rhat_x

    u = flow[:, :, 0]
    v = flow[:, :, 1]
    radial = u * rhat_x + v * rhat_y
    tangential = u * that_x + v * that_y
    return radial, tangential


def compute_flow_jacobian_features(flow):
    u = flow[:, :, 0]
    v = flow[:, :, 1]

    du_dy, du_dx = np.gradient(u)
    dv_dy, dv_dx = np.gradient(v)

    divergence = du_dx + dv_dy
    curl = dv_dx - du_dy
    strain_xx = du_dx
    strain_yy = dv_dy
    strain_xy = 0.5 * (du_dy + dv_dx)
    strain_mag = np.sqrt(strain_xx * strain_xx + strain_yy * strain_yy + 2.0 * strain_xy * strain_xy)

    return {
        "divergence": divergence,
        "curl": curl,
        "strain_mag": strain_mag,
    }


def reconstruct_pseudo_height(divergence, flow_magnitude=None, support_map=None):
    rhs = divergence.astype(np.float32).copy()

    # No-contact guard: tiny flow noise becomes large low-frequency drift after Poisson solve.
    if flow_magnitude is not None:
        mean_mag = float(np.mean(flow_magnitude))
        max_mag = float(np.max(flow_magnitude))
        if mean_mag < 0.03 and max_mag < 0.15:
            return np.zeros_like(rhs, dtype=np.float32)

    support = prepare_support_mask(support_map)
    if support_map is not None:
        if support is None:
            return np.zeros_like(rhs, dtype=np.float32)
        rhs *= support

    noise_floor = max(float(np.std(rhs)) * 1.5, 1e-4)
    rhs[np.abs(rhs) < noise_floor] = 0.0
    if not np.any(rhs):
        return np.zeros_like(rhs, dtype=np.float32)

    # Zero-pad before FFT solve to reduce periodic wrap-around ghosts, then crop back.
    h, w = rhs.shape
    pad_y = max(8, h // 2)
    pad_x = max(8, w // 2)
    rhs_padded = np.pad(rhs, ((pad_y, pad_y), (pad_x, pad_x)), mode="constant")
    rhs_padded = rhs_padded - float(rhs_padded.mean())

    hp, wp = rhs_padded.shape
    fy = np.fft.fftfreq(hp).reshape(-1, 1)
    fx = np.fft.fftfreq(wp).reshape(1, -1)
    denom = (2.0 * np.pi * fx) ** 2 + (2.0 * np.pi * fy) ** 2
    denom[0, 0] = 1.0

    rhs_hat = np.fft.fft2(rhs_padded)
    height_hat = -rhs_hat / denom
    height_hat[0, 0] = 0.0
    height_full = np.real(np.fft.ifft2(height_hat)).astype(np.float32)
    height = height_full[pad_y:pad_y + h, pad_x:pad_x + w]
    height = cv2.GaussianBlur(height, (0, 0), 1.0)
    return height


def build_visualization(img, flow, center, features, radial, tangential, pseudo_height, depth_map, args):
    arrows = orisys.util.draw_arrows(
        img,
        flow,
        threshold=args.flow_threshold,
        grid_spacing=args.grid_spacing,
        arrow_scale=1.0,
    )

    contact_view = ensure_bgr(img)
    cv2.circle(contact_view, center, 6, (0, 255, 0), -1)
    cv2.circle(contact_view, center, 12, (0, 255, 0), 2)

    panel_images = [
        (contact_view, "Image / center"),
        (arrows, "Raw flow"),
        (signed_heatmap(radial), "Radial flow"),
        (signed_heatmap(tangential), "Tangential flow"),
        (signed_heatmap(features["divergence"]), "Divergence"),
        (signed_heatmap(features["curl"]), "Curl"),
        (magnitude_heatmap(features["strain_mag"]), "Strain mag"),
        (signed_heatmap(depth_map), "Depth map"),
        (signed_heatmap(pseudo_height), "Pseudo height"),
    ]
    panels = [make_panel(image, title, args.panel_size) for image, title in panel_images]
    return stack_panels(panels)


def make_writer(path, frame_size, fps):
    if not path:
        return None
    parent = Path(path).parent
    if parent and not parent.exists():
        parent.mkdir(parents=True, exist_ok=True)
    return cv2.VideoWriter(path, cv2.VideoWriter_fourcc(*"mp4v"), fps, frame_size)


def approximate_raw_flow_to_mm(flow_px, plane_width_mm, plane_height_mm, camera_distance_mm, image_shape):
    h, w = image_shape[:2]
    fx = camera_distance_mm * w / max(plane_width_mm, 1e-6)
    fy = camera_distance_mm * h / max(plane_height_mm, 1e-6)
    flow_mm = flow_px.copy().astype(np.float32)
    flow_mm[:, :, 0] *= camera_distance_mm / max(fx, 1e-6)
    flow_mm[:, :, 1] *= camera_distance_mm / max(fy, 1e-6)
    return flow_mm


def compute_metrics(flow_magnitude, raw_flow_mm, radial, tangential, features, depth_map, pseudo_height):
    return {
        "flow_mean": float(flow_magnitude.mean()),
        "flow_mean_mm": float(np.linalg.norm(raw_flow_mm, axis=2).mean()),
        "radial_abs_mean": float(np.mean(np.abs(radial))),
        "tangential_abs_mean": float(np.mean(np.abs(tangential))),
        "divergence_mean": float(features["divergence"].mean()),
        "curl_abs_mean": float(np.mean(np.abs(features["curl"]))),
        "strain_mean": float(features["strain_mag"].mean()),
        "depth_peak": float(np.max(np.abs(depth_map))),
        "pseudo_height_peak": float(np.max(np.abs(pseudo_height))),
    }


def format_metrics(frame_count, fps, metrics):
    line = (
        f"frame={frame_count:04d} fps={fps:6.2f} "
        f"flow_mean={metrics['flow_mean']:.4f} "
        f"flow_mean_mm={metrics['flow_mean_mm']:.4f} "
        f"radial_abs={metrics['radial_abs_mean']:.4f} "
        f"tangential_abs={metrics['tangential_abs_mean']:.4f} "
        f"div_mean={metrics['divergence_mean']:.4f} "
        f"curl_abs={metrics['curl_abs_mean']:.4f} "
        f"strain={metrics['strain_mean']:.4f} "
        f"depth_peak={metrics['depth_peak']:.4f} "
        f"height_peak={metrics['pseudo_height_peak']:.4f}"
    )
    t = metrics.get("_timings")
    if t:
        line += (
            f" | analyze {t['total']:.2f}ms"
            f" (mag:{t['magnitude']:.2f} jac:{t['jacobian']:.2f}"
            f" rad_tan:{t['rad_tan']:.2f} pseudo_h:{t['pseudo_h']:.2f}"
            f" metrics:{t['metrics']:.2f} vis:{t['visualize']:.2f})"
        )
    return line


def read_sensor_frame(sensor):
    sensor.compute_flow()
    sensor.compute_hhd()
    sensor.compute_contact()
    return sensor.read_info(
        sensor.info.FPS,
        sensor.info.VRAW,
        sensor.info.IMG,
        sensor.info.DEPTH,
    )


def analyze_frame(img, flow, depth_map, args):
    t = time.perf_counter()

    flow_magnitude = np.linalg.norm(flow, axis=2)
    t_mag = time.perf_counter()

    features = compute_flow_jacobian_features(flow)
    t_jac = time.perf_counter()

    center = estimate_contact_center(features["divergence"], flow_magnitude)
    radial, tangential = decompose_radial_tangential(flow, center)
    t_decomp = time.perf_counter()

    pseudo_height = reconstruct_pseudo_height(
        features["divergence"],
        flow_magnitude=flow_magnitude,
        support_map=depth_map,
    )
    t_height = time.perf_counter()

    raw_flow_mm = approximate_raw_flow_to_mm(
        flow,
        args.plane_width_mm,
        args.plane_height_mm,
        args.camera_distance_mm,
        img.shape,
    )
    metrics = compute_metrics(flow_magnitude, raw_flow_mm, radial, tangential, features, depth_map, pseudo_height)
    t_metrics = time.perf_counter()

    canvas = build_visualization(
        img,
        flow,
        center,
        features,
        radial,
        tangential,
        pseudo_height,
        depth_map,
        args,
    )
    t_vis = time.perf_counter()

    metrics["_timings"] = {
        "magnitude": (t_mag - t) * 1000,
        "jacobian": (t_jac - t_mag) * 1000,
        "rad_tan": (t_decomp - t_jac) * 1000,
        "pseudo_h": (t_height - t_decomp) * 1000,
        "metrics": (t_metrics - t_height) * 1000,
        "visualize": (t_vis - t_metrics) * 1000,
        "total": (t_vis - t) * 1000,
    }
    return canvas, metrics


def main():
    args = parse_args()
    sensor = create_sensor(args)
    writer = None
    frame_count = 0
    print_every = max(args.print_every, 1)

    try:
        print("\nPress 'q' to quit, 'r' to reset tracking.")

        while True:
            img = sensor.get_img()
            if img is None:
                print("No more frames returned from sensor.")
                break

            fps, flow, img, depth_map = read_sensor_frame(sensor)
            canvas, metrics = analyze_frame(img, flow, depth_map, args)

            if writer is None and args.save:
                writer = make_writer(args.save, (canvas.shape[1], canvas.shape[0]), max(float(fps), 1.0))

            if writer is not None:
                writer.write(canvas)

            frame_count += 1

            if frame_count % print_every == 0:
                print(format_metrics(frame_count, fps, metrics))

            if not args.no_display:
                display_canvas = fit_display_image(
                    canvas,
                    max_width=args.display_max_width,
                    max_height=args.display_max_height,
                )
                cv2.imshow(WINDOW_TITLE, display_canvas)
                key = cv2.waitKey(1) & 0xFF
                if key == ord("q"):
                    print("\nExiting...")
                    break
                if key == ord("r"):
                    sensor.reset()
                    print("Sensor reset.")

            if args.max_frames > 0 and frame_count >= args.max_frames:
                print(f"Reached max_frames={args.max_frames}.")
                break
    finally:
        if writer is not None:
            writer.release()
        sensor.disconnect()
        if not args.no_display:
            cv2.destroyAllWindows()


if __name__ == "__main__":
    main()