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

import argparse
import multiprocessing as mp
import queue
import time
from multiprocessing import Process, Queue

import cv2
import numpy as np

from .paths import ensure_project_paths

ensure_project_paths()


def parse_video_source(value):
    if isinstance(value, str):
        value = value.strip()
        if value.isdigit():
            return int(value)
    return value


def _put_latest(result_queue, data):
    try:
        result_queue.put_nowait(data)
        return
    except queue.Full:
        pass

    try:
        result_queue.get_nowait()
    except Exception:
        pass

    try:
        result_queue.put_nowait(data)
    except queue.Full:
        pass


def _draw_magnitude_map(flow, visual_size, threshold=1.5):
    magnitude = np.sqrt(flow[:, :, 0] ** 2 + flow[:, :, 1] ** 2)
    h, w = magnitude.shape
    if h <= 0 or w <= 0:
        return None

    downsample_factor = 10
    magnitude_downsampled = magnitude[::downsample_factor, ::downsample_factor]
    magnitude_smoothed = cv2.resize(
        magnitude_downsampled,
        (w, h),
        interpolation=cv2.INTER_CUBIC,
    )

    magnitude_thresholded = magnitude_smoothed.copy()
    magnitude_thresholded[magnitude_smoothed < threshold] = 0

    mag_min = float(magnitude_thresholded.min())
    mag_max = float(magnitude_thresholded.max())
    if mag_max > mag_min:
        magnitude_norm = (
            (magnitude_thresholded - mag_min) / (mag_max - mag_min) * 255
        ).astype(np.uint8)
    else:
        magnitude_norm = np.zeros_like(magnitude_thresholded, dtype=np.uint8)

    magnitude_colored = cv2.applyColorMap(magnitude_norm, cv2.COLORMAP_JET)
    return cv2.resize(
        magnitude_colored,
        (visual_size, visual_size),
        interpolation=cv2.INTER_CUBIC,
    )


def _make_visuals(orisys_module, img, flow, visual_size):
    if img is None or flow is None or np.asarray(flow).size == 0:
        return None, None, 0.0, 0.0

    if len(img.shape) == 2:
        img_bgr = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    else:
        img_bgr = img

    h, w = img_bgr.shape[:2]
    bg = np.empty((h, w, 3), dtype=img_bgr.dtype)
    bg[:, :, 0] = 22
    bg[:, :, 1] = 16
    bg[:, :, 2] = 11

    arrows = orisys_module.util.draw_arrows(
        bg,
        flow,
        threshold=2,
        grid_spacing=20,
        arrow_scale=0.5,
        below_threshold_color=(160, 160, 160),
    )
    arrows = cv2.resize(arrows, (visual_size, visual_size), interpolation=cv2.INTER_CUBIC)
    magnitude = _draw_magnitude_map(flow, visual_size)
    flow_mag = np.sqrt(flow[:, :, 0] ** 2 + flow[:, :, 1] ** 2)
    return arrows, magnitude, float(flow_mag.mean()), float(flow_mag.max())


def _sensor_force_worker(
    sensor_id,
    vid_src,
    config_name,
    rotation_config_path,
    result_queue,
    stop_event,
    target_fps,
    cuda,
    isstitch,
    backend,
    motion_threshold,
    include_visuals,
    visual_size,
):
    ensure_project_paths()
    import orisys

    sensor = None
    try:
        sensor = orisys.Sensor(
            vid_src,
            isstitch=isstitch,
            cuda=cuda,
            verbose=False,
            config_name=config_name,
            backend=backend,
            rotation_config_path=rotation_config_path,
        )

        frame_interval = 1.0 / target_fps if target_fps and target_fps > 0 else 0.0

        while not stop_event.is_set():
            t_start = time.perf_counter()
            img = sensor.get_img()
            if img is None:
                time.sleep(0.01)
                continue

            is_contact = sensor.compute_deformation(
                check_motion=True,
                threshold=motion_threshold,
            )
            if include_visuals:
                fps, fnormal, fshearx, fsheary, flow, img_view = sensor.read_info(
                    sensor.info.FPS,
                    sensor.info.FNORMAL,
                    sensor.info.FSHEARX,
                    sensor.info.FSHEARY,
                    sensor.info.VRAW,
                    sensor.info.IMG,
                )
                flow_view, magnitude_view, flow_mean, flow_max = _make_visuals(
                    orisys,
                    img_view,
                    flow,
                    int(visual_size),
                )
            else:
                fps, fnormal, fshearx, fsheary = sensor.read_info(
                    sensor.info.FPS,
                    sensor.info.FNORMAL,
                    sensor.info.FSHEARX,
                    sensor.info.FSHEARY,
                )
                flow_view = None
                magnitude_view = None
                flow_mean = 0.0
                flow_max = 0.0

            data = {
                "sensor_id": sensor_id,
                "fnormal": float(fnormal),
                "fshearx": float(fshearx),
                "fsheary": float(fsheary),
                "fps": float(fps),
                "is_contact": bool(is_contact),
                "timestamp": time.time(),
                "flow_mean": flow_mean,
                "flow_max": flow_max,
            }
            if include_visuals:
                data["flow_view"] = flow_view
                data["magnitude_view"] = magnitude_view

            _put_latest(result_queue, data)

            if frame_interval > 0:
                sleep_time = frame_interval - (time.perf_counter() - t_start)
                if sleep_time > 0.001:
                    time.sleep(sleep_time)

    except Exception as exc:
        _put_latest(
            result_queue,
            {
                "sensor_id": sensor_id,
                "error": str(exc),
                "timestamp": time.time(),
            },
        )
    finally:
        if sensor is not None:
            try:
                sensor.disconnect()
            except Exception:
                pass


