# -*- coding: utf-8 -*-
# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
import argparse
import importlib
import json
import os
import random
import sys
from multiprocessing import Process, Queue

import cv2
import megengine as mge
import numpy as np
from megengine import jit
from megengine.data import DataLoader, SequentialSampler
from tqdm import tqdm

from official.vision.detection.tools.data_mapper import data_mapper
from official.vision.detection.tools.nms import py_cpu_nms

logger = mge.get_logger(__name__)

class DetEvaluator:
    def __init__(self, model):
        self.model = model

    @staticmethod
    def get_hw_by_short_size(im_height, im_width, short_size, max_size):
        """get height and width by short size

           Args:
               im_height(int): height of original image, e.g. 800
               im_width(int): width of original image, e.g. 1000
               short_size(int): short size of transformed image. e.g. 800
               max_size(int): max size of transformed image. e.g. 1333

           Returns:
               resized_height(int): height of transformed image
               resized_width(int): width of transformed image
        """

        im_size_min = np.min([im_height, im_width])
        im_size_max = np.max([im_height, im_width])
        scale = (short_size + 0.0) / im_size_min
        if scale * im_size_max > max_size:
            scale = (max_size + 0.0) / im_size_max

        resized_height, resized_width = (
            int(round(im_height * scale)),
            int(round(im_width * scale)),
        )
        return resized_height, resized_width

    @staticmethod
    def process_inputs(img, short_size, max_size, flip=False):
        original_height, original_width, _ = img.shape
        resized_height, resized_width = DetEvaluator.get_hw_by_short_size(
            original_height, original_width, short_size, max_size
        )
        resized_img = cv2.resize(
            img, (resized_width, resized_height), interpolation=cv2.INTER_LINEAR,
        )
        resized_img = cv2.flip(resized_img, 1) if flip else resized_img
        trans_img = np.ascontiguousarray(
            resized_img.transpose(2, 0, 1)[None, :, :, :], dtype=np.uint8
        )
        im_info = np.array(
            [(resized_height, resized_width, original_height, original_width)],
            dtype=np.float32,
        )
        return trans_img, im_info

    def predict(self, val_func):
        """
        Args:
            val_func(callable): model inference function

        Returns:
            results boxes: detection model output
        """
        model = self.model

        box_cls, box_delta = val_func()
        box_cls, box_delta = box_cls.numpy(), box_delta.numpy()
        dtboxes_all = list()
        all_inds = np.where(box_cls > model.cfg.test_cls_threshold)

        for c in range(0, model.cfg.num_classes):
            inds = np.where(all_inds[1] == c)[0]
            inds = all_inds[0][inds]
            scores = box_cls[inds, c]
            if model.cfg.class_aware_box:
                bboxes = box_delta[inds, c, :]
            else:
                bboxes = box_delta[inds, :]

            dtboxes = np.hstack((bboxes, scores[:, np.newaxis])).astype(np.float32)

            if dtboxes.size > 0:
                keep = py_cpu_nms(dtboxes, model.cfg.test_nms)
                dtboxes = np.hstack(
                    (dtboxes[keep], np.ones((len(keep), 1), np.float32) * c)
                ).astype(np.float32)
                dtboxes_all.extend(dtboxes)

        if len(dtboxes_all) > model.cfg.test_max_boxes_per_image:
            dtboxes_all = sorted(dtboxes_all, reverse=True, key=lambda i: i[4])[
                : model.cfg.test_max_boxes_per_image
            ]

        dtboxes_all = np.array(dtboxes_all, dtype=np.float)
        return dtboxes_all

    @staticmethod
    def format(results, cfg):
        dataset_class = data_mapper[cfg.test_dataset["name"]]

        all_results = []
        for record in results:
            image_filename = record["image_id"]
            boxes = record["det_res"]
            if len(boxes) <= 0:
                continue
            boxes[:, 2:4] = boxes[:, 2:4] - boxes[:, 0:2]
            for box in boxes:
                elem = dict()
                elem["image_id"] = image_filename
                elem["bbox"] = box[:4].tolist()
                elem["score"] = box[4]
                elem["category_id"] = dataset_class.classes_originID[
                    dataset_class.class_names[int(box[5])]
                ]
                all_results.append(elem)
        return all_results

    @staticmethod
    def vis_det(
        img,
        dets,
        is_show_label=True,
        classes=None,
        thresh=0.3,
        name="detection",
        return_img=True,
    ):
        img = np.array(img)
        colors = dict()
        font = cv2.FONT_HERSHEY_SIMPLEX

        for det in dets:
            bb = det[:4].astype(int)
            if is_show_label:
                cls_id = int(det[5])
                score = det[4]

                if cls_id == 0:
                    continue
                if score > thresh:
                    if cls_id not in colors:
                        colors[cls_id] = (
                            random.random() * 255,
                            random.random() * 255,
                            random.random() * 255,
                        )

                    cv2.rectangle(
                        img, (bb[0], bb[1]), (bb[2], bb[3]), colors[cls_id], 3
                    )

                    if classes and len(classes) > cls_id:
                        cls_name = classes[cls_id]
                    else:
                        cls_name = str(cls_id)
                    cv2.putText(
                        img,
                        "{:s} {:.3f}".format(cls_name, score),
                        (bb[0], bb[1] - 2),
                        font,
                        0.5,
                        (255, 255, 255),
                        1,
                    )
            else:
                cv2.rectangle(img, (bb[0], bb[1]), (bb[2], bb[3]), (0, 0, 255), 2)

        if return_img:
            return img
        cv2.imshow(name, img)
        while True:
            c = cv2.waitKey(100000)
            if c == ord("d"):
                return None
            elif c == ord("n"):
                break

def build_dataloader(rank, world_size, data_dir, cfg):
    val_dataset = data_mapper[cfg.test_dataset["name"]](
        os.path.join(data_dir, cfg.test_dataset["name"], cfg.test_dataset["root"]),
        os.path.join(data_dir, cfg.test_dataset["name"], cfg.test_dataset["ann_file"]),
        order=["image", "info"],
    )
    val_sampler = SequentialSampler(val_dataset, 1, world_size=world_size, rank=rank)
    val_dataloader = DataLoader(val_dataset, sampler=val_sampler, num_workers=2)
    return val_dataloader

def worker(
    net_file, model_file, data_dir, worker_id, total_worker, result_queue,
):
    """
    :param net_file: network description file
    :param model_file: file of dump weights
    :param data_dir: the dataset directory
    :param worker_id: the index of the worker
    :param total_worker: number of gpu for evaluation
    :param result_queue: processing queue
    """
    os.environ["CUDA_VISIBLE_DEVICES"] = str(worker_id)

    @jit.trace(symbolic=True, opt_level=2)
    def val_func():
        pred = model(model.inputs)
        return pred

    sys.path.insert(0, os.path.dirname(net_file))
    current_network = importlib.import_module(os.path.basename(net_file).split(".")[0])
    model = current_network.Net(current_network.Cfg(), batch_size=1)
    model.eval()
    evaluator = DetEvaluator(model)
    state_dict = mge.load(model_file)
    if "state_dict" in state_dict:
        state_dict = state_dict["state_dict"]
    model.load_state_dict(state_dict)

    loader = build_dataloader(worker_id, total_worker, data_dir, model.cfg)
    for data_dict in loader:
        data, im_info = DetEvaluator.process_inputs(
            data_dict[0][0],
            model.cfg.test_image_short_size,
            model.cfg.test_image_max_size,
        )
        model.inputs["im_info"].set_value(im_info)
        model.inputs["image"].set_value(data.astype(np.float32))

        pred_res = evaluator.predict(val_func)
        result_queue.put_nowait(
            {
                "det_res": pred_res,
                "image_id": int(data_dict[1][2][0].split(".")[0].split("_")[-1]),
            }
        )

def make_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("-b", "--batch_size", default=1, type=int)
    parser.add_argument("-n", "--ngpus", default=1, type=int)
    parser.add_argument(
        "-f", "--file", default="net.py", type=str, help="net description file"
    )
    parser.add_argument("-d", "--dataset_dir", default="/data/datasets", type=str)
    parser.add_argument("-se", "--start_epoch", default=-1, type=int)
    parser.add_argument("-ee", "--end_epoch", default=-1, type=int)
    parser.add_argument("-m", "--model", default=None, type=str)
    return parser

def main():
    # pylint: disable=import-outside-toplevel
    from pycocotools.coco import COCO
    from pycocotools.cocoeval import COCOeval

    parser = make_parser()
    args = parser.parse_args()

    if args.end_epoch == -1:
        args.end_epoch = args.start_epoch

    for epoch_num in range(args.start_epoch, args.end_epoch + 1):
        if args.model:
            model_file = args.model
        else:
            model_file = "log-of-{}/epoch_{}.pkl".format(
                os.path.basename(args.file).split(".")[0], epoch_num
            )
        logger.info("Load Model : %s completed", model_file)

        results_list = list()
        result_queue = Queue(2000)
        procs = []
        for i in range(args.ngpus):
            proc = Process(
                target=worker,
                args=(
                    args.file,
                    model_file,
                    args.dataset_dir,
                    i,
                    args.ngpus,
                    result_queue,
                ),
            )
            proc.start()
            procs.append(proc)

        for _ in tqdm(range(5000)):
            results_list.append(result_queue.get())
        for p in procs:
            p.join()

        sys.path.insert(0, os.path.dirname(args.file))
        current_network = importlib.import_module(
            os.path.basename(args.file).split(".")[0]
        )
        cfg = current_network.Cfg()
        all_results = DetEvaluator.format(results_list, cfg)
        json_path = "log-of-{}/epoch_{}.json".format(
            os.path.basename(args.file).split(".")[0], epoch_num
        )
        all_results = json.dumps(all_results)

        with open(json_path, "w") as fo:
            fo.write(all_results)
        logger.info("Save to %s finished, start evaluation!", json_path)

        eval_gt = COCO(
            os.path.join(
                args.dataset_dir, cfg.test_dataset["name"], cfg.test_dataset["ann_file"]
            )
        )
        eval_dt = eval_gt.loadRes(json_path)
        cocoEval = COCOeval(eval_gt, eval_dt, iouType="bbox")
        cocoEval.evaluate()
        cocoEval.accumulate()
        cocoEval.summarize()
        metrics = [
            "AP",
            "AP@0.5",
            "AP@0.75",
            "APs",
            "APm",
            "APl",
            "AR@1",
            "AR@10",
            "AR@100",
            "ARs",
            "ARm",
            "ARl",
        ]
        logger.info("mmAP".center(32, "-"))
        for i, m in enumerate(metrics):
            logger.info("|\t%s\t|\t%.03f\t|", m, cocoEval.stats[i])
        logger.info("-" * 32)

if __name__ == "__main__":
    main()


from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals
from __future__ import print_function

import torch as th
from torch.utils.data import Dataset
import pandas as pd
import os
import numpy as np
import librosa

from everything_at_once.dataset.utils import _tokenize_text, create_audio_features, \
    create_text_features, _crop_audio_from_mel_spec, _get_video

