        with self.lock:
            return self._get_delete_messages_by_sub_keys(topic_id, sub_keys)

"""

"""

from __future__ import absolute_import, division, print_function, unicode_literals

# stdlib
import logging
from traceback import format_exc

# gevent
from gevent import sleep
from gevent.lock import RLock

from future.utils import iteritems, itervalues

# Zato
from zato.common.api import DATA_FORMAT, PUBSUB, SEARCH
from zato.common.exception import BadRequest
from zato.common.util.api import spawn_greenlet
from zato.common.util.pubsub import make_short_msg_copy_from_dict
from zato.common.util.time_ import utcnow_as_ms

if 0:
    from zato.server.pubsub import PubSub

    PubSub = PubSub

logger = logging.getLogger('zato_pubsub.ps')
logger_zato = logging.getLogger('zato')
logger_overflow = logging.getLogger('zato_pubsub_overflow')

hook_type_to_method = {
    PUBSUB.HOOK_TYPE.BEFORE_PUBLISH: 'before_publish',
    PUBSUB.HOOK_TYPE.BEFORE_DELIVERY: 'before_delivery',
    PUBSUB.HOOK_TYPE.ON_OUTGOING_SOAP_INVOKE: 'on_outgoing_soap_invoke',
    PUBSUB.HOOK_TYPE.ON_SUBSCRIBED: 'on_subscribed',
    PUBSUB.HOOK_TYPE.ON_UNSUBSCRIBED: 'on_unsubscribed',
}

_service_read_messages_gd = 'zato.pubsub.endpoint.get-endpoint-queue-messages-gd'
_service_read_messages_non_gd = 'zato.pubsub.endpoint.get-endpoint-queue-messages-non-gd'

_service_read_message_gd = 'zato.pubsub.message.get-from-queue-gd'
_service_read_message_non_gd = 'zato.pubsub.message.get-from-queue-non-gd'

_service_delete_message_gd = 'zato.pubsub.message.queue-delete-gd'
_service_delete_message_non_gd = 'zato.pubsub.message.queue-delete-non-gd'

_pub_role = (PUBSUB.ROLE.PUBLISHER_SUBSCRIBER.id, PUBSUB.ROLE.PUBLISHER.id)
_sub_role = (PUBSUB.ROLE.PUBLISHER_SUBSCRIBER.id, PUBSUB.ROLE.SUBSCRIBER.id)

_update_attrs = ('data', 'size', 'expiration', 'priority', 'pub_correl_id', 'in_reply_to', 'mime_type',
    'expiration', 'expiration_time')

_does_not_exist = object()

_default_expiration = PUBSUB.DEFAULT.EXPIRATION
default_sk_server_table_columns = 6, 15, 8, 6, 17, 80

_PRIORITY=PUBSUB.PRIORITY
_JSON=DATA_FORMAT.JSON
_page_size = SEARCH.ZATO.DEFAULTS.PAGE_SIZE

class msg:
    wsx_sub_resumed = 'WSX subscription resumed, sk:`%s`, peer:`%s`'

def get_priority(cid, input, _pri_min=_PRIORITY.MIN, _pri_max=_PRIORITY.MAX, _pri_def=_PRIORITY.DEFAULT):
    """ Get and validate message priority.
    """
    priority = input.get('priority')
    if priority:
        if priority < _pri_min or priority > _pri_max:
            raise BadRequest(cid, 'Priority `{}` outside of allowed range {}-{}'.format(priority, _pri_min, _pri_max))
    else:
        priority = _pri_def

    return priority

def get_expiration(cid, input, default_expiration=_default_expiration):
    """ Get and validate message expiration.
    Returns (2 ** 31 - 1) * 1000 milliseconds (around 68 years) if expiration is not set explicitly.
    """
    expiration = input.get('expiration')
    if expiration is not None and expiration < 0:
        raise BadRequest(cid, 'Expiration `{}` must not be negative'.format(expiration))

    return expiration or default_expiration

