"""Crop any n-dimensional array."""

import numpy as np
from ..input import mouse_loop

class Crop(object):
    """A callback class that will return the input array cropped to the output size. N-dimensional."""

    def __init__(self, output_size=(64, 64, 3), center=None):
        """
        Create the cropping callback class.

        :param output_size: Specified the size the input should be cropped to. Can be redefined later.
        :param center: Specifies the center on the input array to take the crop out of.
        """
        self._output_size = None
        self._center = np.asarray([o // 2 for o in output_size])
        self.odd_center = None
        self.mouse_control = None
        self.input_size = None

        self.center = center
        self.output_size = output_size

    @property
    def output_size(self):
        """Get the output size."""
        return self._output_size

    @output_size.setter
    def output_size(self, set):
        """Set the output size."""
        self._output_size = set
        if self._output_size is not None:
            self._output_size = np.asarray(set)

    @property
    def center(self):
        """Get the center."""
        return self._center

    @center.setter
    def center(self, set):
        """Set the center. Guarded so that colors need not be set."""
        if set is not None:
            for x in range(len(set)):
                self._center[x] = set[x]

    def __call__(self, arr):
        """Crop the input array to the specified output size. output is centered on self.center point on input."""
        self.input_size = arr.shape
        if self.center is None:
            self.center = [int(arr.shape[x]) // 2 for x in range(arr.ndim)]
        self.odd_out = np.array(
            [self.output_size[x] % 2 for x in range(len(self.output_size))]
        )
        self.odd_center = np.array(
            [self.center[x] % 2 for x in range(len(self.center))]
        )

        center = self.center.copy()  # stop opencv from thread breaking us
        top_left_get = [
            min(max(0, center[x] - self.output_size[x] // 2), arr.shape[x] - 1)
            for x in range(arr.ndim)
        ]
        bottom_right_get = [
            min(
                max(0, center[x] + self.output_size[x] // 2 + self.odd_out[x]),
                arr.shape[x],
            )
        ]

        top_left_put = [
            min(
                max(
                    0,
                    -(bottom_right_get[x] - center[x] - self.output_size[x] // 2)
                    + self.odd_out[x],
                ),
                self.output_size[x],
            )
        ]
        bottom_right_put = [
            min(
                max(0, top_left_put[x] + (bottom_right_get[x] - top_left_get[x])),
            )
        ]
        get_slices = [slice(x1, x2) for x1, x2 in zip(top_left_get, bottom_right_get)]
        get_slices = tuple(get_slices)
        put_slices = [slice(x1, x2) for x1, x2 in zip(top_left_put, bottom_right_put)]
        put_slices = tuple(put_slices)
        out_array = np.zeros(self.output_size)
        out_array[put_slices] = arr[get_slices]
        return out_array.astype(arr.dtype)

    def enable_mouse_control(self):
        """Move the mouse to move where the crop is from on the original image."""

        @mouse_loop
        def m_loop(me):
            if self.center is None:
                self.center = [0, 0, 1]
            self.center[:] = [
                int(me.y / self.output_size[0] * self.input_size[0]),
                int(me.x / self.output_size[1] * self.input_size[1]),
                1,
            ]

        self.mouse_control = m_loop
        return self

"""OpenCV windows that will display the arrays."""

import warnings
from threading import Thread
from typing import List, Union, Callable, Any, Dict, Iterable, Optional

import cv2
from localpubsub import NoData

from displayarray.callbacks import global_cv_display_callback
from displayarray._uid import uid_for_source
from displayarray.frame import subscriber_dictionary
from displayarray.frame.frame_updater import FrameCallable
from displayarray.frame.frame_updater import FrameUpdater
from displayarray.input import MouseEvent
from displayarray.window import window_commands
from displayarray._util import WeakMethod
from displayarray.effects.select_channels import SelectChannels

class SubscriberWindows(object):
    """Windows that subscribe to updates to cameras, videos, and arrays."""

    FRAME_DICT: Dict[str, np.ndarray] = {}
    ESC_KEY_CODES = [27]  # ESC key on most keyboards

    def __init__(
        self,
        window_names: Iterable[str] = ("displayarray",),
        video_sources: Iterable[Union[str, int]] = (0,),
        callbacks: Optional[List[Callable[[np.ndarray], Any]]] = None,
        silent: bool = False,
    ):
        """Create the array displaying window."""
        self.source_names: List[Union[str, int]] = []
        self.close_threads: Optional[List[Thread]] = []
        self.frames: List[np.ndarray] = []
        self.input_vid_global_names: List[str] = []
        self.window_names: List[str] = []
        self.input_cams: List[str] = []
        self.exited = False
        self.silent = silent

        if callbacks is None:
            callbacks = []
        for name in video_sources:
            self.add_source(name)
        self.callbacks = callbacks
        if not self.silent:
            for name in window_names:
                self.add_window(name)

        self.update()

    def __bool__(self):
        return not self.exited

    def __iter__(self):
        while not self.exited:
            self.update()
            yield self.frames

    def block(self):
        """Update the window continuously while blocking the outer program."""
        self.loop()
        for ct in self.close_threads:
            ct.join()

    def add_source(self, name):
        """Add another source for this class to display."""
        uid = uid_for_source(name)
        self.source_names.append(uid)
        self.input_vid_global_names.append(uid)
        self.input_cams.append(name)

    def add_window(self, name):
        """Add another window for this class to display sources with. The name will be the title."""
        self.window_names.append(name)
        cv2.namedWindow(name + " (press ESC to quit)")
        m = WeakMethod(self.handle_mouse)
        cv2.setMouseCallback(name + " (press ESC to quit)", m)

    def add_callback(self, callback):
        """Add a callback for this class to apply to videos."""
        self.callbacks.append(callback)

    def __stop_all_cams(self):
        for c in self.source_names:
            subscriber_dictionary.stop_cam(c)

    def handle_keys(self, key_input: int):
        """Capture key input for the escape function and passing to key control subscriber threads."""
        if key_input in self.ESC_KEY_CODES:
            for name in self.window_names:
                cv2.destroyWindow(name + " (press ESC to quit)")
            self.exited = True
            window_commands.quit()
            self.__stop_all_cams()
            return "quit"
        elif key_input not in [-1, 0]:
            try:
                window_commands.key_pub.publish(chr(key_input))
            except ValueError:
                warnings.warn(
                    RuntimeWarning(
                        f"Unknown key code: [{key_input}]. Please report to the displayarray issue page."
                    )
                )

    def handle_mouse(self, event, x, y, flags, param):
        """Capture mouse input for mouse control subscriber threads."""
        mousey = MouseEvent(event, x, y, flags, param)
        window_commands.mouse_pub.publish(mousey)

    def display_frames(self, frames, win_num=0, ids=None):
        if isinstance(frames, Exception):
            raise frames
        for f in range(len(frames)):
            # detect nested:
            if (
                isinstance(frames[f], (list, tuple))
                or frames[f].dtype.num == 17
                or (
                    len(frames[f].shape) != 2
                    and (len(frames[f].shape) != 3 or frames[f].shape[-1] != 3)
                )
            ):
                win_num = self.display_frames(frames[f], win_num, ids)
            else:
                if len(self.window_names) <= win_num:
                    self.add_window(str(win_num))
                cv2.imshow(
                    self.window_names[win_num] + " (press ESC to quit)", frames[f]
                )
                win_num += 1
        return win_num

    def __check_too_many_channels(self):
        for f in range(len(self.frames)):
            if isinstance(self.frames[f], Exception):
                raise self.frames[f]
            if (
                self.frames[f].shape[-1] not in [1, 3]
                and len(self.frames[f].shape) != 2
            ):
                print(
                    f"Too many channels in output. (Got {self.frames[f].shape[-1]} instead of 1 or 3.) "
                    f"Frame selection callback added."
                )
                print(
                    "Ctrl+scroll to change first channel.\n"
                    "Shift+scroll to change second channel.\n"
                    "Alt+scroll to change third channel."
                )
                sel = SelectChannels()
                sel.enable_mouse_control()
                sel.mouse_print_channels = True
                self.callbacks.append(sel)
                for fr in range(len(self.frames)):
                    self.frames[fr] = self.callbacks[-1](self.frames[fr])
                break

    def update_frames(self):
        """Update the windows with the newest data for all frames."""
        self.frames = []
        for i in range(len(self.input_vid_global_names)):
            if self.input_vid_global_names[i] in self.FRAME_DICT and not isinstance(
                self.FRAME_DICT[self.input_vid_global_names[i]], NoData
            ):
                self.frames.append(self.FRAME_DICT[self.input_vid_global_names[i]])
                if isinstance(self.frames, np.ndarray) and len(self.frames.shape) <= 3:
                    self.frames = [self.frames]
                if len(self.callbacks) > 0:
                    for c in self.callbacks:
                        frame = c(self.frames[-1])
                        if frame is not None:
                            self.frames[-1] = frame
                if not self.silent:
                    self.__check_too_many_channels()
                self.FRAME_DICT[self.input_vid_global_names[i]] = NoData()
            self.display_frames(self.frames)

    def update(self, arr: np.ndarray = None, id: str = None):
        """Update window frames once. Optionally add a new input and input id."""
        if arr is not None and id is not None:
            global_cv_display_callback(arr, id)
            if id not in self.input_cams:
                self.add_source(id)
                    self.add_window(id)
        sub_cmd = window_commands.win_cmd_sub()
        self.update_frames()
        msg_cmd = sub_cmd.get()
        key = self.handle_keys(cv2.waitKey(1))
        return msg_cmd, key

    def wait_for_init(self):
        """Update window frames in a loop until they're actually updated. Useful for waiting for cameras to init."""
        msg_cmd = ""
        key = ""
        while msg_cmd != "quit" and key != "quit" and len(self.frames) == 0:
            msg_cmd, key = self.update()

    def end(self):
        """Close all threads. Should be used with non-blocking mode."""
        window_commands.quit(force_all_read=False)
        self.__stop_all_cams()
        for t in self.close_threads:
            t.join()

    def __enter__(self):

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end()

    def __del__(self):

    def __delete__(self, instance):

    def loop(self):
        """Continually update window frame. OpenCV only allows this in the main thread."""
        key = ""
        while msg_cmd != "quit" and key != "quit":
        sub_cmd.release()

def _get_video_callback_dict_threads(
    *vids,
    callbacks: Optional[Dict[Any, Union[FrameCallable, List[FrameCallable]]]] = None,
    fps=float("inf"),
    size=(-1, -1),
):
    assert callbacks is not None
    vid_threads = []
    for v in vids:
        v_name = uid_for_source(v)
        v_callbacks: List[Callable[[np.ndarray], Any]] = []
        if v_name in callbacks:
            if isinstance(callbacks[v_name], List):
                v_callbacks.extend(callbacks[v_name])  # type: ignore
            elif callable(callbacks[v_name]):
                v_callbacks.append(callbacks[v_name])  # type: ignore
        if v in callbacks:
            if isinstance(callbacks[v], List):
                v_callbacks.extend(callbacks[v])  # type: ignore
            elif callable(callbacks[v]):
                v_callbacks.append(callbacks[v])  # type: ignore
        vid_threads.append(
            FrameUpdater(v, callbacks=v_callbacks, fps_limit=fps, request_size=size)
        )
    return vid_threads

def _get_video_threads(
    *vids,
    callbacks: Optional[
        Union[
            Dict[Any, Union[FrameCallable, List[FrameCallable]]],
            List[FrameCallable],
            FrameCallable,
        ]
    ] = None,
):
    vid_threads: List[Thread] = []
    if isinstance(callbacks, Dict):
        vid_threads = _get_video_callback_dict_threads(
            *vids, callbacks=callbacks, fps=fps, size=size
        )
    elif isinstance(callbacks, List):
        for v in vids:
            vid_threads.append(
                FrameUpdater(v, callbacks=callbacks, fps_limit=fps, request_size=size)
            )
    elif callable(callbacks):
                FrameUpdater(v, callbacks=[callbacks], fps_limit=fps, request_size=size)
            )
    else:
            if v is not None:
                vid_threads.append(FrameUpdater(v, fps_limit=fps, request_size=size))

def display(
    *vids,
        Union[
        ]
    ] = None,
    window_names=None,
    blocking=False,
    fps_limit=float("inf"),
    silent=False,
):
    """
    Display all the arrays, cameras, and videos passed in.

    callbacks can be a dictionary linking functions to videos, or a list of function or functions operating on the video
     data before displaying.
    Window names end up becoming the title of the windows
    """
    vid_threads = _get_video_threads(
        *vids, callbacks=callbacks, fps=fps_limit, size=size
    )
    for v in vid_threads:
        v.start()
    if window_names is None:
        window_names = ["window {}".format(i) for i in range(len(vids))]
    if blocking:
        SubscriberWindows(
            window_names=window_names, video_sources=vids, silent=silent
        ).loop()
        for vt in vid_threads:
            vt.join()
    else:
        s = SubscriberWindows(
        )
        s.close_threads = vid_threads
        return s

def breakpoint_display(*args, **kwargs):
    """Display all the arrays, cameras, and videos passed in. Stops code execution until the window is closed."""
    return display(*args, **kwargs, blocking=True)

def read_updates(*args, **kwargs):
    """Read back all frame updates and yield a list of frames. List is empty if no frames were read."""
    return display(*args, **kwargs, silent=True)

"""Get and handle updated frames."""

import threading
import asyncio
from typing import Union, Tuple, Any, Callable, List, Optional, Dict

from displayarray.frame.frame_publishing import pub_cam_thread

FrameCallable = Callable[[np.ndarray], Optional[np.ndarray]]

class FrameUpdater(threading.Thread):
    """Thread for updating frames from a video source."""

        self,
        video_source: Union[int, str, np.ndarray] = 0,
        callbacks: Optional[Union[List[FrameCallable], FrameCallable]] = None,
        request_size: Tuple[int, int] = (-1, -1),
        high_speed: bool = True,
        fps_limit: float = float("inf"),
    ):
        """Create the frame updater thread."""
        super(FrameUpdater, self).__init__(target=self.loop, args=())
        self.cam_id = uid_for_source(video_source)
        self.video_source = video_source
        if callable(callbacks):
            self.callbacks = [callbacks]
        else:
            self.callbacks = callbacks
        self.request_size = request_size
        self.high_speed = high_speed
        self.fps_limit = fps_limit
        self.exception_raised = None

    def __wait_for_cam_id(self):
        while str(self.cam_id) not in subscriber_dictionary.CV_CAMS_DICT:
            continue

    def __apply_callbacks_to_frame(self, frame):
        if frame is not None and not isinstance(frame, NoData):
            try:
                for c in self.callbacks:
                    frame_c = c(frame)
                    if frame_c is not None:
                        frame = frame_c
                if frame.shape[-1] not in [1, 3] and len(frame.shape) != 2:
                    print(
                        f"Too many channels in output. (Got {frame.shape[-1]} instead of 1 or 3.) "
                        f"Frame selection callback added."
                    )
                    print(
                        "Ctrl+scroll to change first channel.\n"
                        "Shift+scroll to change second channel.\n"
                        "Alt+scroll to change third channel."
                    )
                    sel = SelectChannels()
                    sel.enable_mouse_control()
                    sel.mouse_print_channels = True
                    self.callbacks.append(sel)
                    frame = self.callbacks[-1](frame)
            except Exception as e:
                self.exception_raised = e
                frame = self.exception_raised
                subscriber_dictionary.stop_cam(self.cam_id)
                window_commands.quit()
                raise e
            global_cv_display_callback(frame, self.cam_id)

        """Continually get frames from the video publisher, run callbacks on them, and listen to commands."""
        t = pub_cam_thread(
            self.video_source, self.request_size, self.high_speed, self.fps_limit
        )
        self.__wait_for_cam_id()

        sub_cam = subscriber_dictionary.cam_frame_sub(str(self.cam_id))
        sub_owner = subscriber_dictionary.handler_cmd_sub(str(self.cam_id))
        msg_owner = sub_owner.return_on_no_data = ""
        while msg_owner != "quit":
            frame = sub_cam.get(blocking=True, timeout=1.0)  # type: np.ndarray
            self.__apply_callbacks_to_frame(frame)
            msg_owner = sub_owner.get()
        sub_owner.release()
        sub_cam.release()
        subscriber_dictionary.stop_cam(self.cam_id)
        t.join()

    def display(self, callbacks: List[Callable[[np.ndarray], Any]] = None):
        """
        Start default display operation.

        For multiple video sources, please use something outside of this class.

        :param callbacks: List of callbacks to be run on frames before displaying to the screen.
        """
        from displayarray.window import SubscriberWindows

        self.start()
        SubscriberWindows(video_sources=[self.cam_id], callbacks=callbacks).loop()
        self.join()
        if self.exception_raised is not None:
            raise self.exception_raised

'''async def read_updates(
    *vids,
        Union[
        ]
    ] = None,
    end_callback: Callable[[], bool] = lambda: False,
    blocking=True,
):
    """
    Read back all updates from the requested videos.

    Examp#le usage:

    .. co#de-block:: python

      >>#> from examples.videos import test_video
      >>#> f = 0
      >>#> for f, r in enumerate(read_updates(test_video, end_callback=lambda :f==2)):
      ..#.   print(f"Frame:{f}. Array:{r}")

    """
    from displayarray.window import SubscriberWindows
    from displayarray.window.subscriber_windows import _get_video_threads

    vid_names = [uid_for_source(name) for name in vids]
    )
        v.start()

    while not end_callback():
        vid_update_dict = {}
        dict_was_updated = False
        for i in range(len(vid_names)):
            if vid_names[i] in SubscriberWindows.FRAME_DICT and not isinstance(
                SubscriberWindows.FRAME_DICT[vid_names[i]], NoData
            ):
                vid_update_dict[vid_names[i]] = SubscriberWindows.FRAME_DICT[
                    vid_names[i]
                ]
                if (
                    isinstance(vid_update_dict[vid_names[i]], np.ndarray)
                    and len(vid_update_dict[vid_names[i]].shape) <= 3
                ):
                    vid_update_dict[vid_names[i]] = [vid_update_dict[vid_names[i]]]
                dict_was_updated = True
        if dict_was_updated or not blocking:
            yield vid_update_dict
        await asyncio.sleep(0)
    for v in vid_names:
        subscriber_dictionary.stop_cam(v)
        v.join()'''

async def read_updates_zero_mq(
    *topic_names,
    address: str = "tcp://127.0.0.1:5600",
    flags: int = 0,
    copy: bool = True,
    track: bool = False,
    blocking: bool = False,
    end_callback: Callable[[Any], bool] = lambda x: False,
):
    """Read updated frames from ZeroMQ."""
    import zmq

    ctx = zmq.Context()
    s = ctx.socket(zmq.SUB)
    s.connect(address)
    if not blocking:
        flags |= zmq.NOBLOCK

    for topic in topic_names:
        s.setsockopt(zmq.SUBSCRIBE, topic)
    cb_val = False
    while not cb_val:
        try:
            md = s.recv_json(flags=flags)
            msg = s.recv(flags=flags, copy=copy, track=track)
            buf = memoryview(msg)
            arr = np.frombuffer(buf, dtype=md["dtype"])
            arr.reshape(md["shape"])
            name = md["name"]
            cb_val = end_callback(md)
            yield name, arr
        except zmq.ZMQError as e:
            if isinstance(e, zmq.Again):
                pass  # no messages to receive
            else:
                raise e
        finally:
            await asyncio.sleep(0)

async def read_updates_ros(
    dtypes=None,
    listener_node_name=None,
    poll_rate_hz=None,
):
    """Read updated frames from ROS."""
    import rospy
    from rospy.numpy_msg import numpy_msg
    from rospy.client import _WFM
    import std_msgs.msg
    import random
    import string

    if dtypes is None:
        raise ValueError(
            "ROS cannot automatically determine the types of incoming numpy arrays. Please specify.\n"
            "Options are: \n"
            "\tfloat32, float64, bool, char, int16, "
            "\tint32, int64, str, uint16, uint32, uint64, uint8"
        )

    if listener_node_name is None:
        # https://stackoverflow.com/a/2257449
        listener_node_name = "".join(
            random.choices(string.ascii_uppercase + string.digits, k=8)
        )

    rospy.init_node(listener_node_name)

    msg_types = [
        {
            np.float32: std_msgs.msg.Float32(),
            np.float64: std_msgs.msg.Float64(),
            np.bool: std_msgs.msg.Bool(),
            np.char: std_msgs.msg.Char(),
            np.int16: std_msgs.msg.Int16(),
            np.int32: std_msgs.msg.Int32(),
            np.int64: std_msgs.msg.Int64(),
            np.str: std_msgs.msg.String(),
            np.uint16: std_msgs.msg.UInt16(),
            np.uint32: std_msgs.msg.UInt32(),
            np.uint64: std_msgs.msg.UInt64(),
            np.uint8: std_msgs.msg.UInt8(),
            "float32": std_msgs.msg.Float32(),
            "float64": std_msgs.msg.Float64(),
            "bool": std_msgs.msg.Bool(),
            "char": std_msgs.msg.Char(),
            "int16": std_msgs.msg.Int16(),
            "int32": std_msgs.msg.Int32(),
            "int64": std_msgs.msg.Int64(),
            "str": std_msgs.msg.String(),
            "uint16": std_msgs.msg.UInt16(),
            "uint32": std_msgs.msg.UInt32(),
            "uint64": std_msgs.msg.UInt64(),
            "uint8": std_msgs.msg.UInt8(),
        }.get(
            dtype, dtype
        )  # allow users to use their own custom messages in numpy arrays
        for dtype in dtypes
    ]
    s = None
    try:
        wfms = {t: _WFM() for t in topic_names}
        s = {
            t: rospy.Subscriber(t, numpy_msg(msg_types[i]), wfms[t].cb)
            for i, t in enumerate(topic_names)
        }
        while not cb_val:
            while not rospy.core.is_shutdown():
                if poll_rate_hz:
                    await asyncio.sleep(1.0 / poll_rate_hz)
                else:
                    await asyncio.sleep(0)
                for t, w in wfms.items():
                    if w.msg is not None:
                        yield t, w.msg
                        cb_val = end_callback(w.msg)
                        w.msg = None
    except KeyboardInterrupt:
        pass
    finally:
        if s is not None:
            for _, sub in s.items():
                sub.unregister()
    if rospy.core.is_shutdown():
        raise rospy.exceptions.ROSInterruptException("rospy shutdown")

"""
DisplayArray.

Display NumPy arrays.

Usage:
  displayarray (-w <webcam-number> | -v <video-filename> | -t <topic-name>[,dtype])... [-m <msg-backend>]
  displayarray -h
  displayarray --version

Options:
  -h, --help                                           Show this help text.
  --version                                            Show version number.
  -w <webcam-number>, --webcam=<webcam-number>         Display video from a webcam.
  -v <video-filename>, --video=<video-filename>        Display frames from a video file.
  -t <topic-name>, --topic=<topic-name>                Display frames from a topic using the chosen message broker.
  -m <msg-backend>, --message-backend <msg-backend>    Choose message broker backend. [Default: ROS]
                                                       Currently supported: ROS, ZeroMQ
  --ros                                                Use ROS as the backend message broker.
  --zeromq                                             Use ZeroMQ as the backend message broker.
"""

from docopt import docopt

def main(argv=None):
    """Process command line arguments."""
    arguments = docopt(__doc__, argv=argv)
    if arguments["--version"]:
        from displayarray import __version__

        print(f"DisplayArray V{__version__}")
        return
    from displayarray import display

    vids = [int(w) for w in arguments["--webcam"]] + arguments["--video"]
    v_disps = None
    if vids:
        v_disps = display(*vids, blocking=False)
    from displayarray.frame.frame_updater import read_updates_ros, read_updates_zero_mq

    topics = arguments["--topic"]
    topics_split = [t.split(",") for t in topics]
    d = display()

    async def msg_recv():
        nonlocal d
        while d:
            if arguments["--message-backend"] == "ROS":
                async for v_name, frame in read_updates_ros(
                    [t for t, d in topics_split], [d for t, d in topics_split]
                ):
                    d.update(arr=frame, id=v_name)
            if arguments["--message-backend"] == "ZeroMQ":
                async for v_name, frame in read_updates_zero_mq(
                    *[bytes(t, encoding="ascii") for t in topics]
                ):

    async def update_vids():
        while v_disps:
            if v_disps:
                v_disps.update()
                await asyncio.sleep(0)

    async def runner():
        await asyncio.wait([msg_recv(), update_vids()])

    loop = asyncio.get_event_loop()
    loop.run_until_complete(runner())
    loop.close()

if __name__ == "__main__":
    main()

"""Get camera IDs."""

import cv2

from typing import List

def get_cam_ids() -> List[int]:
    """Get all cameras that OpenCV can currently detect."""
    cam_list: List[int] = []

    while True:
        cam = cv2.VideoCapture(len(cam_list))
        if not cam.isOpened():
            break
        cam_list.append(len(cam_list))

    return cam_list

"""Allow OpenCV to handle numpy arrays as input."""

import cv2

class NpCam(object):
    """Add OpenCV camera controls to a numpy array."""

    def __init__(self, img):
        """Create a fake camera for OpenCV based on the initial array."""
        assert isinstance(img, np.ndarray)
        self.__img = img
        self.__is_opened = True
        if len(img.shape) > 0:
            self.__height = self.__img.shape[0]
            if len(self.__img.shape) > 1:
                self.__width = self.__img.shape[1]
            else:
                self.__width = self.__height
        else:
            self.__width = self.__height = 1
        self.__ratio = float(self.__width) / self.__height

        self.__wait_for_ratio = False

    def __handler_ratio(self):
        if self.__width <= 0 or not isinstance(self.__width, int):
            self.__width = int(self.__ratio * self.__height)
        elif self.__height <= 0 or not isinstance(self.__height, int):
            self.__height = int(self.__width / self.__ratio)
        if self.__width > 0 and self.__height > 0:
            self.__img = cv2.resize(self.__img, (self.__width, self.__height))

    def set(self, *args, **kwargs):
        """Set CAP_PROP_FRAME_WIDTH or CAP_PROP_FRAME_HEIGHT to scale a numpy array to that size."""
        if args[0] in [cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT]:
            self.__wait_for_ratio = not self.__wait_for_ratio
            if args[0] == cv2.CAP_PROP_FRAME_WIDTH:
                self.__width = args[1]
            else:
                self.__height = args[1]
            if not self.__wait_for_ratio:
                self.__handler_ratio()

    @staticmethod
    def get(*args, **kwargs):
        """Get OpenCV args. Currently only a fake CAP_PROP_FRAME_COUNT to fix detecting video ends."""
        if args[0] == cv2.CAP_PROP_FRAME_COUNT:
            return float("inf")

    def read(self):
        """Read back the numpy array in standard "did it work", "the array", OpenCV format."""
        return True, self.__img

    def isOpened(self):  # NOSONAR
        """Hack to tell OpenCV we're opened until we call release."""
        return self.__is_opened

    def release(self):
        """Let OpenCV know we're finished."""
        self.__is_opened = False

"""Publish frames so any function within this program can find them."""

import time

import cv2

from .np_to_opencv import NpCam

from typing import Union, Tuple, Optional, Dict, Any, List, Callable

def pub_cam_loop(
    cam_id: Union[int, str, np.ndarray],
    request_size: Tuple[int, int] = (-1, -1),
    high_speed: bool = True,
    fps_limit: float = float("inf"),
) -> bool:
    """
    Publish whichever camera you select to CVCams.<cam_id>.Vid.

    You can send a quit command 'quit' to CVCams.<cam_id>.Cmd
    Status information, such as failure to open, will be posted to CVCams.<cam_id>.Status

    :param high_speed: Selects mjpeg transferring, which most cameras seem to support, so speed isn't limited
    :param fps_limit: Limits the frames per second.
    :param cam_id: An integer representing which webcam to use, or a string representing a video file.
    :param request_size: A tuple with width, then height, to request the video size.
    :return: True if loop ended normally, False if it failed somehow.
    """
    name = uid_for_source(cam_id)

    if isinstance(cam_id, (int, str)):
        cam: Union[NpCam, cv2.VideoCapture] = cv2.VideoCapture(cam_id)
    elif isinstance(cam_id, np.ndarray):
        cam = NpCam(cam_id)
    else:
        raise TypeError(
            "Only strings or ints representing cameras, or numpy arrays representing pictures supported."
        )

    subscriber_dictionary.register_cam(name)

    frame_counter = 0

    sub = subscriber_dictionary.cam_cmd_sub(name)
    sub.return_on_no_data = ""
    msg = ""

    if high_speed:
        cam.set(cv2.CAP_PROP_FOURCC, cv2.CAP_OPENCV_MJPEG)

    cam.set(cv2.CAP_PROP_FRAME_WIDTH, request_size[0])
    cam.set(cv2.CAP_PROP_FRAME_HEIGHT, request_size[1])

    if not cam.isOpened():
        subscriber_dictionary.CV_CAMS_DICT[name].status_pub.publish("failed")
        return False
    now = time.time()
    while msg != "quit":
        time.sleep(1.0 / (fps_limit - (time.time() - now)))
        now = time.time()
        (ret, frame) = cam.read()  # type: Tuple[bool, np.ndarray ]
        if ret is False or not isinstance(frame, np.ndarray):
            cam.release()
            subscriber_dictionary.CV_CAMS_DICT[name].status_pub.publish("failed")
            return False
        if cam.get(cv2.CAP_PROP_FRAME_COUNT) > 0:
            frame_counter += 1
            if frame_counter >= cam.get(cv2.CAP_PROP_FRAME_COUNT):
                frame_counter = 0
                cam = cv2.VideoCapture(cam_id)
        subscriber_dictionary.CV_CAMS_DICT[name].frame_pub.publish(frame)
        msg = sub.get()
    sub.release()

    cam.release()
    return True

def pub_cam_thread(
    cam_id: Union[int, str],
    request_ize: Tuple[int, int] = (-1, -1),
) -> threading.Thread:
    """Run pub_cam_loop in a new thread. Starts on creation."""
    t = threading.Thread(
        target=pub_cam_loop, args=(cam_id, request_ize, high_speed, fps_limit)
    )
    t.start()
    return t

async def publish_updates_zero_mq(
    *vids,
        Union[Dict[Any, FrameCallable], List[FrameCallable], FrameCallable]
    ] = None,
    publishing_address="tcp://127.0.0.1:5600",
    prepend_topic="",
    flags=0,
    copy=True,
    track=False
):
    """Publish frames to ZeroMQ when they're updated."""
    import zmq
    from displayarray import read_updates

    s = ctx.socket(zmq.PUB)
    s.bind(publishing_address)

    try:
        for v in read_updates(vids, callbacks, fps_limit, size, end_callback, blocking):
            if v:
                for vid_name, frame in v.items():
                    md = dict(
                        dtype=str(frame.dtype),
                        shape=frame.shape,
                        name=prepend_topic + vid_name,
                    )
                    s.send_json(md, flags | zmq.SNDMORE)
                    s.send(frame, flags, copy=copy, track=track)
            if fps_limit:
                await asyncio.sleep(1.0 / fps_limit)
            else:
        pass
    finally:
        vid_names = [uid_for_source(name) for name in vids]
        for v in vid_names:
            subscriber_dictionary.stop_cam(v)

async def publish_updates_ros(
    *vids,
    ] = None,
    node_name="displayarray",
    publisher_name="npy",
    rate_hz=None,
    dtype=None
):
    """Publish frames to ROS when they're updated."""

    def get_msg_type(dtype):
        if dtype is None:
            msg_type = {
                np.float32: std_msgs.msg.Float32(),
                np.float64: std_msgs.msg.Float64(),
                np.bool: std_msgs.msg.Bool(),
                np.char: std_msgs.msg.Char(),
                np.int16: std_msgs.msg.Int16(),
                np.int32: std_msgs.msg.Int32(),
                np.int64: std_msgs.msg.Int64(),
                np.str: std_msgs.msg.String(),
                np.uint16: std_msgs.msg.UInt16(),
                np.uint32: std_msgs.msg.UInt32(),
                np.uint64: std_msgs.msg.UInt64(),
                np.uint8: std_msgs.msg.UInt8(),
            }[dtype]
        else:
            msg_type = (
                dtype  # allow users to use their own custom messages in numpy arrays
            )
        return msg_type

    publishers: Dict[str, rospy.Publisher] = {}
    rospy.init_node(node_name, anonymous=True)
    try:
            if v:
                if rospy.is_shutdown():
                    break
                    if vid_name not in publishers:
                        dty = frame.dtype if dtype is None else dtype
                        publishers[vid_name] = rospy.Publisher(
                            publisher_name + vid_name,
                            numpy_msg(get_msg_type(dty)),
                            queue_size=10,
                        )
                    publishers[vid_name].publish(frame)
            if rate_hz:
                await asyncio.sleep(1.0 / rate_hz)
            else:
        pass
    finally:

"""Decorators for creating input loops that OpenCV handles."""

from typing import Callable

class MouseEvent(object):
    """Holds all the OpenCV mouse event information."""

    def __init__(self, event, x, y, flags, param):
        """Create an OpenCV mouse event."""
        self.event = event
        self.x = x
        self.y = y
        self.flags = flags
        self.param = param

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return "event:{}\nx,y:{},{}\nflags:{}\nparam:{}\n".format(
            self.event, self.x, self.y, self.flags, self.param
        )

class _mouse_thread(object):  # NOSONAR
    """Run a function on mouse information that is received by the window."""

    def __init__(self, f):
        self.f = f
        self.sub_mouse = window_commands.mouse_pub.make_sub()

    def __call__(self, *args, **kwargs):
        """Call the function this was set up with."""
        self.f(self.sub_mouse, *args, **kwargs)

class _mouse_loop_thread(object):  # NOSONAR
    """Run a function on mouse information that is received by the window, in the main loop."""

    def __init__(self, f, run_when_no_events=False, fps=60):
        self.f = f
        self.sub_cmd = window_commands.win_cmd_pub.make_sub()
        self.sub_cmd.return_on_no_data = ""
        self.run_when_no_events = run_when_no_events
        self.fps = fps

        """Run the function this was set up with in a loop."""
        while msg_cmd != "quit":
            mouse_xyzclick = self.sub_mouse.get(blocking=True)  # type: MouseEvent
            if mouse_xyzclick is not self.sub_mouse.return_on_no_data:
                self.f(mouse_xyzclick, *args, **kwargs)
            elif self.run_when_no_events:
                self.f(None, *args, **kwargs)
            msg_cmd = self.sub_cmd.get()
            time.sleep(1.0 / self.fps)

class mouse_loop(object):  # NOSONAR
    """
    Run a function on mouse information that is received by the window, continuously in a new thread.

    .. code-block:: python

      >>> @mouse_loop
      ... def fun(mouse_event):
      ...   print("x:{}, y:{}".format(mouse_event.x, mouse_event.y))
    """

        """Start a new mouse thread for the decorated function."""
        self.t = threading.Thread(
            target=_mouse_loop_thread(f, run_when_no_events=False)
        )
        self.t.start()

        """Return the thread that was started with the function passed in."""
        return self.t

class _key_thread(object):  # NOSONAR

        self.f = f
        self.sub_key = window_commands.key_pub.make_sub()

        self.f(self.sub_key, *args, **kwargs)

class _key_loop_thread(object):  # NOSONAR

        self.f = f

            key_chr = self.sub_key.get()  # type: chr
            if key_chr is not self.sub_key.return_on_no_data:
                self.f(key_chr, *args, **kwargs)

class key_loop(object):  # NOSONAR
    """

      >>> @key_loop
      ... def fun(key):
      ...   print("key pressed:{}".format(key))
    """

    def __init__(self, f: Callable[[str], None]):
        """Start a new key thread for the decorated function."""
        self.t = threading.Thread(target=_key_loop_thread(f, run_when_no_events=False))

"""Create lens effects. Currently only 2D+color arrays are supported."""

import cv2

try:
    import torch
except ImportError:
    torch = None  # type: ignore

class ControllableLens(object):
    """A lens callback that can be controlled by the program or the user."""

    def __init__(self, use_bleed=False, zoom=1, center=None):
        """Create the lens callback."""
        self._center = center
        self._zoom = zoom
        self.use_bleed = use_bleed
        self.bleed = None

    def _check_setup_bleed(self, arr):
        if not isinstance(self.bleed, np.ndarray) and self.use_bleed:
            self.bleed = np.zeros_like(arr)

    def run_bleed(self, arr, x, y):
        """Spread color outwards, like food coloring in water."""
        arr[y, ...] = (arr[(y + 1) % len(y), ...] + arr[(y - 1) % len(y), ...]) / 2
        arr[:, x, ...] = (
            arr[:, (x + 1) % len(x), ...] + arr[:, (x - 1) % len(x), ...]
        ) / 2

class Barrel(ControllableLens):
    """A barrel lens distortion callback."""

        self, use_bleed=False, barrel_power=1, pincushion_power=1, zoom=1, center=None
    ):
        """Create the barrel lens distortion callback."""
        super().__init__(use_bleed, zoom, center)
        self._barrel_power = barrel_power

    @property
        """Guarded get center. Limits to within input."""
        if self.input_size is not None:
            self._center[:] = [
                min(max(0, s), self.input_size[i]) for i, s in enumerate(self._center)
            ]

    def center(self, setpoint):
        """Guarded set center. Limits to within input."""
            setpoint = [
                min(max(0, s), self.input_size[i]) for i, s in enumerate(setpoint)
            ]
        self._center[:] = setpoint

    @property
    def zoom(self):
        """Guarded zoom. Avoids divide by zero conditions."""
        if self._zoom == 0:
            return 1e-15
        else:
            return self._zoom

    @property
    def barrel_power(self):
        """Guarded barrel power. Avoids divide by zero conditions."""
        if self._barrel_power == 0:
        else:
            return self._barrel_power

    @barrel_power.setter
    def barrel_power(self, setpoint):
        """Set the barrel power."""
        self._barrel_power = setpoint

    @zoom.setter  # type: ignore
    def zoom(self, setpoint):
        """Set the zoom power."""
        self._zoom = setpoint

    def enable_mouse_control(self, crop_size=None):
        """
        Enable the default mouse controls.

        Move the mouse to center the image
        scroll to increase/decrease barrel
        ctrl+scroll to increase/decrease zoom
        """

            if self.input_size is not None:
                if crop_size is not None:
                    self.center = [
                        me.y * self.input_size[0] / crop_size[0],
                        me.x * self.input_size[1] / crop_size[1],
                    ]
                else:
                    self.center = [me.y, me.x]
                if me.event == cv2.EVENT_MOUSEWHEEL:
                    if me.flags & cv2.EVENT_FLAG_CTRLKEY:
                        if me.flags > 0:
                            self.zoom *= 1.1
                        else:
                            self.zoom /= 1.1
                    else:
                            self.barrel_power *= 1.1
                        else:
                            self.barrel_power /= 1.1
                        print(self.barrel_power)

        """Run the lens distortion algorithm on the input."""
        zoom_out = 1.0 / self.zoom
        self._check_setup_bleed(arr)

        y = np.arange(arr.shape[0])
        x = np.arange(arr.shape[1])
        if self._center is None:
            self._center = [len(y) / 2.0, len(x) / 2.0]
        y2_ = (y - (len(y) / 2.0)) * zoom_out / arr.shape[0]
        x2_ = (x - (len(x) / 2.0)) * zoom_out / arr.shape[1]
        p2 = np.array(np.meshgrid(x2_, y2_))

        cy = self._center[0] / arr.shape[0]
        cx = self._center[1] / arr.shape[1]

        theta = np.arctan2(p2[1], p2[0])

        radius = np.linalg.norm(p2, axis=0, ord=2)

        radius = pow(radius, self.barrel_power)

        x_new = 0.5 * (radius * np.cos(theta) + cx * 2)
        x_new = np.clip(x_new * len(x), 0, len(x) - 1)

        y_new = 0.5 * (radius * np.sin(theta) + cy * 2)
        y_new = np.clip(y_new * len(y), 0, len(y) - 1)

        p = np.array(np.meshgrid(y, x)).astype(np.uint32)

        p_new = np.array((y_new, x_new))
        p_new = p_new.astype(np.uint32)

        if self.use_bleed:
            arr2 = self.bleed.copy()
            self.run_bleed(arr2, x, y)
            arr2[p_new[0], p_new[1], :] = np.swapaxes(arr[p[0], p[1], :], 0, 1)
            self.bleed = arr2
        else:
            arr[p[0], p[1], :] = np.swapaxes(arr[p_new[0], p_new[1], :], 0, 1)

        return arr

class BarrelPyTorch(Barrel):
    """A barrel distortion callback class accelerated by PyTorch."""

        """Run a pytorch accelerated lens distortion algorithm on the input."""
        y = torch.arange(arr.shape[1]).type(torch.FloatTensor).cuda()
        x = torch.arange(arr.shape[0]).type(torch.FloatTensor).cuda()
            self._center = [y.shape[0] / 2.0, x.shape[0] / 2.0]

        y2_ = (y - (y.shape[0] / 2.0)) * zoom_out / arr.shape[1]
        x2_ = (x - (x.shape[0] / 2.0)) * zoom_out / arr.shape[0]
        p2 = torch.stack(torch.meshgrid(x2_, y2_))

        cy = self._center[1] / arr.shape[1]
        cx = self._center[0] / arr.shape[0]

        theta = torch.atan2(p2[1], p2[0])

        radius = torch.norm(p2, dim=0)

        radius = torch.pow(radius, self.barrel_power)

        x_new = 0.5 * (radius * torch.cos(theta) + cx * 2)
        x_new = torch.clamp(x_new * x.shape[0], 0, x.shape[0] - 1)

        y_new = 0.5 * (radius * torch.sin(theta) + cy * 2)
        y_new = torch.clamp(y_new * y.shape[0], 0, y.shape[0] - 1)

        p = torch.stack(torch.meshgrid([x, y])).type(torch.IntTensor)

        p_new = torch.stack((x_new, y_new))
        p_new = p_new.type(torch.IntTensor)

        arr[p[0], p[1], :] = arr[p_new[0], p_new[1], :]

        return arr

class Mustache(ControllableLens):
    """A mustache distortion callback."""

    ):
        """Create the mustache distortion callback."""
        self.zoom = zoom
        self.barrel_power = barrel_power
        self.pincushion_power = pincushion_power

        """
        Enable the default mouse loop to control the mustache distortion.

        ctrl+scroll wheel: zoom in and out
        shift+scroll wheel: adjust pincushion power
        scroll wheel: adjust barrel power
        """

            self.center[:] = [me.y, me.x]
            if me.event == cv2.EVENT_MOUSEWHEEL:
                if me.flags & cv2.EVENT_FLAG_CTRLKEY:
                    if me.flags > 0:
                        self.zoom *= 1.1
                    else:
                        self.zoom /= 1.1
                elif me.flags & cv2.EVENT_FLAG_SHIFTKEY:
                        self.pincushion_power *= 1.1
                    else:
                        self.pincushion_power /= 1.1
                else:
                        self.barrel_power *= 1.1
                    else:
                        self.barrel_power /= 1.1

        """Run the mustache distortion algorithm on the input numpy array."""

            self.center = [len(y) / 2.0, len(x) / 2.0]
        y2_ = (y - self.center[0]) * zoom_out / arr.shape[0]
        x2_ = (x - self.center[1]) * zoom_out / arr.shape[1]

        barrel_power = self.barrel_power
        pincushion_power = self.pincushion_power

        radius = np.linalg.norm(p2, axis=0)
        radius2 = np.linalg.norm(p2, axis=0, ord=4)

        radius = pow(radius, barrel_power)
        radius2 = pow(radius2, pincushion_power)

        x_new = 0.5 * (radius2 * radius * np.cos(theta) + 1)

        y_new = 0.5 * (radius2 * radius * np.sin(theta) + 1)

        p_new = np.array((y_new, x_new)).astype(np.uint32)

        else:
            arr2 = np.zeros_like(arr)

        return arr2