class HowTo_Dataset(Dataset):
    """HowTo100M dataset"""

    def __init__(
            self,
            csv,
            caption,
            we,
            features_path,
            features_path_audio,
            features_path_3D=None,
            min_time=8.0,
            n_video_tokens=None,
            feature_per_sec_2D=1.0,
            feature_per_sec_3D=24.0 / 16.0,
            we_dim=300,
            max_words=20,
            n_clips=1,
            num_audio_STFT_frames=768,
            video_sampling_strategy='clip',
            flag_text_audio_misaligned=False,
            use_2D=True,
            use_3D=True,
    ):
        """
        Args:
        """
        self.csv = pd.read_csv(csv)
        self.caption = caption
        self.we = we

        self.features_path = features_path
        self.features_path_audio = features_path_audio
        self.features_path_3D = features_path_3D if features_path_3D is not None else features_path

        self.video_sampling_strategy = video_sampling_strategy
        self.flag_text_audio_misaligned = flag_text_audio_misaligned

        self.min_time = min_time
        self.n_video_tokens = n_video_tokens if n_video_tokens is not None else int(1.5 * min_time)
        self.feature_per_sec_2D = feature_per_sec_2D
        self.feature_per_sec_3D = feature_per_sec_3D
        self.we_dim = we_dim
        self.max_words = max_words
        self.num_audio_STFT_frames = num_audio_STFT_frames
        self.n_clips = n_clips
        self.use_2D = use_2D
        self.use_3D = use_3D

    def __len__(self):
        return len(self.csv)

    def __getitem__(self, idx):
        # load video
        vid_path = self.csv['path'].values[idx].replace("None/", "")
        video_id = vid_path.split("/")[-1]

        features_2d = None
        features_3d = None
        if self.use_2D:
            path = os.path.join(self.features_path, vid_path, video_id + "_2d.npz")
            if not os.path.exists(path):
                path = os.path.join(self.features_path, vid_path + "_2d.npz")

            features_2d = np.load(path)['features']
            features_2d = th.from_numpy(features_2d).float()
        if self.use_3D:
            path = os.path.join(self.features_path_3D, vid_path, video_id + "_3d.npz")
            if not os.path.exists(path):
                path = os.path.join(self.features_path_3D, vid_path, video_id + ".npz")
            features_3d = np.load(path)['features']
            features_3d = th.from_numpy(features_3d).float()

        # load audio
        audio_path = os.path.join(self.features_path_audio, vid_path, video_id + "_spec.npz")
        mel_spec = np.load(audio_path)['arr_0']

        # load text
        caption = self.caption[video_id]

        video, video_mask, audio, audio_mask, text, text_mask, raw_text, audio_STFT_nframes, paths, ids, starts, ends = \
            _sample_video_audio_text(caption, mel_spec, features_2d, features_3d, vid_path, self.n_clips,
                                     self.feature_per_sec_2D, self.feature_per_sec_3D,
                                     self.min_time, self.max_words, self.we, self.we_dim,
                                     self.num_audio_STFT_frames, self.n_video_tokens,
                                     video_sampling_strategy=self.video_sampling_strategy,
                                     flag_text_audio_misaligned=self.flag_text_audio_misaligned)

        return {'video': video, 'audio': audio, 'text': text,
                'video_mask': video_mask, 'audio_mask': audio_mask, 'text_mask': text_mask,
                'raw_text': raw_text,
                'unroll_clips': 1, 'audio_STFT_nframes': audio_STFT_nframes,
                'meta': {'paths': paths, 'ids': ids, 'dataset': ['HowTo100M'] * len(video)}}

def _sample_video_audio_text(caption, mel_spec, features_2d, features_3d, vid_path, n_clips, fps_2d, fps_3d,
                             min_time, max_words, we, we_dim, num_audio_STFT_frames, n_video_tokens,
                             video_sampling_strategy='clip',
                             flag_text_audio_misaligned=False,
                             ):
    audio, audio_mask, audio_STFT_nframes, starts, ends, text, text_mask, raw_text = \
        _sample_audio_and_text(caption, n_clips, mel_spec,
                               min_time, max_words, we, we_dim, num_audio_STFT_frames,
                               flag_text_audio_misaligned=flag_text_audio_misaligned)

    video, video_mask = _get_video(features_2d, features_3d, fps_2d, fps_3d, starts, ends, n_video_tokens, video_sampling_strategy)

    audio_STFT_nframes = -np.ones(len(audio_STFT_nframes))

    paths = [vid_path] * len(video)
    ids = [vid_path + str(start) for start in starts]

    return video, video_mask, audio, audio_mask, text, text_mask, raw_text, audio_STFT_nframes, paths, ids, starts, ends

def _sample_audio_and_text(caption, n_clips, mel_spec, min_time, max_words, we, we_dim, num_audio_STFT_frames,
                           flag_text_audio_misaligned=False):
    starts = np.zeros(n_clips)
    ends = np.zeros(n_clips)
    text = [0 for _ in range(n_clips)]
    text_mask = [0 for _ in range(n_clips)]
    raw_text = [0 for _ in range(n_clips)]
    audio = [0 for _ in range(n_clips)]
    audio_mask = [0 for _ in range(n_clips)]
    audio_STFT_nframes = np.zeros(n_clips)

    for i in range(n_clips):
        audio[i], audio_mask[i], audio_STFT_nframes[i], starts[i], ends[i], text[i], text_mask[i], raw_text[i] = \
                _get_single_random_audio_and_text(caption, mel_spec, min_time, max_words, we, we_dim, num_audio_STFT_frames,
                                                  flag_text_audio_misaligned=flag_text_audio_misaligned)
    audio = th.stack(audio, dim=0)
    audio_mask = th.stack(audio_mask, dim=0)
    text = th.stack(text, dim=0)
    text_mask = th.stack(text_mask, dim=0)
    return audio, audio_mask, audio_STFT_nframes, starts, ends, text, text_mask, raw_text

def _get_single_random_audio_and_text(caption, mel_spec, min_time,  max_words, we, we_dim, num_audio_STFT_frames,
                                      flag_text_audio_misaligned=False):
    video_duration_seconds = int(librosa.core.frames_to_time(mel_spec.shape[1], sr=16000, hop_length=160, n_fft=400))
    start = np.random.rand() * max(0, video_duration_seconds - min_time)
    end = start + min_time

    # Search if there is ASR narration in the sampled clip
    ind_start = max(0, np.searchsorted(caption['start'], start, side='left') - 1)
    ind_end = max(0, np.searchsorted(caption['start'], end, side='left') - 1)
    if caption['start'][ind_start] <= start <= caption['end'][ind_start]:
        return _get_text_and_audio_by_ind(caption, mel_spec, min_time, max_words, we, we_dim, num_audio_STFT_frames,
                                          ind=ind_start,
                                          flag_text_audio_misaligned=flag_text_audio_misaligned)
    elif caption['start'][ind_end] <= end <= caption['end'][ind_end]:
        return _get_text_and_audio_by_ind(caption, mel_spec, min_time, max_words, we, we_dim, num_audio_STFT_frames,
                                          ind=ind_end,
                                          flag_text_audio_misaligned=flag_text_audio_misaligned)
    else:
        words = [] # no ASR in this clip
        start_frame = max(0, librosa.core.time_to_frames(start, sr=16000, hop_length=160, n_fft=400))
        mel_spec = mel_spec[:, start_frame:start_frame + num_audio_STFT_frames]

        # padded audio features
        audio, audio_mask, audio_STFT_nframes = create_audio_features(mel_spec, num_audio_STFT_frames)

        # create padded text features
        text, text_mask, raw_text = create_text_features(words, max_words, we, we_dim)

        return audio, audio_mask, audio_STFT_nframes, start, end, text, text_mask, raw_text

def _get_text_and_audio_by_ind(caption, mel_spec, min_time, max_words, we, we_dim, num_audio_STFT_frames, ind=None,
                               flag_text_audio_misaligned=False):
    if ind is None:
        n_caption = len(caption['start'])
        ind = np.random.choice(range(n_caption))

    start, end = ind, ind
    words = _tokenize_text(caption['text'][ind])
    diff = caption['end'][end] - caption['start'][start]
    # Extend the video clip if shorter than the minimum desired clip duration
    while diff < min_time:
        if start > 0 and end < len(caption['end']) - 1:
            next_words = _tokenize_text(caption['text'][end + 1])
            prev_words = _tokenize_text(caption['text'][start - 1])
            d1 = caption['end'][end + 1] - caption['start'][start]
            d2 = caption['end'][end] - caption['start'][start - 1]
            # Use the closest neighboring video clip
            if d2 <= d1:
                start -= 1
                words.extend(prev_words)
            else:
                end += 1
                words.extend(next_words)
        # If no video clips after it, use the clip before it
        elif start > 0:
            words.extend(_tokenize_text(caption['text'][start - 1]))
            start -= 1
        # If no video clips before it, use the clip after it.
        elif end < len(caption['end']) - 1:
            words.extend(_tokenize_text(caption['text'][end + 1]))
            end += 1
        # If there's no clips before or after
        else:
            break
        diff = caption['end'][end] - caption['start'][start]

    start, end = caption['start'][start], caption['end'][end]
    if flag_text_audio_misaligned:
        if np.random.rand() > 0:
            start_audio = start + 0.5 * min_time
            end_audio = start_audio + min_time
        else:
            start_audio = max(0, start - 0.5 * min_time)
            end_audio = start_audio + min_time
    else:
        start_audio = start
        end_audio = end

    # create padded audio features
    mel_spec = _crop_audio_from_mel_spec(start_audio, end_audio, mel_spec)
    audio, audio_mask, audio_STFT_nframes = create_audio_features(mel_spec, num_audio_STFT_frames)

    # create padded text features
    text, text_mask, raw_text = create_text_features(words, max_words, we, we_dim)

    return audio, audio_mask, audio_STFT_nframes, start, end, text, text_mask, raw_text


from typing import Any, Mapping, NamedTuple, Tuple

import haiku as hk
import jax
import optax
import typer
from haiku._src.data_structures import (  # pylint: disable=no-name-in-module
    FlatMap,
)

from jdetr import data, losses, models
from jdetr._typing import JaxArray, Number, PRNGKey

def build_forward_fn(num_classes: int, max_boxes: int = 100):
    def detr_forward(
        x: JaxArray, is_training: bool = True
    ) -> Tuple[JaxArray, JaxArray]:
        backbone = models.ResNet50(num_classes=0, resnet_v2=True)
        transformer = models.Transformer(
            feature_dim=256,
            num_heads=8,
            num_encoder_layers=6,
            num_decoder_layers=6,
            feedforward_dim=2048,
            dropout_rate=0.1,
        )
        detr = models.DETR(
            backbone, transformer, num_classes=num_classes, num_queries=max_boxes
        )
        return detr(x, is_training=is_training)

    return detr_forward

class State(NamedTuple):
    step: int
    trainable_params: FlatMap
    stateful_params: FlatMap
    opt_state: FlatMap
    rng: PRNGKey

class Trainer:
    def __init__(self, forward_fn, loss_fn, optimizer):
        self.forward_fn = forward_fn
        self.loss_fn = loss_fn
        self.optimizer = optimizer

    # @ft.partial(jax.jit, static_argnums=(0,))
    def init_states(self, rng: PRNGKey, batch: data.DataTuple) -> State:
        out_rng, init_rng = jax.random.split(rng)
        images, _, _ = batch
        trainable_params, stateful_params = self.forward_fn.init(init_rng, images)
        opt_state = self.optimizer.init(trainable_params)
        return State(
            step=0,
            trainable_params=trainable_params,
            stateful_params=stateful_params,
            opt_state=opt_state,
            rng=out_rng,
        )

    def forward(
        self,
        trainable_params: FlatMap,
        stateful_params: FlatMap,
        rng: PRNGKey,
        batch: data.DataTuple,
    ) -> Tuple[JaxArray, Mapping[str, Any]]:
        images, boxes, labels = batch
        (boxes_pred, labels_pred), stateful_params = self.forward_fn.apply(
            trainable_params, stateful_params, rng, images
        )
        # TODO Add pmap-averaging

        batch_num_valid_boxes = (labels > 0).sum()
        loss = (
            self.loss_fn(boxes_pred, boxes, labels_pred, labels) / batch_num_valid_boxes
        )
        return loss, {
            "loss": loss,
            "boxes_pred": boxes_pred,
            "labels_pred": labels_pred,
            "stateful_params": stateful_params,
            "num_valid_boxes": batch_num_valid_boxes,
        }

    def update_state(
        self, state: State, batch: data.DataTuple
    ) -> Tuple[State, Mapping[str, Number]]:
        new_rng_state, rng_state = jax.random.split(state.rng)
        grads, logs = jax.grad(self.forward, has_aux=True)(
            state.trainable_params, state.stateful_params, rng_state, batch
        )
        # TODO Add pmap grad-averaging
        updates, opt_state = self.optimizer.update(grads, state.opt_state)
        trainable_params = optax.apply_updates(state.trainable_params, updates)

        state = State(
            state.step + 1,
            trainable_params,
            logs.pop("stateful_params"),
            opt_state,
            new_rng_state,
        )

        return state, logs