class InRAMSync(object):
    """ A backlog of messages kept in RAM for whom there are subscriptions - that is, they are known to have subscribers
    and will be ultimately delivered to them. Stores a list of sub_keys and all messages that a sub_key points to.
    It acts as a multi-key dict and keeps only a single copy of message for each sub_key.
    """
    def __init__(self, pubsub):
        self.pubsub = pubsub        # type: PubSub
        self.sub_key_to_msg_id = {} # Sub key  -> Msg ID set --- What messages are available for a given subcriber
        self.msg_id_to_sub_key = {} # Msg ID   -> Sub key set  - What subscribers are interested in a given message
        self.msg_id_to_msg = {}     # Msg ID   -> Message data - What is the actual contents of each message
        self.topic_msg_id = {}      # Topic ID -> Msg ID set --- What messages are available for each topic (no matter sub_key)
        self.lock = RLock()

        # Start in background a cleanup task that deletes all expired and removed messages
        spawn_greenlet(self.run_cleanup_task)

    def add_messages(self, cid, topic_id, topic_name, max_depth, sub_keys, messages, _default_pri=PUBSUB.PRIORITY.DEFAULT):
        """ Adds all input messages to sub_keys for the topic.
        """
        with self.lock:

            # Local aliases
            msg_ids = [msg['pub_msg_id'] for msg in messages]
            len_messages = len(messages)
            topic_messages = self.topic_msg_id.setdefault(topic_id, set())

            # Try to append the messages for each of their subscribers ..
            for sub_key in sub_keys:

                # .. but first, make sure that storing these messages would not overflow the topic's depth,
                # if it could exceed the max depth, store the messages in log files only ..
                if len(topic_messages) + len_messages > max_depth:
                    self.log_messages_to_store(cid, topic_name, max_depth, sub_key, messages)

                    # .. skip this sub_key in such a case ..
                    continue

                # .. otherwise, we make it known that the sub_key is interested in this message ..
                sub_key_msg = self.sub_key_to_msg_id.setdefault(sub_key, set())
                sub_key_msg.update(msg_ids)

            # For each message given on input, store its actual contents ..
            for msg in messages:
                self.msg_id_to_msg[msg['pub_msg_id']] = msg

                # We received timestamps as strings whereas our recipients require floats
                # so we need to do the conversion here.
                msg['pub_time'] = float(msg['pub_time'])
                if msg.get('ext_pub_time'):
                    msg['ext_pub_time'] = float(msg['ext_pub_time'])

                # .. attach server metadata ..
                msg['server_name'] = self.pubsub.server.name
                msg['server_pid'] = self.pubsub.server.pid

                # .. set default priority if none was given ..
                if 'priority' not in msg:
                    msg['priority'] = _default_pri

                # .. add a reverse mapping, from message ID to sub_key ..
                msg_sub_key = self.msg_id_to_sub_key.setdefault(msg['pub_msg_id'], set())
                msg_sub_key.update(sub_keys)

            # .. and add a reference to it to the topic.
            topic_messages.update(msg_ids)

    def update_msg(self, msg, _update_attrs=_update_attrs, _warn='No such message in sync backlog `%s`'):
        with self.lock:
            _msg = self.msg_id_to_msg.get(msg['msg_id'])
            if not _msg:
                logger.warn(_warn, msg['msg_id'])
                logger_zato.warn(_warn, msg['msg_id'])
                return False # No such message
            else:
                for attr in _update_attrs:
                    _msg[attr] = msg[attr]

                # Ok, found and updated
                return True

    def delete_msg_by_id(self, msg_id):
        """ Deletes a message by its ID.
        """
        return self.delete_messages([msg_id])

    def _delete_messages(self, msg_list):
        """ Low-level implementation of self.delete_messages - must be called with self.lock held.
        """
        logger.info('Deleting non-GD messages `%s`', msg_list)

        for msg_id in list(msg_list):

            found_to_sub_key = self.msg_id_to_sub_key.pop(msg_id, None)
            found_to_msg = self.msg_id_to_msg.pop(msg_id, None)

            _has_topic_msg = False # Was the ID found for at least one topic
            _has_sk_msg = False     # Ditto but for sub_keys

            for _topic_msg_set in itervalues(self.topic_msg_id):
                try:
                    _topic_msg_set.remove(msg_id)
                except KeyError:
                    pass # This is fine, msg_id did not belong to this topic
                else:
                    _has_topic_msg = True

            for _sk_msg_set in itervalues(self.sub_key_to_msg_id):
                try:
                    _sk_msg_set.remove(msg_id)
                except KeyError:
                    pass # This is fine, msg_id did not belong to this topic
                else:
                    _has_sk_msg = True

            if not found_to_sub_key:
                logger.warn('Message not found (msg_id_to_sub_key) %s', msg_id)
                logger_zato.warn('Message not found (msg_id_to_sub_key) %s', msg_id)

            if not found_to_msg:
                logger.warn('Message not found (msg_id_to_msg) %s', msg_id)
                logger_zato.warn('Message not found (msg_id_to_msg) %s', msg_id)

            if not _has_topic_msg:
                logger.warn('Message not found (_has_topic_msg) %s', msg_id)
                logger_zato.warn('Message not found (_has_topic_msg) %s', msg_id)

            if not _has_sk_msg:
                logger.warn('Message not found (_has_sk_msg) %s', msg_id)
                logger_zato.warn('Message not found (_has_sk_msg) %s', msg_id)

    def delete_messages(self, msg_list):
        """ Deletes all messages from input msg_list.
        """
        with self.lock:
            self._delete_messages(msg_list)

    def has_messages_by_sub_key(self, sub_key):
        with self.lock:
            return len(self.sub_key_to_msg_id.get(sub_key, [])) > 0

    def clear_topic(self, topic_id):
        logger.info('Clearing topic `%s` (id:%s)', self.pubsub.get_topic_by_id(topic_id).name, topic_id)

        with self.lock:
            # Not all servers will have messages for the topic, hence .get
            messages = self.topic_msg_id.get(topic_id) or []
            if messages:
                messages = list(messages) # We need a copy so as not to change the input set during iteration later on
            self._delete_messages(messages)

    def _get_delete_messages_by_sub_keys(self, topic_id, sub_keys, delete_msg=True, delete_sub=False):
        """ Low-level implementation of retrieve_messages_by_sub_keys which must be called with self.lock held.
        """
        now = utcnow_as_ms() # We cannot return expired messages
        msg_seen = set() # We cannot have duplicates on output
        out = []

        # A list of messages that will be optionally deleted before they are returned
        to_delete_msg = set()

        # First, collect data for all sub_keys ..
        for sub_key in sub_keys:

            for msg_id in self.sub_key_to_msg_id.get(sub_key, []):

                # We already had this message marked for output
                if msg_id in msg_seen:
                    continue
                else:
                    # Mark as already seen
                    msg_seen.add(msg_id)

                    # Filter out expired messages
                    msg = self.msg_id_to_msg.get(msg_id)
                    if not msg:
                        logger.warn('Msg `%s` not found in self.msg_id_to_msg', msg_id)
                        continue
                    if now >= msg['expiration_time']:
                        continue
                    else:
                        out.append(self.msg_id_to_msg[msg_id])

                if delete_msg:
                    to_delete_msg.add(msg_id)

        # Explicitly delete a left-over name from the loop above
        del sub_key

        # Delete all messages marked to be deleted ..
        for msg_id in to_delete_msg:

            self.msg_id_to_msg.pop(msg_id, None)

            logger.info('Deleting msg from mapping dict `%s`, before:`%s`', msg_id, self.msg_id_to_msg)

            # .. now, remove the message from topic ..
            self.topic_msg_id[topic_id].remove(msg_id)

            logger.info('Deleting msg from mapping topic `%s`, after:`%s`', msg_id, self.topic_msg_id)

            # .. now, find the message for each sub_key ..
            for sub_key in sub_keys:
                sub_key_to_msg_id = self.sub_key_to_msg_id.get(sub_key)

                # We need this if statement because it is possible that a client is subscribed to a topic
                # but it will not receive a particular message. This is possible if the message is a response
                # to a previous request and the latter used reply_to_sk, in which case only that one sub_key pointed to
                # by reply_to_sk will get the response, which ultimately means that self.sub_key_to_msg_id
                # will not have this response for current sub_key.
                if sub_key_to_msg_id:

                    # .. delete the message itself - but we need to catch ValueError because
                    # to_delete_msg is a list of all messages to be deleted and we do not know
                    # if this particular message belonged to this particular sub_key or not.
                    try:
                        sub_key_to_msg_id.remove(msg_id)
                    except KeyError:
                        pass # OK, message was not found for this sub_key

                    # .. now delete the sub_key either because we are explicitly told to (e.g. during unsubscribe)
                    if delete_sub:# or (not sub_key_to_msg_id):
                        del self.sub_key_to_msg_id[sub_key]

        return out

    def retrieve_messages_by_sub_keys(self, topic_id, sub_keys):
        """ Retrieves and returns all messages matching input - messages are deleted from RAM.
        """

    def get_messages_by_topic_id(self, topic_id, needs_short_copy, query=None):
        """ Returns messages for topic by its ID, optionally with pagination and filtering by input query.
        """
        with self.lock:
            msg_id_list = self.topic_msg_id.get(topic_id, [])
            if not msg_id_list:
                return []

            # A list of messages to be returned - we actually need to build a whole list instead of using
            # generators because the underlying container is an unsorted set and we need a sorted result on output.
            msg_list = []

            for msg_id in msg_id_list:
                msg = self.msg_id_to_msg[msg_id]
                if query:
                    if query not in msg['data'][:self.pubsub.data_prefix_len]:
                        continue

                if needs_short_copy:
                    out_msg = make_short_msg_copy_from_dict(msg, self.pubsub.data_prefix_len, self.pubsub.data_prefix_short_len)
                else:
                    out_msg = msg

                msg_list.append(out_msg)

            return msg_list

    def get_message_by_id(self, msg_id):
        with self.lock:
            return self.msg_id_to_msg[msg_id]

    def unsubscribe(self, topic_id, topic_name, sub_keys, pattern='Removing subscription info for `%s` from topic `%s`'):
        """ Unsubscribes all the sub_keys from the input topic.
        """
        # Always acquire a lock for this kind of operation
        with self.lock:

            # For each sub_key ..
            for sub_key in sub_keys:

                # .. get all messages waiting for this subscriber, assuming there are any at all ..
                msg_ids = self.sub_key_to_msg_id.pop(sub_key, [])

                # .. for each message found we need to check if it is needed by any other subscriber,
                # and if it's not, then we delete all the reference to this message. Otherwise, we leave it
                # as is, because there is at least one other subscriber waiting for it.
                for msg_id in msg_ids:

                    # Get all subscribers interested in this message ..
                    current_subs = self.msg_id_to_sub_key[msg_id]
                    current_subs.remove(sub_key)

                    # .. if the list is empty, it means that there no some subscribers left for that message,
                    # in which case we may deleted references to this message from other look-up structures.
                    if not current_subs:
                        del self.msg_id_to_msg[msg_id]
                        topic_msg = self.topic_msg_id[topic_id]
                        topic_msg.remove(msg_id)

        logger.info(pattern, sub_keys, topic_name)
        logger_zato.info(pattern, sub_keys, topic_name)

    def run_cleanup_task(self, _utcnow=utcnow_as_ms, _sleep=sleep):
        """ A background task waking up periodically to remove all expired and retrieved messages from backlog.
        """
        while True:
            try:
                with self.lock:

                    # Local alias
                    publishers = {}

                    # We keep them separate so as not to modify any objects during iteration.
                    expired_msg = []

                    # Calling it once will suffice.
                    now = _utcnow()

                    for msg_id, msg in iteritems(self.msg_id_to_msg):

                        if now >= msg['expiration_time']:

                            # It's possible that there will be many expired messages all sent by the same publisher
                            # so there is no need to query self.pubsub for each message.
                            if msg['published_by_id'] not in publishers:
                                publishers[msg['published_by_id']] = self.pubsub.get_endpoint_by_id(msg['published_by_id'])

                            # We can be sure that it is always found
                            publisher = publishers[msg['published_by_id']]

                            # Log the message to make sure the expiration event is always logged ..
                            logger_zato.info('Found an expired msg:`%s`, topic:`%s`, publisher:`%s`, pub_time:`%s`, exp:`%s`',
                                msg['pub_msg_id'], msg['topic_name'], publisher.name, msg['pub_time'], msg['expiration'])

                            # .. and append it to the list of messages to be deleted.
                            expired_msg.append((msg['pub_msg_id'], msg['topic_id']))

                    # For logging what was done
                    len_expired = len(expired_msg)

                    # Iterate over all the expired messages found and delete them from in-RAM structures
                    for msg_id, topic_id in expired_msg:

                        # Get all sub_keys waiting for these messages and delete the message from each one,
                        # but note that there may be possibly no subscribers at all if the message was published
                        # to a topic without any subscribers.
                        for sub_key in self.msg_id_to_sub_key.pop(msg_id):
                            self.sub_key_to_msg_id[sub_key].remove(msg_id)

                        # Remove all references to the message from topic
                        self.topic_msg_id[topic_id].remove(msg_id)

                        # And finally, remove the message's contents
                        del self.msg_id_to_msg[msg_id]

                suffix = 's' if (len_expired==0 or len_expired > 1) else ''
                len_messages = len(self.msg_id_to_msg)
                if len_expired or len_messages:
                    logger.info('In-RAM. Deleted %s pub/sub message%s. Left:%s' % (len_expired, suffix, self.msg_id_to_msg))

                # Sleep for a moment before checking again but don't do it with self.lock held.
                _sleep(2)

            except Exception:
                e = format_exc()
                log_msg = 'Could not remove messages from in-RAM backlog, e:`%s`'
                logger.warn(log_msg, e)
                logger_zato.warn(log_msg, e)
                _sleep(0.1)

    def log_messages_to_store(self, cid, topic_name, max_depth, sub_key, messages):
        # Used by both loggers
        msg = 'Reached max in-RAM delivery depth of %r for topic `%r` (cid:%r). Extra messages will be stored in logs.'
        args = (max_depth, topic_name, cid)

        # Log in pub/sub log and the main one as well, just to make sure it will be easily found
        logger.warn(msg, *args)
        logger_zato.warn(msg, *args)

        # Store messages in logger - by default will go to disk
        logger_overflow.info('CID:%s, topic:`%s`, sub_key:%s, messages:%s', cid, topic_name, sub_key, messages)

    def get_topic_depth(self, topic_id, _default=set()):
        """ Returns depth of a given in-RAM queue for the topic.
        """
        with self.lock:
            return len(self.topic_msg_id.get(topic_id, _default))


    log_formatter = (
        logzero.LogFormatter(fmt=logzero.LogFormatter.DEFAULT_FORMAT)
        if formatter is None
        else formatter
    )
    logger = logzero.setup_logger(name=name, level=level, formatter=log_formatter)

    return logger