"""Reduce many color images to the three colors that your eyeballs can see."""

import cv2

from typing import Iterable

class SelectChannels(object):
    """
    Select channels to display from an array with too many colors.

    :param selected_channels: the list of channels to display.
    """

    def __init__(self, selected_channels: Iterable[int] = None):
        """Select which channels from the input array to display in the output."""
        if selected_channels is None:
            selected_channels = [0, 0, 0]
        self.selected_channels = selected_channels
        self.mouse_print_channels = False
        self.num_input_channels = None

        """Run the channel selector."""
        self.num_input_channels = arr.shape[-1]
        out_arr = [
            arr[..., min(max(0, x), arr.shape[-1] - 1)] for x in self.selected_channels
        ]
        out_arr = np.stack(out_arr, axis=-1)
        return out_arr

        """
        Enable mouse control.

        Alt+Scroll to increase/decrease channel 2.
        Shift+Scroll to increase/decrease channel 1.
        Ctrl+scroll to increase/decrease channel 0.
        """

                        self.selected_channels[0] += 1
                        self.selected_channels[0] = min(
                            self.selected_channels[0], self.num_input_channels - 1
                        )
                    else:
                        self.selected_channels[0] -= 1
                        self.selected_channels[0] = max(self.selected_channels[0], 0)
                    if self.mouse_print_channels:
                        print(f"Channel 0 now maps to {self.selected_channels[0]}.")
                        self.selected_channels[1] += 1
                        self.selected_channels[1] = min(
                            self.selected_channels[1], self.num_input_channels - 1
                        )
                    else:
                        self.selected_channels[1] -= 1
                        self.selected_channels[1] = max(self.selected_channels[1], 0)
                        print(f"Channel 1 now maps to {self.selected_channels[1]}.")
                elif me.flags & cv2.EVENT_FLAG_ALTKEY:
                        self.selected_channels[2] += 1
                        self.selected_channels[2] = min(
                            self.selected_channels[2], self.num_input_channels - 1
                        )
                    else:
                        self.selected_channels[2] -= 1
                        self.selected_channels[2] = max(self.selected_channels[2], 0)
                        print(f"Channel 2 now maps to {self.selected_channels[2]}.")