def train(
    batch_size_per_device: int = 4,
    data_dir: str = "__pycache__/coco",
    max_boxes: int = 100,
    learning_rate: float = 3e-4,
):
    batch_size = batch_size_per_device * 1

    # FIXME Do sth with validation data
    train_data, _ = data.get_coco_datagen(batch_size, data_dir, max_boxes=max_boxes)
    forward_fn = hk.transform_with_state(
        build_forward_fn(train_data.n_classes, max_boxes=max_boxes)
    )
    loss_fn = jax.vmap(losses.SetCriterion())
    optimizer = optax.chain(
        optax.clip_by_global_norm(1), optax.adamw(learning_rate, b1=0.9, b2=0.99)
    )

    rng = jax.random.PRNGKey(428)
    (batch,) = train_data.data.take(1).as_numpy_iterator()
    trainer = Trainer(forward_fn, loss_fn, optimizer)
    state = trainer.init_states(rng, batch)
    for batch in train_data.data.as_numpy_iterator():
        state, logs = trainer.update_state(state, batch)
        print(logs)

if __name__ == "__main__":
    app = typer.Typer()
    app.command()(train)
    app()


"""A variant of anchor_head_single.

The differences are as follows:
* two more options: num_convs, GN
* apply two split convs for regression outputs and classification outputs
* when num_convs == 0, this module should be almost the same as anchor_head_single
* in conv_box/cls, the kernel size is modified to 3 instead of 1
"""

import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn

from .anchor_head_template import AnchorHeadTemplate

def convbn(in_planes, out_planes, kernel_size, stride, pad, dilation, gn=False, groups=32):
    return nn.Sequential(nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=dilation if dilation > 1 else pad, dilation=dilation, bias=False),
                         nn.BatchNorm2d(out_planes) if not gn else nn.GroupNorm(groups, out_planes))