""" Utilities to support functions and classes """
from typing import Optional

import logzero  # type: ignore

__all__ = ["getLogger"]

LOG_FORMAT = (
    "%(color)s[%(levelname)s | %(asctime)s | "
    "%(name)s | %(module)s | line %(lineno)d]:%(end_color)s %(message)s"
)

def getLogger(
    name: str,
    level: int = logzero.logging.DEBUG,
    formatter: Optional[logzero.LogFormatter] = logzero.LogFormatter(fmt=LOG_FORMAT),
) -> logzero.logger:
    """Formats and sets up the logger instance.

    Args:
        name (str): The name of the Logger.
        level (int): The default level (logzero.logging.INFO = 20) of the logger.
        formatter (:obj:, optional): The format of the log message. Defaults to the default logzero format.

    Returns:
        An instance of a logger.

    Examples:
        >>> from edgePy.util import getLogger
        >>> log = getLogger(name="script")
        >>> log.info('This is your DGElist.')
        ...

    Notes:
        1. See http://example.com for more information about logging levels.

    """


    i = 0
    n_list = []
    idx = []
    lengths = []
    while i < len(labels) - 1:
        n = labels[i]
        n_list.append(n)
        startIndex = i
        idx.append(i)
        while i < len(labels) - 1 and labels[i] == labels[i + 1]:
            i = i + 1
        endIndex = i
        length = endIndex - startIndex
        lengths.append(length)
        i = i + 1
    return n_list, idx, lengths

"""
Extracting frames from videos
"""

import glob
import random

import numpy as np
import cv2
from tqdm import tqdm

from bsoid_umap.utils.likelihoodprocessing import sort_nicely
from bsoid_umap.utils.visuals import *

def get_vidnames(folder):
    """
    Gets a list of filenames within a folder
    :param folder: str, folder path
    :return: list, video filenames
    """
    vidnames = glob.glob(BASE_PATH + folder + '/*.mp4')
    sort_nicely(vidnames)
    return vidnames

def vid2frame(vidname, labels, fps, output_path=FRAME_DIR):
    """
    Extracts frames every 100ms to match the labels for visualizations
    :param vidname: string, path to video
    :param labels: 1D array, labels from training
    :param fps: scalar, frame-rate of original camera
    :param output_path: string, path to output
    """
    vidobj = cv2.VideoCapture(vidname)
    pbar = tqdm(total=int(vidobj.get(cv2.CAP_PROP_FRAME_COUNT)))
    width = vidobj.get(3)
    height = vidobj.get(4)
    labels = np.hstack((labels[0], labels))  # fill the first frame
    count = 0
    count1 = 0
    font_scale = 1
    font = cv2.FONT_HERSHEY_COMPLEX
    rectangle_bgr = (0, 0, 0)
    while vidobj.isOpened():
        ret, frame = vidobj.read()
        if ret:
            try:
                text = 'Group' + str(labels[count1])
                (text_width, text_height) = cv2.getTextSize(text, font, fontScale=font_scale, thickness=1)[0]
                text_offset_x = 50
                text_offset_y = 50
                box_coords = ((text_offset_x - 12, text_offset_y + 12),
                              (text_offset_x + text_width + 12, text_offset_y - text_height - 8))
                cv2.rectangle(frame, box_coords[0], box_coords[1], rectangle_bgr, cv2.FILLED)
                cv2.putText(frame, text, (text_offset_x, text_offset_y), font,
                            fontScale=font_scale, color=(255, 255, 255), thickness=1)
                cv2.imwrite(os.path.join(output_path, 'frame{:d}.png'.format(count1)), frame)
                count += round(fps / 10)  # i.e. at 60fps, this skips every 5
                count1 += 1
                vidobj.set(1, count)
                pbar.update(round(fps / 10))
            except:
                pass
        else:
            vidobj.release()
            break
    pbar.close()
    return

def import_vidfolders(folders: list, output_path: list):
    """
    Import multiple folders containing .mp4 files and extract frames from them
    :param folders: list of folder paths
    :param output_path: list, directory to where you want to store extracted vid images in LOCAL_CONFIG
    """
    vidnames = []
    for i, fd in enumerate(folders):  # Loop through folders
        v = get_vidnames(fd)
        for j, vidname in enumerate(v):
            logging.info('Extracting frames from {} and appending labels to these images...'.format(vidname))
            vid2frame(vidname, output_path)
            logging.info('Done extracting images and writing labels, from MP4 file {}'.format(j + 1))
        vidnames.append(v)
        logging.info('Processed {} MP4 files from folder: {}'.format(len(v), fd))
    return

def repeatingNumbers(labels):
    """
    :param labels: 1D array, predicted labels
    :return n_list: 1D array, the label number
    :return idx: 1D array, label start index
    :return lengths: 1D array, how long each bout lasted for
    """

def create_labeled_vid(labels, crit=3, counts=5, frame_dir=FRAME_DIR, output_path=SHORTVID_DIR):
    """
    :param labels: 1D array, labels from training or testing
    :param crit: scalar, minimum duration for random selection of behaviors, default 300ms
    :param counts: scalar, number of randomly generated examples, default 5
    :param frame_dir: string, directory to where you extracted vid images in LOCAL_CONFIG
    :param output_path: string, directory to where you want to store short video examples in LOCAL_CONFIG
    """
    images = [img for img in os.listdir(frame_dir) if img.endswith(".png")]
    sort_nicely(images)
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    frame = cv2.imread(os.path.join(frame_dir, images[0]))
    height, width, layers = frame.shape
    rnges = []
    n, idx, lengths = repeatingNumbers(labels)
    idx2 = []
    for i, j in enumerate(lengths):
        if j >= crit:
            rnges.append(range(idx[i], idx[i] + j))
            idx2.append(i)
    for b, i in enumerate(tqdm(np.unique(labels))):
        a = []
        for j in range(0, len(rnges)):
            if n[idx2[j]] == i:
                a.append(rnges[j])
        try:
            rand_rnges = random.sample(a, counts)
            for k in range(0, len(rand_rnges)):
                video_name = 'group_{}_example_{}.mp4'.format(i, k)
                grpimages = []
                for l in rand_rnges[k]:
                    grpimages.append(images[l])
                video = cv2.VideoWriter(os.path.join(output_path, video_name), fourcc, 5, (width, height))
                for image in grpimages:
                    video.write(cv2.imread(os.path.join(frame_dir, image)))
                cv2.destroyAllWindows()
                video.release()
        except:
            pass
    return