class TwoSensorForceReader:
    def __init__(
        self,
        video1=0,
        video2=1,
        config1="./config/ddjx01.json",
        config2="./config/ddjx01.json",
        rotation1="./config/rotation_config_0.json",
        rotation2="./config/rotation_config_1.json",
        target_fps=30,
        cuda=True,
        isstitch=True,
        backend="auto",
        motion_threshold=0,
        queue_size=5,
        include_visuals=False,
        visual_size=320,
    ):
        self.sensor_configs = [
            {
                "sensor_id": 0,
                "vid_src": parse_video_source(video1),
                "config_name": config1,
                "rotation_config_path": rotation1,
            },
            {
                "sensor_id": 1,
                "vid_src": parse_video_source(video2),
                "config_name": config2,
                "rotation_config_path": rotation2,
            },
        ]
        self.target_fps = target_fps
        self.cuda = cuda
        self.isstitch = isstitch
        self.backend = backend
        self.motion_threshold = motion_threshold
        self.queue_size = queue_size
        self.include_visuals = include_visuals
        self.visual_size = visual_size

        self.result_queues = []
        self.stop_events = []
        self.processes = []
        self.latest = [None, None]

    def start(self):
        if self.processes:
            return

        self.result_queues = [Queue(maxsize=self.queue_size) for _ in self.sensor_configs]
        self.stop_events = [mp.Event() for _ in self.sensor_configs]
        self.latest = [None] * len(self.sensor_configs)

        for idx, config in enumerate(self.sensor_configs):
            process = Process(
                target=_sensor_force_worker,
                args=(
                    config["sensor_id"],
                    config["vid_src"],
                    config["config_name"],
                    config["rotation_config_path"],
                    self.result_queues[idx],
                    self.stop_events[idx],
                    self.target_fps,
                    self.cuda,
                    self.isstitch,
                    self.backend,
                    self.motion_threshold,
                    self.include_visuals,
                    self.visual_size,
                ),
            )
            process.start()
            self.processes.append(process)

    def update(self):
        for idx, result_queue in enumerate(self.result_queues):
            data = None
            while True:
                try:
                    data = result_queue.get_nowait()
                except queue.Empty:
                    break
                except Exception:
                    break

            if data is not None:
                if "error" in data:
                    raise RuntimeError(
                        f"sensor {data.get('sensor_id', idx)} failed: {data['error']}"
                    )
                self.latest[idx] = data

        return self.latest

    def get_samples(self, timeout=0.0):
        deadline = time.perf_counter() + max(0.0, float(timeout))
        while True:
            self.update()
            if all(sample is not None for sample in self.latest):
                return self.latest[0], self.latest[1]
            if time.perf_counter() >= deadline:
                return self.latest[0], self.latest[1]
            time.sleep(0.005)

    def get_forces(self, timeout=0.0):
        left, right = self.get_samples(timeout=timeout)
        left_force = None if left is None else left["fnormal"]
        right_force = None if right is None else right["fnormal"]
        return left_force, right_force

    def stop(self):
        for event in self.stop_events:
            event.set()

        for process in self.processes:
            process.join(timeout=2.0)
            if process.is_alive():
                process.terminate()
                process.join(timeout=1.0)

        self.processes = []
        self.stop_events = []
        self.result_queues = []

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, exc_type, exc, traceback):
        self.stop()


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--video1", "-v1", type=str, default="0")
    parser.add_argument("--video2", "-v2", type=str, default="1")
    parser.add_argument("--config1", type=str, default="./config/ddjx01.json")
    parser.add_argument("--config2", type=str, default="./config/ddjx01.json")
    parser.add_argument("--rotation1", "-r1", type=str, default="./config/rotation_config_0.json")
    parser.add_argument("--rotation2", "-r2", type=str, default="./config/rotation_config_1.json")
    parser.add_argument("--target-fps", type=float, default=30)
    parser.add_argument("--motion-threshold", type=float, default=0)
    parser.add_argument("--cpu", action="store_true")
    args = parser.parse_args()

    reader = TwoSensorForceReader(
        video1=args.video1,
        video2=args.video2,
        config1=args.config1,
        config2=args.config2,
        rotation1=args.rotation1,
        rotation2=args.rotation2,
        target_fps=args.target_fps,
        cuda=not args.cpu,
        motion_threshold=args.motion_threshold,
    )

    print("Starting two-sensor force reader. Press Ctrl+C to stop.")
    reader.start()
    try:
        while True:
            left_sample, right_sample = reader.get_samples(timeout=1.0)
            left_force, right_force = reader.get_forces()
            left_fps = 0.0 if left_sample is None else left_sample["fps"]
            right_fps = 0.0 if right_sample is None else right_sample["fps"]
            left_contact = False if left_sample is None else left_sample["is_contact"]
            right_contact = False if right_sample is None else right_sample["is_contact"]

            print(
                "left={:>10} right={:>10} fps=({:5.1f}, {:5.1f}) contact=({}, {})".format(
                    "None" if left_force is None else f"{left_force:.4f}",
                    "None" if right_force is None else f"{right_force:.4f}",
                    left_fps,
                    right_fps,
                    int(left_contact),
                    int(right_contact),
                )
            )
            time.sleep(0.05)
    except KeyboardInterrupt:
        print("\nStopping.")
    finally:
        reader.stop()


if __name__ == "__main__":
    mp.freeze_support()
    main()