class DetHead(AnchorHeadTemplate):
    def __init__(self, model_cfg, input_channels, num_class, class_names, grid_size, point_cloud_range,
                 predict_boxes_when_training=True):
        super().__init__(
            model_cfg=model_cfg, num_class=num_class, class_names=class_names, grid_size=grid_size, point_cloud_range=point_cloud_range,
            predict_boxes_when_training=predict_boxes_when_training
        )

        self.num_anchors_per_location = sum(self.num_anchors_per_location)
        self.num_convs = model_cfg.NUM_CONVS
        self.GN = model_cfg.GN
        self.xyz_for_angles = model_cfg.xyz_for_angles
        self.hwl_for_angles = model_cfg.hwl_for_angles

        if self.num_convs > 0:
            self.rpn3d_cls_convs = []
            self.rpn3d_bbox_convs = []
            for _ in range(self.num_convs):
                self.rpn3d_cls_convs.append(
                    nn.Sequential(
                        convbn(input_channels, input_channels, 3, 1, 1, 1, gn=self.GN),
                        nn.ReLU(inplace=True))
                )
                self.rpn3d_bbox_convs.append(
                    nn.Sequential(
                        convbn(input_channels, input_channels, 3, 1, 1, 1, gn=self.GN),
                        nn.ReLU(inplace=True))
                )
            assert len(self.rpn3d_cls_convs) == self.num_convs
            assert len(self.rpn3d_bbox_convs) == self.num_convs
            self.rpn3d_cls_convs = nn.Sequential(*self.rpn3d_cls_convs)
            self.rpn3d_bbox_convs = nn.Sequential(*self.rpn3d_bbox_convs)

        cls_feature_channels = input_channels
        cls_groups = 1

        self.conv_cls = nn.Conv2d(
            cls_feature_channels, self.num_anchors_per_location * self.num_class,
            kernel_size=3, padding=1, stride=1, groups=cls_groups
        )
        if self.xyz_for_angles and self.hwl_for_angles:
            box_dim = self.num_anchors_per_location * self.box_coder.code_size
        elif not self.xyz_for_angles and not self.hwl_for_angles:
            box_dim = self.num_class * 6 + self.num_anchors_per_location * (self.box_coder.code_size - 6)
        else:
            box_dim = self.num_class * 3 + self.num_anchors_per_location * (self.box_coder.code_size - 3)
        self.conv_box = nn.Conv2d(
            input_channels, box_dim,
            kernel_size=3, padding=1, stride=1
        )

        self.num_angles = self.num_anchors_per_location // self.num_class

        if self.model_cfg.get('USE_DIRECTION_CLASSIFIER', None) is not None:
            self.conv_dir_cls = nn.Conv2d(
                cls_feature_channels,
                self.num_anchors_per_location * self.model_cfg.NUM_DIR_BINS,
                kernel_size=1,
                groups=cls_groups
            )
        else:
            self.conv_dir_cls = None

        self.init_weights()

    def init_weights(self):
        pi = 0.01
        nn.init.normal_(self.conv_cls.weight, std=0.1)
        nn.init.normal_(self.conv_box.weight, std=0.02)
        nn.init.constant_(self.conv_cls.bias, -np.log((1 - pi) / pi))
        # nn.init.normal_(self.conv_box.weight, mean=0, std=0.001)

    def forward(self, data_dict):
        # NOTE: clear forward ret dict to avoid potential bugs
        self.forward_ret_dict.clear()

        spatial_features_2d = data_dict['spatial_features_2d']

        if self.do_feature_imitation and self.training:
            if 'gt_boxes' in data_dict:
                self.forward_ret_dict['gt_boxes'] = data_dict['gt_boxes']
            self.forward_ret_dict['imitation_features_pairs'] = []
            imitation_conv_layers = [self.conv_imitation] if len(self.imitation_configs) == 1 else self.conv_imitation
            for cfg, imitation_conv in zip(self.imitation_configs, imitation_conv_layers):
                lidar_feature_name = cfg.lidar_feature_layer
                stereo_feature_name = cfg.stereo_feature_layer
                self.forward_ret_dict['imitation_features_pairs'].append(
                    dict(
                        config=cfg,
                        stereo_feature_name=stereo_feature_name,
                        lidar_feature_name=lidar_feature_name,
                        gt=data_dict['lidar_outputs'][lidar_feature_name],
                        pred=imitation_conv(data_dict[stereo_feature_name])
                    )
                )

            # for k in data_dict:
            #     if k in ["lidar_batch_cls_preds", "lidar_batch_box_preds"]:
            #         self.forward_ret_dict[k] = data_dict[k]

        cls_features = spatial_features_2d
        reg_features = spatial_features_2d
        if self.num_convs > 0:
            cls_features = self.rpn3d_cls_convs(cls_features)
            reg_features = self.rpn3d_bbox_convs(reg_features)
        box_preds = self.conv_box(reg_features)
        box_preds = box_preds.permute(0, 2, 3, 1).contiguous()  # [N, H, W, C]

        data_dict['reg_features'] = reg_features

        cls_preds = self.conv_cls(cls_features)
        cls_preds = cls_preds.permute(0, 2, 3, 1).contiguous()  # [N, H, W, C]

        if not self.xyz_for_angles or not self.hwl_for_angles:
            # TODO: here we assume that for each class, there are only anchors with difference angles
            if self.xyz_for_angles:
                xyz_dim = self.num_anchors_per_location * 3
                xyz_shapes = (self.num_class, self.num_anchors_per_location // self.num_class, 3)
            else:
                xyz_dim = self.num_class * 3
                xyz_shapes = (self.num_class, 1, 3)
            if self.hwl_for_angles:
                hwl_dim = self.num_anchors_per_location * 3
                hwl_shapes = (self.num_class, self.num_anchors_per_location // self.num_class, 3)
            else:
                hwl_dim = self.num_class * 3
                hwl_shapes = (self.num_class, 1, 3)
            rot_dim = self.num_anchors_per_location * (self.box_coder.code_size - 6)
            rot_shapes = (self.num_class, self.num_anchors_per_location // self.num_class, (self.box_coder.code_size - 6))
            assert box_preds.shape[-1] == xyz_dim + hwl_dim + rot_dim
            xyz_preds, hwl_preds, rot_preds = torch.split(box_preds, [xyz_dim, hwl_dim, rot_dim], dim=-1)
            # anchors [Nz, Ny, Nx, N_cls*N_size=3*1, N_rot, 7]
            xyz_preds = xyz_preds.view(*xyz_preds.shape[:3], *xyz_shapes)
            hwl_preds = hwl_preds.view(*hwl_preds.shape[:3], *hwl_shapes)
            rot_preds = rot_preds.view(*rot_preds.shape[:3], *rot_shapes)
            # expand xyz and hwl
            if not self.xyz_for_angles:
                xyz_preds = xyz_preds.repeat(1, 1, 1, 1, rot_preds.shape[4] // xyz_preds.shape[4], 1)
            if not self.hwl_for_angles:
                hwl_preds = hwl_preds.repeat(1, 1, 1, 1, rot_preds.shape[4] // hwl_preds.shape[4], 1)
            box_preds = torch.cat([xyz_preds, hwl_preds, rot_preds], dim=-1)
            box_preds = box_preds.view(*box_preds.shape[:3], -1)

        self.forward_ret_dict['cls_preds'] = cls_preds
        self.forward_ret_dict['box_preds'] = box_preds

        if 'valids' in data_dict:
            self.forward_ret_dict['valids'] = data_dict['valids'].any(1)

        if self.conv_dir_cls is not None:
            dir_cls_preds = self.conv_dir_cls(cls_features)
            dir_cls_preds = dir_cls_preds.permute(0, 2, 3, 1).contiguous()
            self.forward_ret_dict['dir_cls_preds'] = dir_cls_preds
        else:
            dir_cls_preds = None

        if self.training or 'gt_boxes' in data_dict:
            targets_dict = self.assign_targets(
                gt_boxes=data_dict['gt_boxes']
            )
            data_dict.update(targets_dict)
            data_dict['anchors'] = self.anchors
            self.forward_ret_dict.update(targets_dict)

        if not self.training or self.predict_boxes_when_training:
            batch_cls_preds, batch_box_preds = self.generate_predicted_boxes(
                batch_size=data_dict['batch_size'],
                cls_preds=cls_preds, box_preds=box_preds, dir_cls_preds=dir_cls_preds
            )
            data_dict['batch_cls_preds'] = batch_cls_preds
            data_dict['batch_box_preds'] = batch_box_preds
            # TODO: check the code here, we add sigmoid in the generate predicted boxes, so set normalized to be True
            data_dict['cls_preds_normalized'] = False

        return data_dict


import numpy as np
import copy as cp
from dicepool import dpool

class Skills(object):
    """
    Class holding the attributes of a character
    """
    def __init__(self,
                 academics = 0, computer = 0, crafts = 0, investigation = 0,
                 medicine = 0, occult = 0, politics = 0, science = 0,
                 athletics = 0, brawl = 0, drive = 0, firearms = 0, larceny = 0,
                 stealth = 0, survival = 0, weaponry = 0,
                 animal_ken = 0, empathy = 0, expression = 0, intimidation = 0,
                 persuasion = 0, socialize = 0, streetwise = 0, subterfuge = 0
                 ):
        self.skills = {}

        self.list = {}
        self.list['physical'] = ['academics', 'computer', 'crafts', 'investigation',
        'medicine', 'occult', 'politics', 'science']

        self.list['psychical'] = ['athletics', 'brawl', 'drive', 'firearms',
        'larceny', 'stealth', 'survival', 'weaponry']

        self.list['social'] = ['animal_ken', 'empathy', 'expression', 'intimidation',
        'persuasion', 'socialize', 'streetwise', 'subterfuge']

        self.skills['academics'] = academics
        self.skills['computer'] = computer
        self.skills['crafts'] = crafts
        self.skills['investigation'] = investigation
        self.skills['medicine'] = medicine
        self.skills['occult'] = occult
        self.skills['politics'] = politics
        self.skills['science'] = science

        self.skills['athletics'] = athletics
        self.skills['brawl'] = brawl
        self.skills['drive'] = drive
        self.skills['firearms'] = firearms
        self.skills['larceny'] = larceny
        self.skills['stealth'] = stealth
        self.skills['survival'] = survival
        self.skills['weaponry'] = weaponry

        self.skills['animal_ken'] = animal_ken
        self.skills['empathy'] = empathy
        self.skills['expression'] = expression
        self.skills['intimidation'] = intimidation
        self.skills['persuasion'] = persuasion
        self.skills['socialize'] = socialize
        self.skills['streetwise'] = streetwise
        self.skills['subterfuge'] = subterfuge

    def get_sum(self, name):
        return sum([self.skills[key] for key in self.list[name]])

    def get_sum_psychical(self):
        return self.get_sum('psychical')

    def get_sum_physical(self):
        return self.get_sum('physical')

    def get_sum_social(self):
        return self.get_sum('social')

    def finish(self):
        for key in self.list:
            while self.get_sum(name=key) < 4:
                sk = np.random.choice(self.list[key])
                self.skills[sk] = self.skills[sk] + 1

        rem = []
        for key in self.list:
            if self.get_sum(key) == 4:
                rem.append(key)

        selected = np.random.choice(rem)

        remaining = [key for key in self.list]
        remaining.remove(selected)

        for key in remaining:
            while self.get_sum(name=key) < 7:
                sk = np.random.choice(self.list[key])
                self.skills[sk] = self.skills[sk] + 1

        rem = []
        for key in remaining:
            if self.get_sum(key) == 7:
                rem.append(key)

        selected = np.random.choice(rem)
        remaining.remove(selected)

        for key in remaining:
            while self.get_sum(name=key) < 11:
                sk = np.random.choice(self.list[key])
                self.skills[sk] = self.skills[sk] + 1

    def missing(self):
        physical = self.get_sum_physical()
        psychical = self.get_sum_psychical()
        social = self.get_sum_social()

        maximum = max(physical, psychical, social)
        minimum = min(physical, psychical, social)
        middle = physical + psychical + social - minimum - maximum

        return maximum - 11, middle - 7, minimum - 4

    def skill(self, skill):
        return self.skills[skill]

if __name__ == "__main__":

    sk = Skills()

    print(sk.get_sum_psychical())
    print(sk.get_sum_social())
    print(sk.get_sum_physical())

    sk.finish()

    print(sk.get_sum_psychical())
    print(sk.get_sum_social())
    print(sk.get_sum_physical())

    print(
        sk.skill('academics'), sk.skill('computer'), sk.skill('crafts'), sk.skill('investigation'),
        sk.skill('medicine'), sk.skill('occult'), sk.skill('politics'), sk.skill('science'),
        sk.skill('athletics'), sk.skill('brawl'), sk.skill('drive'), sk.skill('firearms'), sk.skill('larceny'),
        sk.skill('stealth'), sk.skill('survival'), sk.skill('weaponry'), sk.skill('animal_ken'),
        sk.skill('empathy'), sk.skill('expression'), sk.skill('intimidation'), sk.skill('persuasion'),
        sk.skill('socialize'), sk.skill('streetwise'), sk.skill('subterfuge')
        )

    print (sk.missing())


"""Tests for the move scheduler."""
import pytest
from typing import List, Any
from numpy import float64
from mock import AsyncMock, call, MagicMock
from opentrons_hardware.firmware_bindings import ArbitrationId, ArbitrationIdParts

from opentrons_hardware.firmware_bindings.constants import NodeId
from opentrons_hardware.drivers.can_bus.can_messenger import MessageListenerCallback
from opentrons_hardware.firmware_bindings.messages.message_definitions import (
    AddLinearMoveRequest,
    HomeRequest,
    MoveCompleted,
)
from opentrons_hardware.firmware_bindings.messages.payloads import (
    AddLinearMoveRequestPayload,
    MoveCompletedPayload,
    ExecuteMoveGroupRequestPayload,
    HomeRequestPayload,
)
from opentrons_hardware.hardware_control.constants import (
    interrupts_per_sec,
)
from opentrons_hardware.hardware_control.motion import (
    MoveGroups,
    MoveGroupSingleAxisStep,
    MoveType,
    MoveStopCondition,
)
from opentrons_hardware.hardware_control.move_group_runner import (
    MoveGroupRunner,
    MoveScheduler,
    _CompletionPacket,
)
from opentrons_hardware.hardware_control.types import NodeMap
from opentrons_hardware.firmware_bindings.messages import (
    message_definitions as md,
    MessageDefinition,
)
from opentrons_hardware.firmware_bindings.utils import (
    UInt8Field,
    Int32Field,
    UInt32Field,
)

def calc_duration(step: MoveGroupSingleAxisStep) -> int:
    """Calculate duration."""
    return int(step.duration_sec * interrupts_per_sec)

def calc_velocity(step: MoveGroupSingleAxisStep) -> int:
    """Calculate velocity."""
    return int(step.velocity_mm_sec / interrupts_per_sec * (2**31))

def calc_acceleration(step: MoveGroupSingleAxisStep) -> int:
    """Calculate acceleration."""
    return int(
        step.acceleration_mm_sec_sq
        / interrupts_per_sec
        / interrupts_per_sec
        * (2**31)
    )

@pytest.fixture
def mock_can_messenger() -> AsyncMock:
    """Mock communication."""
    return AsyncMock()

@pytest.fixture
def move_group_single() -> MoveGroups:
    """Move group with one move."""
    return [
        [
            {
                NodeId.head: MoveGroupSingleAxisStep(
                    distance_mm=float64(246),
                    velocity_mm_sec=float64(2),
                    duration_sec=float64(123),
                )
            }
        ]
    ]

@pytest.fixture
def move_group_home_single() -> MoveGroups:
    """Home Request."""
    return [
        # Group 0
        [
            {
                NodeId.head: MoveGroupSingleAxisStep(
                    distance_mm=float64(0),
                    velocity_mm_sec=float64(235),
                    duration_sec=float64(2142),
                    acceleration_mm_sec_sq=float64(1000),
                    stop_condition=MoveStopCondition.limit_switch,
                    move_type=MoveType.home,
                ),
            }
        ]
    ]

@pytest.fixture
def move_group_multiple() -> MoveGroups:
    """Move group with multiple moves."""
    return [
        # Group 0
        [
            {
                NodeId.head: MoveGroupSingleAxisStep(
                    distance_mm=float64(229),
                    velocity_mm_sec=float64(235),
                    duration_sec=float64(2142),
                    acceleration_mm_sec_sq=float64(1000),
                ),
            }
        ],
        # Group 1
        [
            {
                NodeId.gantry_x: MoveGroupSingleAxisStep(
                    distance_mm=float64(522),
                    velocity_mm_sec=float64(22),
                    duration_sec=float64(1),
                    acceleration_mm_sec_sq=float64(1000),
                ),
                NodeId.gantry_y: MoveGroupSingleAxisStep(
                    distance_mm=float64(25),
                    velocity_mm_sec=float64(23),
                    duration_sec=float64(0),
                    acceleration_mm_sec_sq=float64(1000),
                ),
            }
        ],
        # Group 2
        [
            {
                NodeId.pipette_left: MoveGroupSingleAxisStep(
                    distance_mm=float64(12),
                    velocity_mm_sec=float64(-23),
                    duration_sec=float64(1234),
                    acceleration_mm_sec_sq=float64(1000),
                ),
            },
            {
                NodeId.pipette_left: MoveGroupSingleAxisStep(
                    distance_mm=float64(12),
                    velocity_mm_sec=float64(23),
                    duration_sec=float64(1234),
                    acceleration_mm_sec_sq=float64(1000),
                ),
            },
        ],
    ]

async def test_no_groups_do_nothing(mock_can_messenger: AsyncMock) -> None:
    """It should not send any commands if there are no moves."""
    subject = MoveGroupRunner(move_groups=[])
    position = await subject.run(mock_can_messenger)
    mock_can_messenger.send.assert_not_called()
    assert position == {}

async def test_single_group_clear(
    mock_can_messenger: AsyncMock, move_group_single: MoveGroups
) -> None:
    """It should send a clear group command before setup."""
    subject = MoveGroupRunner(move_groups=move_group_single)
    await subject._clear_groups(can_messenger=mock_can_messenger)
    mock_can_messenger.send.assert_has_calls(
        [call(node_id=NodeId.broadcast, message=md.ClearAllMoveGroupsRequest())],
    )

async def test_multi_group_clear(
    mock_can_messenger: AsyncMock, move_group_multiple: MoveGroups
) -> None:
    """It should send a clear group command before setup."""
    subject = MoveGroupRunner(move_groups=move_group_multiple)
    await subject.prep(can_messenger=mock_can_messenger)
    mock_can_messenger.send.assert_has_calls(
        [call(node_id=NodeId.broadcast, message=md.ClearAllMoveGroupsRequest())],
    )

async def test_home(
    mock_can_messenger: AsyncMock, move_group_home_single: MoveGroups
) -> None:
    """Test Home Request Functionality."""
    subject = MoveGroupRunner(move_groups=move_group_home_single)
    await subject.prep(can_messenger=mock_can_messenger)
    step = move_group_home_single[0][0].get(NodeId.head)
    assert isinstance(step, MoveGroupSingleAxisStep)
    mock_can_messenger.send.assert_any_call(
        node_id=NodeId.head,
        message=HomeRequest(
            payload=HomeRequestPayload(
                group_id=UInt8Field(0),
                seq_id=UInt8Field(0),
                velocity=Int32Field(calc_velocity(step)),
                duration=UInt32Field(calc_duration(step)),
            )
        ),
    )

async def test_single_send_setup_commands(
    mock_can_messenger: AsyncMock, move_group_single: MoveGroups
) -> None:
    """It should send all the move group set up commands."""
    subject = MoveGroupRunner(move_groups=move_group_single)
    await subject.prep(can_messenger=mock_can_messenger)
    step = move_group_single[0][0].get(NodeId.head)
    assert isinstance(step, MoveGroupSingleAxisStep)
    mock_can_messenger.send.assert_any_call(
        node_id=NodeId.head,
        message=AddLinearMoveRequest(
            payload=AddLinearMoveRequestPayload(
                group_id=UInt8Field(0),
                seq_id=UInt8Field(0),
                request_stop_condition=UInt8Field(0),
                velocity=Int32Field(calc_velocity(step)),
                acceleration=Int32Field(calc_acceleration(step)),
                duration=UInt32Field(calc_duration(step)),
            )
        ),
    )

async def test_multi_send_setup_commands(
    mock_can_messenger: AsyncMock, move_group_multiple: MoveGroups
) -> None:
    """It should send all the move group set up commands."""
    subject = MoveGroupRunner(move_groups=move_group_multiple)
    await subject.prep(can_messenger=mock_can_messenger)

    # Group 0
    step = move_group_multiple[0][0].get(NodeId.head)
    assert isinstance(step, MoveGroupSingleAxisStep)
    mock_can_messenger.send.assert_any_call(
        node_id=NodeId.head,
        message=AddLinearMoveRequest(
            payload=AddLinearMoveRequestPayload(
                group_id=UInt8Field(0),
                seq_id=UInt8Field(0),
                request_stop_condition=UInt8Field(0),
                velocity=Int32Field(calc_velocity(step)),
                acceleration=Int32Field(calc_acceleration(step)),
                duration=UInt32Field(calc_duration(step)),
            )
        ),
    )

    # Group 1
    step = move_group_multiple[1][0].get(NodeId.gantry_x)
    assert isinstance(step, MoveGroupSingleAxisStep)
    mock_can_messenger.send.assert_any_call(
        node_id=NodeId.gantry_x,
        message=AddLinearMoveRequest(
            payload=AddLinearMoveRequestPayload(
                group_id=UInt8Field(1),
                seq_id=UInt8Field(0),
                request_stop_condition=UInt8Field(0),
                velocity=Int32Field(calc_velocity(step)),
                acceleration=Int32Field(calc_acceleration(step)),
                duration=UInt32Field(calc_duration(step)),
            )
        ),
    )

    step = move_group_multiple[1][0].get(NodeId.gantry_y)
    assert isinstance(step, MoveGroupSingleAxisStep)
    mock_can_messenger.send.assert_any_call(
        node_id=NodeId.gantry_y,
        message=AddLinearMoveRequest(
            payload=AddLinearMoveRequestPayload(
                group_id=UInt8Field(1),
                seq_id=UInt8Field(0),
                request_stop_condition=UInt8Field(0),
                velocity=Int32Field(calc_velocity(step)),
                acceleration=Int32Field(calc_acceleration(step)),
                duration=UInt32Field(calc_duration(step)),
            )
        ),
    )

    # Group 2
    step = move_group_multiple[2][0].get(NodeId.pipette_left)
    assert isinstance(step, MoveGroupSingleAxisStep)
    mock_can_messenger.send.assert_any_call(
        node_id=NodeId.pipette_left,
        message=AddLinearMoveRequest(
            payload=AddLinearMoveRequestPayload(
                group_id=UInt8Field(2),
                seq_id=UInt8Field(0),
                request_stop_condition=UInt8Field(0),
                velocity=Int32Field(calc_velocity(step)),
                acceleration=Int32Field(calc_acceleration(step)),
                duration=UInt32Field(calc_duration(step)),
            )
        ),
    )

    step = move_group_multiple[2][1].get(NodeId.pipette_left)
    assert isinstance(step, MoveGroupSingleAxisStep)
    mock_can_messenger.send.assert_any_call(
        node_id=NodeId.pipette_left,
        message=AddLinearMoveRequest(
            payload=AddLinearMoveRequestPayload(
                group_id=UInt8Field(2),
                seq_id=UInt8Field(1),
                request_stop_condition=UInt8Field(0),
                velocity=Int32Field(calc_velocity(step)),
                acceleration=Int32Field(calc_acceleration(step)),
                duration=UInt32Field(calc_duration(step)),
            )
        ),
    )

async def test_move() -> None:
    """It should register to listen for messages."""
    subject = MoveGroupRunner(move_groups=[])
    mock_can_messenger = MagicMock()
    await subject._move(mock_can_messenger)
    mock_can_messenger.add_listener.assert_called_once()
    mock_can_messenger.remove_listener.assert_called_once()

class MockSendMoveCompleter:
    """Side effect mock of CanMessenger.send that immediately completes moves."""

    def __init__(
        self, move_groups: MoveGroups, listener: MessageListenerCallback
    ) -> None:
        """Constructor."""
        self._move_groups = move_groups
        self._listener = listener

    async def mock_send(
        self,
        node_id: NodeId,
        message: MessageDefinition,
    ) -> None:
        """Mock send function."""
        if isinstance(message, md.ExecuteMoveGroupRequest):
            # Iterate through each move in each sequence and send a move
            # completed for it.
            for seq_id, moves in enumerate(
                self._move_groups[message.payload.group_id.value]
            ):
                for node, move in moves.items():
                    assert isinstance(move, MoveGroupSingleAxisStep)
                    payload = MoveCompletedPayload(
                        group_id=message.payload.group_id,
                        seq_id=UInt8Field(seq_id),
                        current_position_um=UInt32Field(int(move.distance_mm * 1000)),
                        encoder_position=UInt32Field(int(move.distance_mm * 4000)),
                        ack_id=UInt8Field(1),
                    )
                    arbitration_id = ArbitrationId(
                        parts=ArbitrationIdParts(originating_node_id=node)
                    )
                    self._listener(md.MoveCompleted(payload=payload), arbitration_id)

async def test_single_move(
    mock_can_messenger: AsyncMock, move_group_single: MoveGroups
) -> None:
    """It should send a start group command."""
    subject = MoveScheduler(move_groups=move_group_single)
    mock_sender = MockSendMoveCompleter(move_group_single, subject)
    mock_can_messenger.send.side_effect = mock_sender.mock_send
    position = await subject.run(can_messenger=mock_can_messenger)

    mock_can_messenger.send.assert_has_calls(
        calls=[
            call(
                node_id=NodeId.broadcast,
                message=md.ExecuteMoveGroupRequest(
                    payload=ExecuteMoveGroupRequestPayload(
                        group_id=UInt8Field(0),
                        cancel_trigger=UInt8Field(0),
                        start_trigger=UInt8Field(0),
                    )
                ),
            )
        ]
    )
    assert len(position) == 1
    assert position[0][1].payload.current_position_um.value == 246000

async def test_multi_group_move(
    mock_can_messenger: AsyncMock, move_group_multiple: MoveGroups
) -> None:
    """It should start next group once the prior has completed."""
    subject = MoveScheduler(move_groups=move_group_multiple)
    mock_sender = MockSendMoveCompleter(move_group_multiple, subject)
    mock_can_messenger.send.side_effect = mock_sender.mock_send
    position = await subject.run(can_messenger=mock_can_messenger)

    mock_can_messenger.send.assert_has_calls(
        calls=[
            call(
                node_id=NodeId.broadcast,
                message=md.ExecuteMoveGroupRequest(
                    payload=ExecuteMoveGroupRequestPayload(
                        group_id=UInt8Field(0),
                        cancel_trigger=UInt8Field(0),
                        start_trigger=UInt8Field(0),
                    )
                ),
            ),
            call(
                node_id=NodeId.broadcast,
                message=md.ExecuteMoveGroupRequest(
                    payload=ExecuteMoveGroupRequestPayload(
                        group_id=UInt8Field(1),
                        cancel_trigger=UInt8Field(0),
                        start_trigger=UInt8Field(0),
                    )
                ),
            ),
            call(
                node_id=NodeId.broadcast,
                message=md.ExecuteMoveGroupRequest(
                    payload=ExecuteMoveGroupRequestPayload(
                        group_id=UInt8Field(2),
                        cancel_trigger=UInt8Field(0),
                        start_trigger=UInt8Field(0),
                    )
                ),
            ),
        ]
    )
    assert len(position) == 5
    assert position[0][1].payload.current_position_um.value == 229000
    assert position[1][1].payload.current_position_um.value == 522000
    assert position[2][1].payload.current_position_um.value == 25000
    assert position[3][1].payload.current_position_um.value == 12000
    assert position[4][1].payload.current_position_um.value == 12000

def _build_arb(from_node: NodeId) -> ArbitrationId:
    return ArbitrationId(ArbitrationIdParts(originating_node_id=from_node))

@pytest.mark.parametrize(
    "completions,position_map",
    [
        (
            # one axis, completions reversed compared to execution order
            [
                (
                    _build_arb(NodeId.gantry_x),
                    MoveCompleted(
                        payload=MoveCompletedPayload(
                            ack_id=UInt8Field(0),
                            group_id=UInt8Field(2),
                            seq_id=UInt8Field(2),
                            current_position_um=UInt32Field(10000),
                            encoder_position=UInt32Field(10000 * 4),
                        )
                    ),
                ),
                (
                    _build_arb(NodeId.gantry_x),
                    MoveCompleted(
                        payload=MoveCompletedPayload(
                            ack_id=UInt8Field(0),
                            group_id=UInt8Field(2),
                            seq_id=UInt8Field(1),
                            current_position_um=UInt32Field(20000),
                            encoder_position=UInt32Field(10000 * 4),
                        )
                    ),
                ),
                (
                    _build_arb(NodeId.gantry_x),
                    MoveCompleted(
                        payload=MoveCompletedPayload(
                            ack_id=UInt8Field(0),
                            group_id=UInt8Field(1),
                            seq_id=UInt8Field(2),
                            current_position_um=UInt32Field(30000),
                            encoder_position=UInt32Field(10000 * 4),
                        )
                    ),
                ),
            ],
            {NodeId.gantry_x: 10},
        ),
        (
            # multiple axes with different numbers of completions
            [
                (
                    _build_arb(NodeId.gantry_x),
                    MoveCompleted(
                        payload=MoveCompletedPayload(
                            ack_id=UInt8Field(0),
                            group_id=UInt8Field(2),
                            seq_id=UInt8Field(2),
                            current_position_um=UInt32Field(10000),
                            encoder_position=UInt32Field(10000 * 4),
                        )
                    ),
                ),
                (
                    _build_arb(NodeId.gantry_x),
                    MoveCompleted(
                        payload=MoveCompletedPayload(
                            ack_id=UInt8Field(0),
                            group_id=UInt8Field(2),
                            seq_id=UInt8Field(1),
                            current_position_um=UInt32Field(20000),
                            encoder_position=UInt32Field(10000 * 4),
                        )
                    ),
                ),
                (
                    _build_arb(NodeId.gantry_y),
                    MoveCompleted(
                        payload=MoveCompletedPayload(
                            ack_id=UInt8Field(0),
                            group_id=UInt8Field(1),
                            seq_id=UInt8Field(2),
                            current_position_um=UInt32Field(30000),
                            encoder_position=UInt32Field(10000 * 4),
                        )
                    ),
                ),
            ],
            {NodeId.gantry_x: 10, NodeId.gantry_y: 30},
        ),
        (
            # empty base case
            [],
            {},
        ),
    ],
)
def test_accumulate_move_completions(
    completions: List[_CompletionPacket], position_map: NodeMap[float]
) -> None:
    """Build correct move results."""
    assert MoveGroupRunner._accumulate_move_completions(completions) == position_map

@pytest.mark.parametrize("empty_group", [[], [[]], [[{}]]])
async def test_empty_groups(
    mock_can_messenger: AsyncMock, empty_group: List[Any]
) -> None:
    """Test that various kinds of empty groups result in no calls."""
    mg = MoveGroupRunner(empty_group)
    await mg.run(mock_can_messenger)
    mock_can_messenger.send.assert_not_called()


# file: __init__.py
#===============================================================================
# Copyright 2014-2018 Intel Corporation.
#
# This software and the related documents are Intel copyrighted  materials,  and
# your use of  them is  governed by the  express license  under which  they were
# provided to you (License).  Unless the License provides otherwise, you may not
# use, modify, copy, publish, distribute,  disclose or transmit this software or
# the related documents without Intel's prior written permission.
#
# This software and the related documents  are provided as  is,  with no express
# or implied  warranties,  other  than those  that are  expressly stated  in the
# License.
#===============================================================================

##

from __future__ import print_function

import sys
import numpy as np

from daal.data_management import (
    CSRNumericTable, NumericTableIface, readOnly, BlockDescriptor, packed_mask,
    DataSourceIface, FileDataSource, HomogenTensor, SubtensorDescriptor
)

CRC_REM = [
    0x00000000, 0x741B8CD6, 0xE83719AC, 0x9C2C957A, 0xA475BF8E, 0xD06E3358, 0x4C42A622, 0x38592AF4,
    0x3CF0F3CA, 0x48EB7F1C, 0xD4C7EA66, 0xA0DC66B0, 0x98854C44, 0xEC9EC092, 0x70B255E8, 0x04A9D93E,
    0x79E1E794, 0x0DFA6B42, 0x91D6FE38, 0xE5CD72EE, 0xDD94581A, 0xA98FD4CC, 0x35A341B6, 0x41B8CD60,
    0x4511145E, 0x310A9888, 0xAD260DF2, 0xD93D8124, 0xE164ABD0, 0x957F2706, 0x0953B27C, 0x7D483EAA,
    0xF3C3CF28, 0x87D843FE, 0x1BF4D684, 0x6FEF5A52, 0x57B670A6, 0x23ADFC70, 0xBF81690A, 0xCB9AE5DC,
    0xCF333CE2, 0xBB28B034, 0x2704254E, 0x531FA998, 0x6B46836C, 0x1F5D0FBA, 0x83719AC0, 0xF76A1616,
    0x8A2228BC, 0xFE39A46A, 0x62153110, 0x160EBDC6, 0x2E579732, 0x5A4C1BE4, 0xC6608E9E, 0xB27B0248,
    0xB6D2DB76, 0xC2C957A0, 0x5EE5C2DA, 0x2AFE4E0C, 0x12A764F8, 0x66BCE82E, 0xFA907D54, 0x8E8BF182,
    0x939C1286, 0xE7879E50, 0x7BAB0B2A, 0x0FB087FC, 0x37E9AD08, 0x43F221DE, 0xDFDEB4A4, 0xABC53872,
    0xAF6CE14C, 0xDB776D9A, 0x475BF8E0, 0x33407436, 0x0B195EC2, 0x7F02D214, 0xE32E476E, 0x9735CBB8,
    0xEA7DF512, 0x9E6679C4, 0x024AECBE, 0x76516068, 0x4E084A9C, 0x3A13C64A, 0xA63F5330, 0xD224DFE6,
    0xD68D06D8, 0xA2968A0E, 0x3EBA1F74, 0x4AA193A2, 0x72F8B956, 0x06E33580, 0x9ACFA0FA, 0xEED42C2C,
    0x605FDDAE, 0x14445178, 0x8868C402, 0xFC7348D4, 0xC42A6220, 0xB031EEF6, 0x2C1D7B8C, 0x5806F75A,
    0x5CAF2E64, 0x28B4A2B2, 0xB49837C8, 0xC083BB1E, 0xF8DA91EA, 0x8CC11D3C, 0x10ED8846, 0x64F60490,
    0x19BE3A3A, 0x6DA5B6EC, 0xF1892396, 0x8592AF40, 0xBDCB85B4, 0xC9D00962, 0x55FC9C18, 0x21E710CE,
    0x254EC9F0, 0x51554526, 0xCD79D05C, 0xB9625C8A, 0x813B767E, 0xF520FAA8, 0x690C6FD2, 0x1D17E304,
    0x5323A9DA, 0x2738250C, 0xBB14B076, 0xCF0F3CA0, 0xF7561654, 0x834D9A82, 0x1F610FF8, 0x6B7A832E,
    0x6FD35A10, 0x1BC8D6C6, 0x87E443BC, 0xF3FFCF6A, 0xCBA6E59E, 0xBFBD6948, 0x2391FC32, 0x578A70E4,
    0x2AC24E4E, 0x5ED9C298, 0xC2F557E2, 0xB6EEDB34, 0x8EB7F1C0, 0xFAAC7D16, 0x6680E86C, 0x129B64BA,
    0x1632BD84, 0x62293152, 0xFE05A428, 0x8A1E28FE, 0xB247020A, 0xC65C8EDC, 0x5A701BA6, 0x2E6B9770,
    0xA0E066F2, 0xD4FBEA24, 0x48D77F5E, 0x3CCCF388, 0x0495D97C, 0x708E55AA, 0xECA2C0D0, 0x98B94C06,
    0x9C109538, 0xE80B19EE, 0x74278C94, 0x003C0042, 0x38652AB6, 0x4C7EA660, 0xD052331A, 0xA449BFCC,
    0xD9018166, 0xAD1A0DB0, 0x313698CA, 0x452D141C, 0x7D743EE8, 0x096FB23E, 0x95432744, 0xE158AB92,
    0xE5F172AC, 0x91EAFE7A, 0x0DC66B00, 0x79DDE7D6, 0x4184CD22, 0x359F41F4, 0xA9B3D48E, 0xDDA85858,
    0xC0BFBB5C, 0xB4A4378A, 0x2888A2F0, 0x5C932E26, 0x64CA04D2, 0x10D18804, 0x8CFD1D7E, 0xF8E691A8,
    0xFC4F4896, 0x8854C440, 0x1478513A, 0x6063DDEC, 0x583AF718, 0x2C217BCE, 0xB00DEEB4, 0xC4166262,
    0xB95E5CC8, 0xCD45D01E, 0x51694564, 0x2572C9B2, 0x1D2BE346, 0x69306F90, 0xF51CFAEA, 0x8107763C,
    0x85AEAF02, 0xF1B523D4, 0x6D99B6AE, 0x19823A78, 0x21DB108C, 0x55C09C5A, 0xC9EC0920, 0xBDF785F6,
    0x337C7474, 0x4767F8A2, 0xDB4B6DD8, 0xAF50E10E, 0x9709CBFA, 0xE312472C, 0x7F3ED256, 0x0B255E80,
    0x0F8C87BE, 0x7B970B68, 0xE7BB9E12, 0x93A012C4, 0xABF93830, 0xDFE2B4E6, 0x43CE219C, 0x37D5AD4A,
    0x4A9D93E0, 0x3E861F36, 0xA2AA8A4C, 0xD6B1069A, 0xEEE82C6E, 0x9AF3A0B8, 0x06DF35C2, 0x72C4B914,
    0x766D602A, 0x0276ECFC, 0x9E5A7986, 0xEA41F550, 0xD218DFA4, 0xA6035372, 0x3A2FC608, 0x4E344ADE
]

def readTextFile(dataset_filename):
    data_list = []
    with open(dataset_filename, 'r') as f:
        lines = f.readlines()
        for line in lines:
            data_list.append(tuple([float(x) for x in line.split(',')[:-1]]))
    return np.array(data_list, dtype=np.ubyte).flatten()

def createSparseTable(dataset_filename, ntype=np.float64):

    try:
        f = open(dataset_filename, 'r')
    except Exception:
        sys.exit("Unable to open '{}'.".format(dataset_filename))

    lines = f.readlines()
    row_index_line = lines[0]
    columns_line = lines[1]
    values_line = lines[2]

    row_offsets = np.array(row_index_line.rstrip('\n,').split(','), dtype=np.uint64)
    num_vectors = len(row_offsets) - 1
    col_indices = np.array(columns_line.rstrip('\n,').split(','), dtype=np.uint64)
    num_columns = len(col_indices)
    data = np.array(values_line.rstrip('\n,').split(','), dtype=ntype)
    num_nonzeros = len(data)

    max_col = 0

    for i in range(num_columns):
        if col_indices[i] > max_col:
            max_col = col_indices[i]
    if sys.version_info[0] < 3:
        num_features = long(max_col)
    else:
        num_features = int(max_col)

    if (
        num_columns != num_nonzeros or
        num_nonzeros != int(row_offsets[num_vectors]) - 1 or
        num_features == 0 or
        num_vectors == 0
    ):
        print("Incorrect format of file")
        sys.exit()

    numeric_table = CSRNumericTable(data, col_indices, row_offsets, num_features, num_vectors)

    return numeric_table

def printAprioriItemsets(large_itemsets_table, large_itemsets_support_table,
                         itemsets_to_print=20):
    large_itemset_count = large_itemsets_support_table.getNumberOfRows()
    num_items_in_large_itemsets = large_itemsets_table.getNumberOfRows()

    large_itemsets = large_itemsets_table.getBlockOfRowsAsDouble(0, num_items_in_large_itemsets).flatten()
    large_itemsets_support_data = large_itemsets_support_table.getBlockOfRowsAsDouble(0, large_itemset_count).flatten()

    large_itemsets_array = [[] for x in range(large_itemset_count)]

    for i in range(num_items_in_large_itemsets):
        large_itemsets_array[int(large_itemsets[2 * i])].append(large_itemsets[2 * i + 1])

    support_array = [0] * large_itemset_count

    for i in range(large_itemset_count):
        support_array[int(large_itemsets_support_data[2 * i])] = large_itemsets_support_data[2 * i + 1]

    print("\nApriori example program results\n")
    print("Last " + str(itemsets_to_print) + " large itemsets: \n")
    print("Itemset\t\t\tSupport")

    if large_itemset_count > itemsets_to_print and itemsets_to_print != 0:
        i_min = large_itemset_count - itemsets_to_print
    else:
        i_min = 0

    for i in range(i_min, large_itemset_count):
        print("{", end='')
        for l in range(len(large_itemsets_array[i]) - 1):
            print("{:.0f}".format(large_itemsets_array[i][l]), end=', ')
        print("{:.0f}}}\t\t".format(large_itemsets_array[i][len(large_itemsets_array[i]) - 1]), end='')

        print("{:.0f}".format(support_array[i]))

def printAprioriRules(left_items_table, right_items_table, confidence_table,
                      num_rules_to_print=20):

    num_rules = confidence_table.getNumberOfRows()
    num_left_items = left_items_table.getNumberOfRows()
    num_right_items = right_items_table.getNumberOfRows()

    left_items = left_items_table.getBlockOfRowsAsDouble(0, num_left_items).flatten()
    right_items = right_items_table.getBlockOfRowsAsDouble(0, num_right_items).flatten()
    confidence = confidence_table.getBlockOfRowsAsDouble(0, num_rules).flatten()

    left_items_array = [[] for x in range(num_rules)]

    if num_rules == 0:
        print("\nNo association rules were found ")
        return

    for i in range(num_left_items):
        left_items_array[int(left_items[2 * i])].append(left_items[2 * i + 1])

    right_items_array = [[] for x in range(num_right_items)]

    for i in range(num_right_items):
        right_items_array[int(right_items[2 * i])].append(right_items[2 * i + 1])

    confidence_array = [0] * num_rules

    for i in range(num_rules):
        confidence_array[i] = confidence[i]

    print("\nLast {} association rules: \n".format(num_rules_to_print))
    print("Rule\t\t\t\tConfidence")
    if num_rules > num_rules_to_print and num_rules_to_print != 0:
        i_min = num_rules - num_rules_to_print
    else:
        i_min = 0

    for i in range(i_min, num_rules):
        print("{", end='')
        for l in range(len(left_items_array[i]) - 1):
            print("{:.0f}, ".format(left_items_array[i][l]), end='')
        print("{:.0f}}} => {{".format(left_items_array[i][len(left_items_array[i]) - 1]), end='')

        for l in range(len(right_items_array[i]) - 1):
            print("{:.0f}, ".format(right_items_array[i][l]), end='')
        print("{:.0f}}}\t\t".format(right_items_array[i][len(right_items_array[i]) - 1]), end='')

        print("{0:.6g}".format(confidence_array[i]))

def isFull(layout):
    layout_int = int(layout)
    if packed_mask & layout_int:
        return False
    return True

def printArray(array, num_printed_cols, num_printed_rows, num_cols, message,
               interval=10, flt64=True):
    print(message)
    flat_array = array.flatten()
    decimals = '3' if flt64 else '0'
    for i in range(num_printed_rows):
        for j in range(num_printed_cols):
            print("{:<{width}.{dec}f}".format(
                flat_array[i * num_cols + j], width=interval, dec=decimals), end=''
            )
        print()
    print()

def isUpper(layout):
    if (
        layout == NumericTableIface.upperPackedSymmetricMatrix or
        layout == NumericTableIface.upperPackedTriangularMatrix
    ):
        return True
    return False

def isLower(layout):
    if (
        layout == NumericTableIface.lowerPackedSymmetricMatrix or
        layout == NumericTableIface.lowerPackedTriangularMatrix
    ):
        return True
    return False

def printLowerArray(array, num_printed_rows, message, interval=10):
    cols = 1
    print(message)
    for i in range(num_printed_rows):
        for j in range(cols):
            print("{:<{width}.3f}".format(array[i][j], width=interval), end='')
        print()
        cols += 1
    print()

def printUpperArray(array, num_printed_cols, num_printed_rows, num_cols,
                    message, interval=10):
    print(message)
    for i in range(num_printed_rows):
        for j in range(i):
            print(' ' * 10, end='')
        for j in range(i, num_printed_cols):
            print("{:<{width}.3f}".format(array[i][j], width=interval), end='')
        print()
    print()

def printNumericTable(data_table, message='', num_printed_rows=0, num_printed_cols=0,
                      interval=10):
    num_rows = data_table.getNumberOfRows()
    num_cols = data_table.getNumberOfColumns()
    layout = data_table.getDataLayout()

    if num_printed_rows != 0:
        num_printed_rows = min(num_rows, num_printed_rows)
    else:
        num_printed_rows = num_rows

    if num_printed_cols != 0:
        num_printed_cols = min(num_cols, num_printed_cols)
    else:
        num_printed_cols = num_cols

    block = BlockDescriptor()
    if isFull(layout) or layout == NumericTableIface.csrArray:
        data_table.getBlockOfRows(0, num_rows, readOnly, block)
        printArray(block.getArray(), num_printed_cols, num_printed_rows,
                   num_cols, message, interval)
        data_table.releaseBlockOfRows(block)
    else:
        packed_table = data_table.getBlockOfRowsAsDouble(0, num_rows)

        if isLower(layout):
            printLowerArray(packed_table, num_printed_rows, message, interval)
        elif isUpper(layout):
            printUpperArray(packed_table, num_printed_cols, num_printed_rows,
                            num_cols, message, interval)

def printNumericTables(data_table_1, data_table_2, title_1='', title_2='',
                       message='', num_printed_rows=0, interval=15, flt64=True):
    num_rows_1 = data_table_1.getNumberOfRows()
    num_rows_2 = data_table_2.getNumberOfRows()
    num_cols_1 = data_table_1.getNumberOfColumns()
    num_cols_2 = data_table_2.getNumberOfColumns()

    num_rows = min(num_rows_1, num_rows_2)
    if num_printed_rows != 0:
        num_rows = min(min(num_rows_1, num_rows_2), num_printed_rows)

    block1 = BlockDescriptor()
    block2 = BlockDescriptor()
    data_table_1.getBlockOfRows(0, num_rows, readOnly, block1)
    data_table_2.getBlockOfRows(0, num_rows, readOnly, block2)

    data_float64_1 = block1.getArray()
    data_float64_2 = block2.getArray()

    data_float64_1 = data_float64_1.flatten()
    data_float64_2 = data_float64_2.flatten()

    decimals = '3' if flt64 else '0'

    print(message)
    print("{:<{width}}".format(title_1, width=(interval * num_cols_1)), end='')
    print("{:<{width}}".format(title_2, width=(interval * num_cols_2)))
    for i in range(num_rows):
        for j in range(num_cols_1):
            print("{:<{width}.{dec}f}".format(data_float64_1[i * num_cols_1 + j], width=interval, dec=decimals), end='')
        for j in range(num_cols_2):
            print("{:<{width}.0f}".format(data_float64_2[i * num_cols_2 + j], width=interval), end='')
        print()
    print()

    data_table_1.releaseBlockOfRows(block1)
    data_table_2.releaseBlockOfRows(block2)

def getCRC32(input, prevRes=0):
    from binascii import crc32
    return crc32(input, prevRes)

def copyBytes(dst, src, size):
    for i in range(size):
        dst[i] = src[i]

def printALSRatings(usersOffsetTable, itemsOffsetTable, ratings):
    nUsers = ratings.getNumberOfRows()
    nItems = ratings.getNumberOfColumns()
    ratingsData = ratings.getBlockOfRowsAsDouble(0, nUsers).flatten()
    usersOffset = usersOffsetTable.getBlockOfRowsAsInt(0, 1).flatten()[0]
    itemsOffset = itemsOffsetTable.getBlockOfRowsAsInt(0, 1).flatten()[0]

    print(" User ID, Item ID, rating")
    for i in range(nUsers):
        for j in range(nItems):
            print("{}, {}, {:.6g}".format(i + usersOffset, j + itemsOffset, ratingsData[i * nItems + j]))

def printTensor(dataTable, message="", nPrintedRows=0, nPrintedCols=0, interval=10):
    dims = dataTable.getDimensions()
    nRows = int(dims[0])

    if nPrintedRows != 0:
        nPrintedRows = min(nRows, nPrintedRows)
    else:
        nPrintedRows = nRows

    block = SubtensorDescriptor()

    dataTable.getSubtensor([], 0, nPrintedRows, readOnly, block)

    nCols = int(block.getSize() / nPrintedRows)

    if nPrintedCols != 0:
        nPrintedCols = min(nCols, nPrintedCols)
    else:
        nPrintedCols = nCols

    printArray(block.getArray(), int(nPrintedCols), int(nPrintedRows), int(nCols), message, interval)
    dataTable.releaseSubtensor(block)

def printTensors(dataTable1, dataTable2, title1="", title2="", message="", nPrintedRows=0, interval=15):
    dims1 = dataTable1.getDimensions()
    nRows1 = int(dims1[0])

    if nPrintedRows != 0:
        nPrintedRows = min(nRows1, nPrintedRows)
    else:
        nPrintedRows = nRows1

    block1 = SubtensorDescriptor()
    dataTable1.getSubtensor([], 0, nPrintedRows, readOnly, block1)
    nCols1 = int(block1.getSize() / nPrintedRows)

    dims2 = dataTable2.getDimensions()
    nRows2 = int(dims2[0])

    if nPrintedRows != 0:
        nPrintedRows = min(nRows2, nPrintedRows)
    else:
        nPrintedRows = nRows2

    block2 = SubtensorDescriptor()
    dataTable2.getSubtensor([], 0, nPrintedRows, readOnly, block2)
    nCols2 = int(block2.getSize() / nPrintedRows)

    dataType1 = block1.getArray().flatten()
    dataType2 = block2.getArray().flatten()

    print(message)
    print("{:<{width}}".format(title1, width=(interval * nCols1)), end='')
    print("{:<{width}}".format(title2, width=(interval * nCols2)))

    for i in range(nPrintedRows):
        for j in range(nCols1):
            print("{v:<{width}.0f}".format(v=dataType1[i * nCols1 + j], width=interval), end='')

        for j in range(nCols2):
            print("{:<{width}.3f}".format(dataType2[i * nCols2 + j], width=int(interval / 2)), end='')
        print()
    print()

    dataTable1.releaseSubtensor(block1)
    dataTable2.releaseSubtensor(block2)

def printTensor3d(dataTable, message="", nFirstDim=0, nSecondDim=0, interval=10):
    dims = dataTable.getDimensions()
    nRows = int(dims[0])
    nCols = int(dims[1])

    if nFirstDim != 0:
        nFirstDim = min(nRows, nFirstDim)
    else:
        nFirstDim = nRows

    if nSecondDim != 0:
        nSecondDim = min(nCols, nSecondDim)
    else:
        nSecondDim = nCols

    block = SubtensorDescriptor()

    print(message)
    for i in range(nFirstDim):
        dataTable.getSubtensor([i], 0, nSecondDim, readOnly, block)

        nThirdDim = block.getSize() / nSecondDim

        printArray(block.getArray(), int(nThirdDim), int(nSecondDim), int(nThirdDim), "", interval)

        dataTable.releaseSubtensor(block)

def readTensorFromCSV(datasetFileName, allowOneColumn=False):
    dataSource = FileDataSource(datasetFileName,
                                DataSourceIface.doAllocateNumericTable,
                                DataSourceIface.doDictionaryFromContext)
    dataSource.loadDataBlock()

    nt = dataSource.getNumericTable()
    size = nt.getNumberOfRows()
    block = BlockDescriptor()
    nt.getBlockOfRows(0, size, readOnly, block)
    blockData = block.getArray().flatten()

    dims = [size]
    if nt.getNumberOfColumns() > 1 or allowOneColumn:
        dims.append(nt.getNumberOfColumns())
        size *= dims[1]

    tensorData = np.array(blockData, dtype=np.float32)

    nt.releaseBlockOfRows(block)

    tensorData.shape = dims
    tensor = HomogenTensor(tensorData, ntype=np.float32)

    return tensor


import cv2
import datetime
import imutils
import numpy as np

protopath = "MobileNetSSD_deploy.prototxt"
modelpath = "MobileNetSSD_deploy.caffemodel"
detector = cv2.dnn.readNetFromCaffe(prototxt=protopath, caffeModel=modelpath)

CLASSES = ["background", "aeroplane", "bicycle", "bird", "boat",
           "bottle", "bus", "car", "cat", "chair", "cow", "diningtable",
           "dog", "horse", "motorbike", "person", "pottedplant", "sheep",
           "sofa", "train", "tvmonitor"]

def main():
    cap = cv2.VideoCapture('test_video.mp4')

    fps_start_time = datetime.datetime.now()
    fps = 0

    total_frames = 0

    while True:
        ret, frame = cap.read()
        frame = imutils.resize(frame, width=600)
        total_frames = total_frames + 1

        (H, W) = frame.shape[:2]

        blob = cv2.dnn.blobFromImage(frame, 0.007843, (W, H), 127.5)

        detector.setInput(blob)
        person_detections = detector.forward()

        for i in np.arange(0, person_detections.shape[2]):
            confidence = person_detections[0, 0, i, 2]
            if confidence > 0.5:
                idx = int(person_detections[0, 0, i, 1])

                if CLASSES[idx] != "person":
                    continue

                person_box = person_detections[0, 0, i, 3:7] * np.array([W, H, W, H])
                (startX, startY, endX, endY) = person_box.astype("int")

                cv2.rectangle(frame, (startX, startY), (endX, endY), (0, 0, 255), 2)

        fps_end_time = datetime.datetime.now()
        time_diff = fps_end_time - fps_start_time
        if time_diff.seconds == 0:
            fps = 0.0
        else:
            fps = (total_frames / time_diff.seconds)

        fps_text = "FPS: {:.2f}".format(fps)

        cv2.putText(frame, fps_text, (5, 30), cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 0, 255), 1)

        cv2.imshow("Application", frame)
        key = cv2.waitKey(1)
        if key == ord('q'):
            break

    cv2.destroyAllWindows()

main()


import os
from collections import Counter

import numpy as np
from keras.engine.saving import load_model

from utils.utils import from_env, get_root
from webserver.storage.factory import StorageFactory

model = None
MODEL_TYPE = from_env('MODEL_TYPE', 'uk_target_only_native_speakers')
MODEL_NUM = from_env('MODEL_NUM', "64812b64080b4668ac824c9ca75b6c04")

def predict_class_audio(MFCCs):
    '''
    Predict class based on MFCC samples
    :param MFCCs: Numpy array of MFCCs
    :param model: Trained model
    :return: Predicted class of MFCC segment group
    '''
    global model
    MFCCs = MFCCs.reshape(MFCCs.shape[0], MFCCs.shape[1], MFCCs.shape[2], 1)
    y_predicted = model.predict_classes(MFCCs, verbose=0)
    return Counter(list(y_predicted)).most_common(1)[0][0]

def load(from_cloud=True):
    # The current served model based on the experiment type

    global model

    storage = StorageFactory.default()
    file_path = storage.load_model(MODEL_TYPE, MODEL_NUM)
    model = load_model(file_path)

    # BUG fix - initializing the model with an empty vector
    model.predict(np.zeros((1, 13, 30, 1)))

load()


import win32api
import win32con
import win32clipboard
from random import randint, random
import sys

import pyautogui
from .directkeys import *
import numpy as np

from .utils import get_config

class InputHandler:
    def __init__(self, resolution):
        self.res_w = resolution[0]
        self.res_h = resolution[1]
        self.config = get_config("input_handler")

    def click(
        self, x_top, y_top, x_bot, y_bot, button="left", speed_factor=3, raw=False
    ):
        if not raw:
            x_top *= self.res_w
            x_bot *= self.res_w
            y_top *= self.res_h
            y_bot *= self.res_h

        x = x_top + random() * (x_bot - x_top)
        y = y_top + random() * (x_bot - x_top)
        pyautogui.moveTo(
            x,
            y,
            self._rnd(
                min=100 / speed_factor,
                mean=150 / speed_factor,
                sigma=100 / speed_factor,
            )
            / 1000,
            pyautogui.easeOutQuad,
        )
        self._hw_move(x, y)
        if button is not None:
            self.rnd_sleep(mean=self.config["mouse_click_interval"] / speed_factor)
            pyautogui.click(x=None, y=None, button=button)

    def click_keys(self, keys, speed_factor=1):
        self._press_keys(keys, speed_factor)
        self._release_keys(keys, speed_factor)

    def click_hotkey(self, hotkey):
        pyautogui.press(hotkey)
        self.rnd_sleep(mean=self.config["key_press_interval"])

    def _press_keys(self, keys, speed_factor=1):
        for key in keys:
            PressKey(key)
            self.rnd_sleep(mean=self.config["key_press_interval"] / speed_factor)

    def _release_keys(self, keys, speed_factor=1):
        for key in keys:
            ReleaseKey(key)
            self.rnd_sleep(mean=self.config["key_press_interval"] / speed_factor)

    def type(self, text):
        self.click_keys([0x1C])
        pyautogui.typewrite(text, interval=self.config["key_press_interval"] / 1000)
        self.click_keys([0x1C])

    def _hw_move(self, x, y):
        nx = int(x * 65535 / win32api.GetSystemMetrics(0))
        ny = int(y * 65535 / win32api.GetSystemMetrics(1))
        win32api.mouse_event(
            win32con.MOUSEEVENTF_ABSOLUTE | win32con.MOUSEEVENTF_MOVE, nx, ny
        )

    def rnd_sleep(self, min=50, mean=300, sigma=300):
        time.sleep(self._rnd(min, mean, sigma=sigma) / 1000)

    def _rnd(self, min, mean, sigma=200):
        r = np.random.normal(mean, sigma)
        while r < min:
            r = np.random.normal(mean, sigma)
        return r

    def inventory_click(
        self,
        slot_x,
        slot_y,
        inventory_base,
        ctrl_click=False,
        button="left",
        speed_factor=2,
    ):
        if ctrl_click:
            self._press_keys([0x1D])
        self.click(
            inventory_base[0] + slot_x * 0.02735 - 0.006625,
            inventory_base[1] + slot_y * 0.04930555 - 0.009777,
            inventory_base[0] + slot_x * 0.02735 + 0.006625,
            inventory_base[1] + slot_y * 0.04930555 + 0.009777,
            speed_factor=speed_factor,
            button=button,
        )
        if ctrl_click:
            self._release_keys([0x1D])

    def inventory_copy(self, slot_x, slot_y, inventory_base, speed_factor=1):
        self.click(
            inventory_base[0] + slot_x * 0.02735 - 0.006625,
            inventory_base[1] + slot_y * 0.04930555 - 0.009777,
            inventory_base[0] + slot_x * 0.02735 + 0.005625,
            inventory_base[1] + slot_y * 0.04930555 + 0.007777,
            speed_factor=speed_factor,
            button=None,
        )
        #reset the clipboard to make sure we actually have the right data
        win32clipboard.OpenClipboard()
        win32clipboard.EmptyClipboard()
        win32clipboard.CloseClipboard()
        self._copy(speed_factor)
        win32clipboard.OpenClipboard()
        data = win32clipboard.GetClipboardData()
        win32clipboard.CloseClipboard()
        return data

    def _copy(self, speed_factor=1):
        # Copying seems a tad unreliable, so we do it twice
        self.click_keys([0x1D, 0x2E], speed_factor)

    def zoom(self, clicks):
        # Did not work in testing
        dx = int(self.res_w / 2 * 65535 / win32api.GetSystemMetrics(0))
        dy = int(self.res_h / 2 * 65535 / win32api.GetSystemMetrics(1))
        self.click(0.499, 0.499, 0.501, 0.501, button=None)
        win32api.mouse_event(win32con.MOUSEEVENTF_WHEEL, dx, dy, clicks, 0)

    def drag(self, x, y, delta=False, speed_factor=1):
        if delta:
            pyautogui.drag(
                x,
                y,
                self._rnd(
                    min=500 / speed_factor,
                    mean=800 / speed_factor,
                    sigma=400 / speed_factor,
                )
                / 1000,
                pyautogui.easeOutQuad,
                button="left",
            )
        else:
            pyautogui.dragTo(
                x,
                y,
                self._rnd(
                    min=500 / speed_factor,
                    mean=800 / speed_factor,
                    sigma=400 / speed_factor,
                )
                / 1000,
                pyautogui.easeOutQuad,
                button="left",
            )


# Copyright (c) Facebook, Inc. and its affiliates.

# install `vqa-maskrcnn-benchmark` from
# https://github.com/ronghanghu/vqa-maskrcnn-benchmark-m4c
import argparse
import os
import sys

import cv2
import numpy as np
import torch
import tqdm
from maskrcnn_benchmark.config import cfg
from maskrcnn_benchmark.layers import nms
from maskrcnn_benchmark.modeling.detector import build_detection_model
from maskrcnn_benchmark.structures.image_list import to_image_list
from maskrcnn_benchmark.utils.model_serialization import load_state_dict
from PIL import Image

sys.path.append("/private/home/ronghanghu/workspace/vqa-maskrcnn-benchmark")  # NoQA

def load_detection_model(yaml_file, yaml_ckpt):
    cfg.merge_from_file(yaml_file)
    cfg.freeze()

    model = build_detection_model(cfg)
    checkpoint = torch.load(yaml_ckpt, map_location=torch.device("cpu"))

    load_state_dict(model, checkpoint.pop("model"))

    model.to("cuda")
    model.eval()
    return model

def _image_transform(image_path):
    img = Image.open(image_path)
    im = np.array(img).astype(np.float32)
    # handle a few corner cases
    if im.ndim == 2:  # gray => RGB
        im = np.tile(im[:, :, None], (1, 1, 3))
    if im.shape[2] > 3:  # RGBA => RGB
        im = im[:, :, :3]

    im = im[:, :, ::-1]  # RGB => BGR
    im -= np.array([102.9801, 115.9465, 122.7717])
    im_shape = im.shape
    im_size_min = np.min(im_shape[0:2])
    im_size_max = np.max(im_shape[0:2])
    im_scale = float(800) / float(im_size_min)
    # Prevent the biggest axis from being more than max_size
    if np.round(im_scale * im_size_max) > 1333:
        im_scale = float(1333) / float(im_size_max)
    im = cv2.resize(
        im, None, None, fx=im_scale, fy=im_scale, interpolation=cv2.INTER_LINEAR
    )
    img = torch.from_numpy(im).permute(2, 0, 1)
    return img, im_scale

def _process_feature_extraction(output, im_scales, feat_name="fc6"):
    batch_size = len(output[0]["proposals"])
    n_boxes_per_image = [len(_) for _ in output[0]["proposals"]]
    score_list = output[0]["scores"].split(n_boxes_per_image)
    score_list = [torch.nn.functional.softmax(x, -1) for x in score_list]
    feats = output[0][feat_name].split(n_boxes_per_image)
    cur_device = score_list[0].device

    feat_list = []
    bbox_list = []

    for i in range(batch_size):
        dets = output[0]["proposals"][i].bbox / im_scales[i]
        scores = score_list[i]

        max_conf = torch.zeros(scores.shape[0]).to(cur_device)

        for cls_ind in range(1, scores.shape[1]):
            cls_scores = scores[:, cls_ind]
            keep = nms(dets, cls_scores, 0.5)
            max_conf[keep] = torch.where(
                cls_scores[keep] > max_conf[keep], cls_scores[keep], max_conf[keep]
            )

        keep_boxes = torch.argsort(max_conf, descending=True)[:100]
        feat_list.append(feats[i][keep_boxes])
        bbox_list.append(output[0]["proposals"][i].bbox[keep_boxes])
    return feat_list, bbox_list

def extract_features(detection_model, image_path, input_boxes=None, feat_name="fc6"):
    im, im_scale = _image_transform(image_path)
    if input_boxes is not None:
        if isinstance(input_boxes, np.ndarray):
            input_boxes = torch.from_numpy(input_boxes.copy())
        input_boxes *= im_scale
    img_tensor, im_scales = [im], [im_scale]
    current_img_list = to_image_list(img_tensor, size_divisible=32)
    current_img_list = current_img_list.to("cuda")
    with torch.no_grad():
        output = detection_model(current_img_list, input_boxes=input_boxes)

    if input_boxes is None:
        feat_list, bbox_list = _process_feature_extraction(output, im_scales, feat_name)
        feat = feat_list[0].cpu().numpy()
        bbox = bbox_list[0].cpu().numpy() / im_scale
    else:
        feat = output[0][feat_name].cpu().numpy()
        bbox = output[0]["proposals"][0].bbox.cpu().numpy() / im_scale

    return feat, bbox

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--detection_cfg",
        type=str,
        default="/private/home/ronghanghu/workspace/pythia/data/"
        + "frcn_feature_extraction/detectron_model.yaml",
        help="Detectron config file; download it from "
        + "https://dl.fbaipublicfiles.com/pythia/detectron_model/"
        + "detectron_model.yaml",
    )
    parser.add_argument(
        "--detection_model",
        type=str,
        default="/private/home/ronghanghu/workspace/pythia/data/"
        + "frcn_feature_extraction/detectron_model.pth",
        help="Detectron model file; download it"
        + " from https://dl.fbaipublicfiles.com/pythia/detectron_model/"
        + "detectron_model.pth",
    )
    parser.add_argument(
        "--imdb_file",
        type=str,
        default="/private/home/ronghanghu/workspace/pythia/data/"
        + "imdb/m4c_textvqa/imdb_train_ocr_en.npy",
        help="The imdb to extract features",
    )
    parser.add_argument(
        "--image_dir",
        type=str,
        default="/private/home/ronghanghu/workspace/DATASETS/TextVQA",
        help="The directory containing images",
    )
    parser.add_argument(
        "--save_dir",
        type=str,
        default="/private/home/ronghanghu/workspace/pythia/data/"
        + "m4c_textvqa_ocr_en_frcn_features_2/train_images",
        help="The directory to save extracted features",
    )
    args = parser.parse_args()

    DETECTION_YAML = args.detection_cfg
    DETECTION_CKPT = args.detection_model
    IMDB_FILE = args.imdb_file
    IMAGE_DIR = args.image_dir
    SAVE_DIR = args.save_dir

    # imdb = np.load(IMDB_FILE, allow_pickle=True)[1:]
    imdb = np.load(IMDB_FILE, allow_pickle=True)
    # keep only one entry per image_id
    image_id2info = {info["image_id"]: info for info in imdb}
    imdb = list(image_id2info[k] for k in sorted(image_id2info))

    detection_model = load_detection_model(DETECTION_YAML, DETECTION_CKPT)
    print("Faster R-CNN OCR features")
    print("\textracting from", IMDB_FILE)
    print("\tsaving to", SAVE_DIR)
    for _, info in enumerate(tqdm.tqdm(imdb)):
        image_path = os.path.join(IMAGE_DIR, info["image_path"])
        save_feat_path = os.path.join(SAVE_DIR, info["feature_path"])
        save_info_path = save_feat_path.replace(".npy", "_info.npy")
        os.makedirs(os.path.dirname(save_feat_path), exist_ok=True)

        w = info["image_width"]
        h = info["image_height"]
        ocr_normalized_boxes = np.array(info["ocr_normalized_boxes"])
        ocr_boxes = ocr_normalized_boxes.reshape(-1, 4) * [w, h, w, h]
        ocr_tokens = info["ocr_tokens"]
        if len(ocr_boxes) > 0:
            extracted_feat, _ = extract_features(
                detection_model, image_path, input_boxes=ocr_boxes
            )
        else:
            extracted_feat = np.zeros((0, 2048), np.float32)

        np.save(save_info_path, {"ocr_boxes": ocr_boxes, "ocr_tokens": ocr_tokens})
        np.save(save_feat_path, extracted_feat)

if __name__ == "__main__":
    main()