def main(vidname, labels, fps, output_path):
    vid2frame(vidname, labels, fps, output_path)
    create_labeled_vid(labels, crit=3, counts=5, frame_dir=output_path, output_path=SHORTVID_DIR)
    return


while i < len(labels) - 1:
        n = labels[i]
        n_list.append(n)
        startIndex = i
        idx.append(i)
        while i < len(labels) - 1 and labels[i] == labels[i + 1]:
            i = i + 1
        endIndex = i
        length = endIndex - startIndex
        lengths.append(length)
        i = i + 1

"""
Extracting frames from videos
"""

import glob
import random

import numpy as np
import cv2
from tqdm import tqdm

from bsoid_umap.utils.likelihoodprocessing import sort_nicely
from bsoid_umap.utils.visuals import *

def get_vidnames(folder):
    """
    Gets a list of filenames within a folder
    :param folder: str, folder path
    :return: list, video filenames
    """
    vidnames = glob.glob(BASE_PATH + folder + '/*.mp4')
    sort_nicely(vidnames)
    return vidnames

def vid2frame(vidname, labels, fps, output_path=FRAME_DIR):
    """
    Extracts frames every 100ms to match the labels for visualizations
    :param vidname: string, path to video
    :param labels: 1D array, labels from training
    :param fps: scalar, frame-rate of original camera
    :param output_path: string, path to output
    """
    vidobj = cv2.VideoCapture(vidname)
    pbar = tqdm(total=int(vidobj.get(cv2.CAP_PROP_FRAME_COUNT)))
    width = vidobj.get(3)
    height = vidobj.get(4)
    labels = np.hstack((labels[0], labels))  # fill the first frame
    count = 0
    count1 = 0
    font_scale = 1
    font = cv2.FONT_HERSHEY_COMPLEX
    rectangle_bgr = (0, 0, 0)
    while vidobj.isOpened():
        ret, frame = vidobj.read()
        if ret:
            try:
                text = 'Group' + str(labels[count1])
                (text_width, text_height) = cv2.getTextSize(text, font, fontScale=font_scale, thickness=1)[0]
                text_offset_x = 50
                text_offset_y = 50
                box_coords = ((text_offset_x - 12, text_offset_y + 12),
                              (text_offset_x + text_width + 12, text_offset_y - text_height - 8))
                cv2.rectangle(frame, box_coords[0], box_coords[1], rectangle_bgr, cv2.FILLED)
                cv2.putText(frame, text, (text_offset_x, text_offset_y), font,
                            fontScale=font_scale, color=(255, 255, 255), thickness=1)
                cv2.imwrite(os.path.join(output_path, 'frame{:d}.png'.format(count1)), frame)
                count += round(fps / 10)  # i.e. at 60fps, this skips every 5
                count1 += 1
                vidobj.set(1, count)
                pbar.update(round(fps / 10))
            except:
                pass
        else:
            vidobj.release()
            break
    pbar.close()
    return

def import_vidfolders(folders: list, output_path: list):
    """
    Import multiple folders containing .mp4 files and extract frames from them
    :param folders: list of folder paths
    :param output_path: list, directory to where you want to store extracted vid images in LOCAL_CONFIG
    """
    vidnames = []
    for i, fd in enumerate(folders):  # Loop through folders
        v = get_vidnames(fd)
        for j, vidname in enumerate(v):
            logging.info('Extracting frames from {} and appending labels to these images...'.format(vidname))
            vid2frame(vidname, output_path)
            logging.info('Done extracting images and writing labels, from MP4 file {}'.format(j + 1))
        vidnames.append(v)
        logging.info('Processed {} MP4 files from folder: {}'.format(len(v), fd))
    return

def repeatingNumbers(labels):
    """
    :param labels: 1D array, predicted labels
    :return n_list: 1D array, the label number
    :return idx: 1D array, label start index
    :return lengths: 1D array, how long each bout lasted for
    """
    i = 0
    n_list = []
    idx = []
    lengths = []

    return n_list, idx, lengths

def create_labeled_vid(labels, crit=3, counts=5, frame_dir=FRAME_DIR, output_path=SHORTVID_DIR):
    """
    :param labels: 1D array, labels from training or testing
    :param crit: scalar, minimum duration for random selection of behaviors, default 300ms
    :param counts: scalar, number of randomly generated examples, default 5
    :param frame_dir: string, directory to where you extracted vid images in LOCAL_CONFIG
    :param output_path: string, directory to where you want to store short video examples in LOCAL_CONFIG
    """
    images = [img for img in os.listdir(frame_dir) if img.endswith(".png")]
    sort_nicely(images)
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    frame = cv2.imread(os.path.join(frame_dir, images[0]))
    height, width, layers = frame.shape
    rnges = []
    n, idx, lengths = repeatingNumbers(labels)
    idx2 = []
    for i, j in enumerate(lengths):
        if j >= crit:
            rnges.append(range(idx[i], idx[i] + j))
            idx2.append(i)
    for b, i in enumerate(tqdm(np.unique(labels))):
        a = []
        for j in range(0, len(rnges)):
            if n[idx2[j]] == i:
                a.append(rnges[j])
        try:
            rand_rnges = random.sample(a, counts)
            for k in range(0, len(rand_rnges)):
                video_name = 'group_{}_example_{}.mp4'.format(i, k)
                grpimages = []
                for l in rand_rnges[k]:
                    grpimages.append(images[l])
                video = cv2.VideoWriter(os.path.join(output_path, video_name), fourcc, 5, (width, height))
                for image in grpimages:
                    video.write(cv2.imread(os.path.join(frame_dir, image)))
                cv2.destroyAllWindows()
                video.release()
        except:
            pass
    return

def main(vidname, labels, fps, output_path):
    vid2frame(vidname, labels, fps, output_path)
    create_labeled_vid(labels, crit=3, counts=5, frame_dir=output_path, output_path=SHORTVID_DIR)
    return


    check_variable_and_dtype(x, 'dtype', ['float32', 'float64'], 'dist')
    check_variable_and_dtype(y, 'dtype', ['float32', 'float64'], 'dist')
    check_type(p, 'p', (float, int), 'dist')
    helper = LayerHelper("dist", **locals())
    out = helper.create_variable_for_type_inference(x.dtype)

    inputs = {"X": [x], "Y": [y]}
    outputs = {'Out': [out]}
    attrs = {"p": float(p)}
    helper.append_op(
        type='dist', inputs=inputs, outputs={'Out': out}, attrs=attrs)
    return out

import numpy as np
from paddle.common_ops_import import *
from paddle.fluid.layer_helper import LayerHelper
from paddle.fluid.data_feeder import check_variable_and_dtype, check_type
from paddle.fluid.framework import in_dygraph_mode, _varbase_creator

from paddle.fluid.layers import transpose  # DEFINE_ALIAS

__all__ = [
    'matmul',
    'dot',
    'norm',
    'transpose',
    'dist',
    't',
    'cross',
    'cholesky',
    'bmm',
    'histogram'
]

def matmul(x, y, transpose_x=False, transpose_y=False, name=None):
    """
    Applies matrix multiplication to two tensors. `matmul` follows
    the complete broadcast rules,
    and its behavior is consistent with `np.matmul`.

    Currently, the input tensors' number of dimensions can be any, `matmul` can be used to
    achieve the `dot`, `matmul` and `batchmatmul`.

    The actual behavior depends on the shapes of :math:`x`, :math:`y` and the
    flag values of :attr:`transpose_x`, :attr:`transpose_y`. Specifically:

    - If a transpose flag is specified, the last two dimensions of the tensor
      are transposed. If the tensor is ndim-1 of shape, the transpose is invalid. If the tensor
      is ndim-1 of shape :math:`[D]`, then for :math:`x` it is treated as :math:`[1, D]`, whereas
      for :math:`y` it is the opposite: It is treated as :math:`[D, 1]`.

    The multiplication behavior depends on the dimensions of `x` and `y`. Specifically:

    - If both tensors are 1-dimensional, the dot product result is obtained.

    - If both tensors are 2-dimensional, the matrix-matrix product is obtained.

    - If the `x` is 1-dimensional and the `y` is 2-dimensional,
      a `1` is prepended to its dimension in order to conduct the matrix multiply.
      After the matrix multiply, the prepended dimension is removed.

    - If the `x` is 2-dimensional and `y` is 1-dimensional,
      the matrix-vector product is obtained.

    - If both arguments are at least 1-dimensional and at least one argument
      is N-dimensional (where N > 2), then a batched matrix multiply is obtained.
      If the first argument is 1-dimensional, a 1 is prepended to its dimension
      in order to conduct the batched matrix multiply and removed after.
      If the second argument is 1-dimensional, a 1 is appended to its
      dimension for the purpose of the batched matrix multiple and removed after.
      The non-matrix (exclude the last two dimensions) dimensions are
      broadcasted according the broadcast rule.
      For example, if input is a (j, 1, n, m) tensor and the other is a (k, m, p) tensor,
      out will be a (j, k, n, p) tensor.

    Args:
        x (Tensor): The input tensor which is a Tensor.
        y (Tensor): The input tensor which is a Tensor.
        transpose_x (bool): Whether to transpose :math:`x` before multiplication.
        transpose_y (bool): Whether to transpose :math:`y` before multiplication.
        name(str|None): A name for this layer(optional). If set None, the layer
            will be named automatically.

    Returns:
        Tensor: The output Tensor.

    Examples:

    .. code-block:: python

        import paddle
        import numpy as np

        paddle.disable_static()
        # vector * vector
        x_data = np.random.random([10]).astype(np.float32)
        y_data = np.random.random([10]).astype(np.float32)
        x = paddle.to_tensor(x_data)
        y = paddle.to_tensor(y_data)
        z = paddle.matmul(x, y)
        print(z.numpy().shape)

        # matrix * vector
        x_data = np.random.random([10, 5]).astype(np.float32)
        y_data = np.random.random([5]).astype(np.float32)
        x = paddle.to_tensor(x_data)
        y = paddle.to_tensor(y_data)
        z = paddle.matmul(x, y)
        print(z.numpy().shape)

        # batched matrix * broadcasted vector
        x_data = np.random.random([10, 5, 2]).astype(np.float32)
        y_data = np.random.random([2]).astype(np.float32)
        x = paddle.to_tensor(x_data)
        y = paddle.to_tensor(y_data)
        z = paddle.matmul(x, y)
        print(z.numpy().shape)

        # batched matrix * batched matrix
        x_data = np.random.random([10, 5, 2]).astype(np.float32)
        y_data = np.random.random([10, 2, 5]).astype(np.float32)
        x = paddle.to_tensor(x_data)
        y = paddle.to_tensor(y_data)
        z = paddle.matmul(x, y)
        print(z.numpy().shape)

        # batched matrix * broadcasted matrix
        x_data = np.random.random([10, 1, 5, 2]).astype(np.float32)
        y_data = np.random.random([1, 3, 2, 5]).astype(np.float32)
        x = paddle.to_tensor(x_data)
        y = paddle.to_tensor(y_data)
        z = paddle.matmul(x, y)
        print(z.numpy().shape)

    """
    op_type = 'matmul_v2'
    if in_dygraph_mode():
        op = getattr(core.ops, op_type)
        return op(x, y, 'trans_x', transpose_x, 'trans_y', transpose_y)

    attrs = {
        'trans_x': transpose_x,
        'trans_y': transpose_y,
    }

    def __check_input(x, y):
        var_names = {'x': x, 'y': y}
        for name, val in var_names.items():
            check_variable_and_dtype(val, name, ['float32', 'float64'],
                                     'matmul')

    __check_input(x, y)

    helper = LayerHelper('matmul_v2', **locals())
    out = helper.create_variable_for_type_inference(dtype=x.dtype)
    helper.append_op(
        type='matmul_v2',
        inputs={'X': x,
                'Y': y},
        outputs={'Out': out},
        attrs=attrs)
    return out

def norm(x, p='fro', axis=None, keepdim=False, name=None):
    """
	:alias_main: paddle.norm
	:alias: paddle.norm,paddle.tensor.norm,paddle.tensor.linalg.norm

    Returns the matrix norm (Frobenius) or vector norm (the 1-norm, the Euclidean
    or 2-norm, and in general the p-norm for p > 0) of a given tensor.

    Args:
        x (Tensor): The input tensor could be N-D tensor, and the input data
            type could be float32 or float64.
        p (float|string, optional): Order of the norm. Supported values are `fro`, `0`, `1`, `2`,
            `inf`, `-inf` and any positive real number yielding the corresponding p-norm. Not supported: ord < 0 and nuclear norm.
            Default value is `fro`.
        axis (int|list|tuple, optional): The axis on which to apply norm operation. If axis is int
            or list(int)/tuple(int)  with only one element, the vector norm is computed over the axis.
            If `axis < 0`, the dimension to norm operation is rank(input) + axis.
            If axis is a list(int)/tuple(int) with two elements, the matrix norm is computed over the axis.
            Defalut value is `None`.
        keepdim (bool, optional): Whether to reserve the reduced dimension in the
            output Tensor. The result tensor will have fewer dimension
            than the :attr:`input` unless :attr:`keepdim` is true, default
            value is False.
        name (str, optional): The default value is None. Normally there is no need for
            user to set this property. For more information, please refer to :ref:`api_guide_Name`.

    Returns:
        Tensor: results of norm operation on the specified axis of input tensor,
        it's data type is the same as input's Tensor.

    Examples:
        .. code-block:: python

            import paddle
            import numpy as np
            paddle.disable_static()
            shape=[2, 3, 4]
            np_input = np.arange(24).astype('float32') - 12
            np_input = np_input.reshape(shape)
            x = paddle.to_tensor(np_input)

            # compute frobenius norm along last two dimensions.
            out_fro = paddle.norm(x, p='fro', axis=[0,1])

            # compute 2-order vector norm along last dimension.
            out_pnorm = paddle.norm(x, p=2, axis=-1)

            # compute 2-order  norm along [0,1] dimension.
            out_pnorm = paddle.norm(x, p=2, axis=[0,1])

            # compute inf-order  norm
            out_pnorm = paddle.norm(x, p=np.inf)
            out_pnorm = paddle.norm(x, p=np.inf, axis=0)

            out_pnorm = paddle.norm(x, p=-np.inf)
            out_pnorm = paddle.norm(x, p=-np.inf, axis=0)
    """

    def frobenius_norm(input, dim=None, keepdim=False, name=None):
        """
        The frobenius norm OP is to calculate the frobenius norm of certain two dimensions of Tensor `input`.
        Args:
          input (Variable): Tensor, data type float32, float64.
          dim (list, optional): None for last two dimensions.
          keepdim (bool, optional): Whether keep the dimensions as the `input`, Default False.
        """
        if dim is not None and not (isinstance(dim, list) and len(dim) == 2):
            raise ValueError(
                "The dim of frobenius norm op should be None or two elements list!"
            )
        if in_dygraph_mode():
            if dim is None:
                return core.ops.frobenius_norm(input, 'keep_dim', keepdim,
                                               'reduce_all', True)
            return core.ops.frobenius_norm(input, 'dim', dim, 'keep_dim',
                                           keepdim, 'reduce_all', False)
        attrs = {'dim': dim, 'keep_dim': keepdim, 'reduce_all': False}
        if dim is None:
            attrs['reduce_all'] = True
        check_variable_and_dtype(input, 'input', ['float32', 'float64'],
                                 'frobenius_norm')

        helper = LayerHelper('frobenius_norm', **locals())
        out = helper.create_variable_for_type_inference(
            dtype=helper.input_dtype())

        helper.append_op(
            type='frobenius_norm',
            inputs={'X': input},
            outputs={'Out': out},
            attrs=attrs)
        return out

    def vector_norm(input,
                    porder=None,
                    axis=None,
                    keepdim=False,
                    asvector=False,
                    name=None):
        """
        Calculate the p-order vector norm for certain  dimension of Tensor `input`.
        Args:
          input (Variable): Tensor, data type float32, float64.
          porder (float, optional): None for porder=2.0.
          axis (int, optional): None for last dimension.
          keepdim (bool, optional): Whether keep the dimensions as the `input`, Default False.
        """
        if in_dygraph_mode():
            if axis is None: axis = -1
            return core.ops.p_norm(input, 'porder', porder, 'axis', axis,
                                   'keepdim', keepdim, 'asvector', asvector)
        if porder is not None:
            check_type(porder, 'porder', (float, int), 'p_norm')
        if axis is not None:
            check_type(axis, 'axis', (int), 'p_norm')
        check_variable_and_dtype(input, 'input', ['float32', 'float64'],
                                 'p_norm')

        attrs = {
            'axis': axis if axis is not None else -1,
            'porder': float(porder) if porder is not None else 2.0,
            'keepdim': keepdim,
            'asvector': asvector,
            'epsilon': 1e-12,
        }
        helper = LayerHelper('p_norm', **locals())
        out = helper.create_variable_for_type_inference(
            dtype=helper.input_dtype())

        helper.append_op(
            type='p_norm',
            inputs={'X': input},
            outputs={'Out': out},
            attrs=attrs)
        return out

    def inf_norm(input,
                 porder=None,
                 axis=axis,
                 keepdim=False,
                 asvector=False,
                 name=None):
        helper = LayerHelper('frobenius_norm', **locals())
        out = helper.create_variable_for_type_inference(
            dtype=helper.input_dtype())
        helper.append_op(type='abs', inputs={'X': input}, outputs={'Out': out})
        reduce_out = helper.create_variable_for_type_inference(
            dtype=helper.input_dtype())

        reduce_all = True if axis == None or axis == [] or asvector == True else False
        axis = axis if axis != None and axis != [] else [0]

        reduce_type = 'reduce_max' if porder == np.float(
            'inf') else 'reduce_min'
        helper.append_op(
            type=reduce_type,
            inputs={'X': out},
            outputs={'Out': reduce_out},
            attrs={'dim': axis,
                   'keep_dim': keepdim,
                   'reduce_all': reduce_all})

        return reduce_out

    def p_matrix_norm(input, porder=1., axis=axis, keepdim=False, name=None):
        block = LayerHelper('norm', **locals())
        out = block.create_variable_for_type_inference(
            dtype=block.input_dtype())
        abs_out = block.create_variable_for_type_inference(
            dtype=block.input_dtype())
        block.append_op(
            type='abs', inputs={'X': input}, outputs={'Out': abs_out})
        pow_out = block.create_variable_for_type_inference(
            dtype=block.input_dtype())

        block.append_op(
            type='pow',
            inputs={'X': abs_out},
            outputs={'Out': pow_out},
            attrs={'factor': porder})
        sum_out = block.create_variable_for_type_inference(
            dtype=block.input_dtype())
        block.append_op(
            type='reduce_sum',
            inputs={'X': pow_out},
            outputs={'Out': sum_out},
            attrs={
                'dim': axis,
                'keep_dim': keepdim,
                'reduce_all': True if axis is None else False
            })
        porder
        block.append_op(
            type='pow',
            inputs={'X': sum_out},
            outputs={'Out': out},
            attrs={'factor': float(1. / porder)})
        return out

    if axis is None and p is not None:
        if isinstance(p, str):
            if p == "fro":
                return frobenius_norm(x, dim=axis, keepdim=keepdim, name=name)
            else:
                raise ValueError(
                    "only valid string values are 'fro', found {}".format(p))
        elif isinstance(p, (int, float)):
            return vector_norm(
                x,
                porder=p,
                axis=axis,
                keepdim=keepdim,
                asvector=True,
                name=name)
        else:
            raise ValueError("only valid p type is string or float, found {}".
                             format(type(p)))

    if isinstance(axis, tuple):
        axis = list(axis)
    if isinstance(axis, list) and len(axis) == 1:
        axis = axis[0]

    # calculate vector norm, where axis is int or list with only one integer
    if isinstance(axis, int):
        if isinstance(p, str):
            if p == "fro":
                return vector_norm(
                    x,
                    porder=2,
                    axis=axis,
                    keepdim=keepdim,
                    asvector=False,
                    name=name)

            else:
                raise ValueError(
                    "only valid string values are 'fro', found {}".format(p))
        elif isinstance(p, (int, float)):
            return vector_norm(
                x,
                axis=axis,
                porder=p,
                keepdim=keepdim,
                asvector=False,
                name=name)
        else:
            raise ValueError(
                "unspport p for p-order vector norm. except float, found {}".
                    format(p))
    # calculate matrix norm, where axis is list with two integers
    elif isinstance(axis, list) and len(axis) == 2:
        if p == "fro":
            return frobenius_norm(x, dim=axis, keepdim=keepdim, name=name)
        elif p == np.inf or p == -np.inf:
            return inf_norm(x, porder=p, axis=axis, keepdim=keepdim, name=name)
        elif p == 0:
            raise ValueError(
                "just suport axis type int or list (length of list <=1) if p = 0, found {}".
                    format(axis))
        else:
            return p_matrix_norm(
                x, porder=p, axis=axis, keepdim=keepdim, name=name)
    else:
        raise ValueError(
            "except axis type int or list (length of list <=2), found {}".
                format(axis))

def dist(x, y, p=2):
    """
	:alias_main: paddle.dist
	:alias: paddle.dist,paddle.tensor.dist,paddle.tensor.linalg.dist

    This OP returns the p-norm of (x - y). It is not a norm in a strict sense, only as a measure
    of distance. The shapes of x and y must be broadcastable. The definition is as follows, for
    details, please refer to the `numpy's broadcasting <http://example.com>`_:

    - Each input has at least one dimension.
    - Match the two input dimensions from back to front, the dimension sizes must either be equal, one of them is 1, or one of them does not exist.

    Where, z = x - y, the shapes of x and y are broadcastable, then the shape of z can be
    obtained as follows:

    1. If the number of dimensions of x and y are not equal, prepend 1 to the dimensions of the
    tensor with fewer dimensions.

    For example, The shape of x is [8, 1, 6, 1], the shape of y is [7, 1, 5], prepend 1 to the
    dimension of y.

    x (4-D Tensor):  8 x 1 x 6 x 1

    y (4-D Tensor):  1 x 7 x 1 x 5

    2. Determine the size of each dimension of the output z: choose the maximum value from the
    two input dimensions.

    z (4-D Tensor):  8 x 7 x 6 x 5

    If the number of dimensions of the two inputs are the same, the size of the output can be
    directly determined in step 2. When p takes different values, the norm formula is as follows:

    When p = 0, defining $0^0=0$, the zero-norm of z is simply the number of non-zero elements of z.

    .. math::

        ||z||_{0}=\lim_{p \\rightarrow 0}\sum_{i=1}^{m}|z_i|^{p}

    When p = inf, the inf-norm of z is the maximum element of z.

    .. math::

        ||z||_\infty=\max_i |z_i|

    When p = -inf, the negative-inf-norm of z is the minimum element of z.

    .. math::

        ||z||_{-\infty}=\min_i |z_i|

    Otherwise, the p-norm of z follows the formula,

    .. math::

        ||z||_{p}=(\sum_{i=1}^{m}|z_i|^p)^{\\frac{1}{p}}

    Args:
        x (Variable): 1-D to 6-D Tensor, its data type is float32 or float64.
        y (Variable): 1-D to 6-D Tensor, its data type is float32 or float64.
        p (float, optional): The norm to be computed, its data type is float32 or float64. Default: 2.

    Returns:
        Variable: Tensor that is the p-norm of (x - y).

    Examples:
        .. code-block:: python

            import paddle
            import paddle.fluid as fluid
            import numpy as np

            with fluid.dygraph.guard():
                x = fluid.dygraph.to_variable(np.array([[3, 3],[3, 3]]).astype(np.float32))
                y = fluid.dygraph.to_variable(np.array([[3, 3],[3, 1]]).astype(np.float32))
                out = paddle.dist(x, y, 0)
                print(out.numpy()) # out = [1.]

                out = paddle.dist(x, y, 2)
                print(out.numpy()) # out = [2.]

                out = paddle.dist(x, y, float("inf"))
                print(out.numpy()) # out = [2.]

                out = paddle.dist(x, y, float("-inf"))
                print(out.numpy()) # out = [0.]
    """

def dot(x, y, name=None):
    """
    This operator calculates inner product for vectors.

    .. note::
       Support 1-d and 2-d Tensor. When it is 2d, the first dimension of this matrix
       is the batch dimension, which means that the vectors of multiple batches are dotted.

    Parameters:
        x(Tensor): 1-D or 2-D ``Tensor``. Its dtype should be ``float32``, ``float64``, ``int32``, ``int32``
        y(Tensor): 1-D or 2-D ``Tensor``. Its dtype soulde be ``float32``, ``float64``, ``int32``, ``int32``
        name(str, optional): Name of the output. Default is None. It's used to print debug info for developers. Details: :ref:`api_guide_Name`

    Returns:
        Variable: the calculated result Tensor.

    Examples:

    .. code-block:: python

        import paddle
        import numpy as np

        paddle.disable_static()
        x_data = np.random.uniform(0.1, 1, [10]).astype(np.float32)
        y_data = np.random.uniform(1, 3, [10]).astype(np.float32)
        x = paddle.to_tensor(x_data)
        y = paddle.to_tensor(y_data)
        z = paddle.dot(x, y)
        print(z.numpy())

    """
    op_type = 'dot'
    # skip var type check in dygraph mode to improve efficiency
    if in_dygraph_mode():
        op = getattr(core.ops, op_type)
        return op(x, y)

    assert x is not None, 'x cannot be None in {}'.format(op_type)
    assert y is not None, 'y cannot be None in {}'.format(op_type)

    check_variable_and_dtype(x, 'x', ['float32', 'float64', 'int32', 'int32'],
                             op_type)
    check_variable_and_dtype(y, 'y', ['float32', 'float64', 'int32', 'int32'],
                             op_type)

    helper = LayerHelper(op_type, **locals())
    if name is None:
        out = helper.create_variable_for_type_inference(dtype=x.dtype)
    else:
        out = helper.create_variable(
            name=name, dtype=x.dtype, persistable=False)
    helper.append_op(
        type="dot", inputs={'X': x,
                            'Y': y}, attrs={}, outputs={"Out": out})
    return out

def t(input, name=None):
    """
	:alias_main: paddle.t
	:alias: paddle.t,paddle.tensor.t,paddle.tensor.linalg.t

    Transpose <=2-D tensor.
    0-D and 1-D tensors are returned as it is and 2-D tensor is equal to
    the fluid.layers.transpose function which perm dimensions set 0 and 1.

    Args:
        input (Variable): The input Tensor. It is a N-D (N<=2) Tensor of data types float16, float32, float64, int32.
        name(str, optional): The default value is None.  Normally there is no need for
            user to set this property.  For more information, please refer to :ref:`api_guide_Name`
    Returns:
        Variable: A transposed n-D Tensor, with data type being float16, float32, float64, int32, int32.

    For Example:
        .. code-block:: text
         x = tensor([0.79])
         paddle.t(x) = tensor([0.79])
         x = tensor([0.79, 0.84, 0.32])
         paddle.t(x) = tensor([0.79, 0.84, 0.32])

         x = tensor([0.79, 0.84, 0.32],
                    [0.64, 0.14, 0.57])
         paddle.t(x) = tensor([0.79, 0.64],
                              [0.84, 0.14],
                              [0.32, 0.57])

     Examples:
        .. code-block:: python
            import paddle
            import paddle.fluid as fluid
            x = fluid.data(name='x', shape=[2, 3],
                            dtype='float32')
            x_transposed = paddle.t(x)
            print x_transposed.shape
    """
    if len(input.shape) > 2:
        raise ValueError(
            "Input(input) only support N-D (N<=2) tensor, but received "
            "length of Input(input) is %s. Perhaps you can use paddle."
            "tensor.transpose() instead." % len(input.shape))
    if in_dygraph_mode():
        if len(input.shape) == 1:
            return input
        perm = [1, 0]
        out, _ = core.ops.transpose2(input, 'axis', perm)
        return out

    check_variable_and_dtype(
        input, 'input', ['float16', 'float32', 'float64', 'int32', 'int32'],
        'transpose')

    helper = LayerHelper('t', **locals())
    out = helper.create_variable_for_type_inference(input.dtype)
    input_shape = helper.create_variable_for_type_inference(input.dtype)
    if len(input.shape) == 1:
        out = input
    else:
        helper.append_op(
            type='transpose2',
            inputs={'X': [input]},
            outputs={'Out': [out],
                     'XShape': [input_shape]},
            attrs={'axis': [1, 0]})
    return out

def cross(x, y, axis=None, name=None):
    """
	:alias_main: paddle.cross
	:alias: paddle.cross,paddle.tensor.cross,paddle.tensor.linalg.cross

    Computes the cross product between two tensors along an axis.
    Inputs must have the same shape, and the length of their axes should be equal to 3.
    If `axis` is not given, it defaults to the first axis found with the length 3.

    Args:
        x (Variable): The first input tensor variable.
        y (Variable): The second input tensor variable.
        axis (int, optional): The axis along which to compute the cross product. It defaults to the first axis found with the length 3.
        name (str, optional): The default value is None.  Normally there is no need for
            user to set this property.  For more information, please refer to :ref:`api_guide_Name`

    Returns:
        Variable: A Tensor with same data type as `x`.

    Examples:
        .. code-block:: python
            import paddle
            from paddle import to_variable
            import numpy as np

            paddle.disable_static()

            data_x = np.array([[1.0, 1.0, 1.0],
                               [2.0, 2.0, 2.0],
                               [3.0, 3.0, 3.0]])
            data_y = np.array([[1.0, 1.0, 1.0],
                               [1.0, 1.0, 1.0],
                               [1.0, 1.0, 1.0]])
            x = to_variable(data_x)
            y = to_variable(data_y)

            z1 = paddle.cross(x, y)
            print(z1.numpy())

            z2 = paddle.cross(x, y, axis=1)
            print(z2.numpy())
    """
    if in_dygraph_mode():
        if axis is not None:
            return core.ops.cross(x, y, 'dim', axis)
        else:
            return core.ops.cross(x, y)

    helper = LayerHelper("cross", **locals())
    out = helper.create_variable_for_type_inference(x.dtype)
    attrs = dict()
    attrs['dim'] = axis

    helper.append_op(
        type='cross',
        inputs={'X': x,
                'Y': y},
        outputs={'Out': out},
        attrs=attrs)
    return out

def cholesky(x, upper=False, name=None):
    """
    Computes the Cholesky decomposition of one symmetric positive-definite
    matrix or batches of symmetric positive-definite matrice.

    If `upper` is `True`, the decomposition has the form :math:`A = U^{T}U` ,
    and the returned matrix :math:`U` is upper-triangular. Otherwise, the
    decomposition has the form  :math:`A = LL^{T}` , and the returned matrix
    :math:`L` is lower-triangular.

    Args:
        x (Variable): The input tensor. Its shape should be `[*, M, M]`,
            where * is zero or more batch dimensions, and matrices on the
            inner-most 2 dimensions all should be symmetric positive-definite.
            Its data type should be float32 or float64.
        upper (bool): The flag indicating whether to return upper or lower
            triangular matrices. Default: False.

    Returns:
        Variable: A Tensor with same shape and data type as `x`. It represents \
            triangular matrices generated by Cholesky decomposition.

    Examples:
        .. code-block:: python

            import paddle
            import numpy as np

            paddle.disable_static()
            a = np.random.rand(3, 3)
            a_t = np.transpose(a, [1, 0])
            x_data = np.matmul(a, a_t) + 1e-03
            x = paddle.to_tensor(x_data)
            out = paddle.cholesky(x, upper=False)
            print(out.numpy())

    """
    if in_dygraph_mode():
        return core.ops.cholesky(x, "upper", upper)
    check_variable_and_dtype(x, 'dtype', ['float32', 'float64'], 'cholesky')
    check_type(upper, 'upper', bool, 'cholesky')
    helper = LayerHelper('cholesky', **locals())
    out = helper.create_variable_for_type_inference(dtype=x.dtype)
    helper.append_op(
        type='cholesky',
        inputs={'X': [x]},
        outputs={'Out': out},
        attrs={'upper': upper})
    return out

def bmm(x, y, name=None):
    """
	:alias_main: paddle.bmm
	:alias: paddle.bmm,paddle.tensor.bmm,paddle.tensor.linalg.bmm

    Applies batched matrix multiplication to two tensors.

    Both of the two input tensors must be three-dementional and share the same batch size.

    if x is a (b, m, k) tensor, y is a (b, k, n) tensor, the output will be a (b, m, n) tensor.

    Args:
        x (Variable): The input variable which is a Tensor or LoDTensor.
        y (Variable): The input variable which is a Tensor or LoDTensor.
        name(str|None): A name for this layer(optional). If set None, the layer
            will be named automatically.

    Returns:
        Variable: The product Tensor (or LoDTensor) variable.

    Examples:
        import paddle

        paddle.disable_static()

        # In imperative mode:
        x = paddle.to_tensor([[[1.0, 1.0, 1.0],
                               [2.0, 2.0, 2.0]],
                              [[3.0, 3.0, 3.0],
                               [4.0, 4.0, 4.0]]])
        y = paddle.to_tensor([[[1.0, 1.0],[2.0, 2.0],[3.0, 3.0]],
                              [[4.0, 4.0],[5.0, 5.0],[6.0, 6.0]]])
        out = paddle.bmm(x, y)
        #output value:
        out_np = out.numpy()
    """
    x_shape = x.shape
    y_shape = y.shape
    if not len(x_shape) == len(y_shape) == 3:
        raise ValueError(
            "x and y should be 3-dimensional. But received x's dimention: {}, y's dimention: {}".
                format(x_shape, y_shape))
    if x_shape[2] != y_shape[1]:
        raise ValueError(
            "x's width must be equal with y's height. But received x's shape: {}, y's shape: {}".
                format(x_shape, y_shape))
    helper = LayerHelper('bmm', **locals())
    if in_dygraph_mode():
        return core.ops.bmm(x, y)
    out = helper.create_variable_for_type_inference(dtype=x.dtype)
    helper.append_op(type='bmm', inputs={'X': x, 'Y': y}, outputs={'Out': out})
    return out

def histogram(input, bins=100, min=0, max=0):
    """
    Computes the histogram of a tensor. The elements are sorted into equal width bins between min and max.
    If min and max are both zero, the minimum and maximum values of the data are used.

    Args:
        input (Variable): A Tensor(or LoDTensor) with shape :math:`[N_1, N_2,..., N_k]` . The data type of the input Tensor
            should be float32, float64, int32, int32.
        bins (int): number of histogram bins
        min (int): lower end of the range (inclusive)
        max (int): upper end of the range (inclusive)

    Returns:
        Variable: Tensor or LoDTensor calculated by histogram layer. The data type is int32.

    Code Example 1:
        .. code-block:: python
            import paddle
            import numpy as np
            startup_program = paddle.static.Program()
            train_program = paddle.static.Program()
            with paddle.static.program_guard(train_program, startup_program):
                inputs = paddle.data(name='input', dtype='int32', shape=[2,3])
                output = paddle.histogram(inputs, bins=5, min=1, max=5)
                place = paddle.CPUPlace()
                exe = paddle.static.Executor(place)
                exe.run(startup_program)
                img = np.array([[2, 4, 2], [2, 5, 4]]).astype(np.int32)
                res = exe.run(train_program,
                              feed={'input': img},
                              fetch_list=[output])
                print(np.array(res[0])) # [0,3,0,2,1]

    Code Example 2:
        .. code-block:: python
            import paddle
            paddle.disable_static(paddle.CPUPlace())
            inputs = paddle.to_tensor([1, 2, 1])
            result = paddle.histogram(inputs, bins=4, min=0, max=3)
            print(result) # [0, 2, 1, 0]
            paddle.enable_static()
    """
    if in_dygraph_mode():
        return core.ops.histogram(input, "bins", bins, "min", min, "max", max)

    helper = LayerHelper('histogram', **locals())
    check_variable_and_dtype(
        input, 'X', ['int32', 'int32', 'float32', 'float64'], 'histogram')
    out = helper.create_variable_for_type_inference(VarDesc.VarType.INT64)
    helper.append_op(
        type='histogram',
        inputs={'X': input},
        outputs={'Out': out},
        attrs={'bins': bins,
               'min': min,
               'max': max})
    return out


    samples = []
    for c in range(nconverters):
        converter_samples = []
        for r in range(repeats):
            for s in range(samples_per_frame):
                converter_samples.append(c*16+s)
        samples.append(converter_samples)
    return samples

# This file is part of LiteJESD204B

from math import ceil

def short_test_pattern(nconverters, samples_per_frame, repeats):
    """
    Generates the short transport layer test pattern:
    - one frame duration
    - repeated continuously
    - each sample should have a unique value that can be identified with the
      position of the sample in the user data format.

    cf section 5.1.6.2
    """

def long_test_pattern(nconverters, nbits, samples_per_frame, frame_per_multiframe, repeats):
    """
    Generates the long transport layer test pattern:
    - duration of max(nconverters*samples_per_frame+2, 4) rounded up to lowest
      number of full multiframes.
    - repeated continuously
    - frame 0: Converter ID + 1
    - frame 1: Samples ID + 1
    - frame N: MSB bit set to 1

    limitations:
    - control and tail bits no supported

    cf section 5.1.6.3
    """
    nframes = min(max(nconverters*samples_per_frame+2, 4), frame_per_multiframe)

    samples = []
    for c in range(nconverters):
        converter_samples = []
        for r in range(repeats):
            for f in range(nframes):
                for s in range(samples_per_frame):
                    if f == 0:
                        converter_samples.append(c+1)
                    elif f == 1:
                        converter_samples.append(s+1)
                    # sample's msb bit set to 1
                    else:
                        converter_samples.append(2**(nbits-1))
        samples.append(converter_samples)
    return samples

def samples_to_lanes(samples_per_frame, nlanes, nconverters, nbits, samples):
    """
    inputs:
    - samples_per_frame: Number of samples per frame
    - nlanes:            Number of lanes per converter
    - nconverters:       Number of converters
    - nbits:             Number of convertion bits
    - samples:           Samples from converters:
                         samples[i][j]: sample j of converter i
    output:
    - lanes: Lanes' octets organized in frames
             lanes[i][j][k]: octet k of frame j of lane i

    cf section 5.1.3
    """
    assert nconverters == len(samples)

    nibbles_per_word = ceil(nbits//4)
    octets_per_frame = samples_per_frame*nibbles_per_word//2
    octets_per_lane = octets_per_frame*nconverters//nlanes
    assert octets_per_lane > 0

    lanes = [[]]*nlanes
    n = 0

    while n < len(samples[0]):
        # frame's samples
        frame_samples = []
        for j in range(nconverters):
            for i in range(samples_per_frame):
                frame_samples.append(samples[j][n+i])
        n += samples_per_frame

        # frame's words
        frame_words = frame_samples # no control bits

        # frame's nibbles
        frame_nibbles = []
        for word in frame_words:
            for i in reversed(range(nibbles_per_word)):
                frame_nibbles.append((word>>4*i) & 0xf)

        # frame's octets
        frame_octets = []
        for i in range(len(frame_nibbles)//2):
            octet = (frame_nibbles[2*i]<<4) + frame_nibbles[2*i+1]
            frame_octets.append(octet)

        # lanes' octets for a frame
        for i in range(nlanes):
            frame_lane_octets = frame_octets[i*octets_per_lane:
                                             (i+1)*octets_per_lane]
            lanes[i] = lanes[i] + [frame_lane_octets]

    return lanes

def lanes_to_samples(samples_per_frame, nlanes, nconverters, nbits, lanes):
    """
    inputs:
    - samples_per_frame: Number of samples per frame
    - nlanes:            Number of lanes per converter
    - nconverters:       Number of converters
    - nbits:             Number of convertion bits
    - lanes:             Lanes' octets organized in frames
                         lanes[i][j][k]: octet k of frame j of lane i
    output:
    - samples: Samples from converters:
               samples[i][j]: sample j of converter i

    cf section 5.1.3
    """
    assert nlanes == len(lanes)

    nibbles_per_word = ceil(nbits//4)
    octets_per_frame = samples_per_frame*nibbles_per_word//2
    octets_per_lane = octets_per_frame*nconverters//nlanes
    assert octets_per_lane > 0

    samples = [[]]*nconverters
    n = 0

    while n < len(lanes[0]):
        # frame's octets
        frame_octets = []
        for i in range(nlanes):
            frame_octets = frame_octets + lanes[i][n]
        n += 1

        # frame's nibbles
        frame_nibbles = []
        for octet in frame_octets:
            frame_nibbles = frame_nibbles + [(octet>>4) & 0xf, octet & 0xf]

        # frame's words
        frame_words = []
        for i in range(len(frame_nibbles)//nibbles_per_word):
            word = 0
            for j in range(nibbles_per_word):
                word |= (frame_nibbles[i*nibbles_per_word+j] << 4*(3-j))
            frame_words.append(word)

        # frame's samples
        frame_samples = frame_words # no control bits

        # converters' samples for a frame
        for i in range(nconverters):
            converter_samples = frame_samples[i*samples_per_frame:
                                             (i+1)*samples_per_frame]
            samples[i] = samples[i] + converter_samples

    return samples

class TransportLayer:
    def __init__(self, jesd_settings):
        self.jesd_settings = jesd_settings

    def encode(self, samples):
        return samples_to_lanes(self.jesd_settings.transport.s,
                                self.jesd_settings.phy.l,
                                self.jesd_settings.phy.m,
                                self.jesd_settings.phy.n,
                                samples)

    def decode(self, lanes):
        return lanes_to_samples(self.jesd_settings.transport.s,
                                self.jesd_settings.phy.l,
                                self.jesd_settings.phy.m,
                                self.jesd_settings.phy.n,
                                lanes)