# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
import os.path as osp
import numpy as np
import tqdm
import math
import cv2
from multiprocessing.pool import ThreadPool
import paddle.fluid as fluid
import paddlex.utils.logging as logging
import paddlex
from paddlex.cv.transforms import arrange_transforms
from paddlex.cv.datasets import generate_minibatch
from collections import OrderedDict
from .base import BaseAPI
from .utils.seg_eval import ConfusionMatrix
from .utils.visualize import visualize_segmentation

class DeepLabv3p(BaseAPI):
    """实现DeepLabv3+网络的构建并进行训练、评估、预测和模型导出。

    Args:
        num_classes (int): 类别数。
        backbone (str): DeepLabv3+的backbone网络，实现特征图的计算，取值范围为['Xception65', 'Xception41',
            'MobileNetV2_x0.25', 'MobileNetV2_x0.5', 'MobileNetV2_x1.0', 'MobileNetV2_x1.5',
            'MobileNetV2_x2.0', 'MobileNetV3_large_x1_0_ssld']。默认'MobileNetV2_x1.0'。
        output_stride (int): backbone 输出特征图相对于输入的下采样倍数，一般取值为8或16。默认16。
        aspp_with_sep_conv (bool):  在asspp模块是否采用separable convolutions。默认True。
        decoder_use_sep_conv (bool)： decoder模块是否采用separable convolutions。默认True。
        encoder_with_aspp (bool): 是否在encoder阶段采用aspp模块。默认True。
        enable_decoder (bool): 是否使用decoder模块。默认True。
        use_bce_loss (bool): 是否使用bce loss作为网络的损失函数，只能用于两类分割。可与dice loss同时使用。默认False。
        use_dice_loss (bool): 是否使用dice loss作为网络的损失函数，只能用于两类分割，可与bce loss同时使用，
            当use_bce_loss和use_dice_loss都为False时，使用交叉熵损失函数。默认False。
        class_weight (list/str): 交叉熵损失函数各类损失的权重。当class_weight为list的时候，长度应为
            num_classes。当class_weight为str时， weight.lower()应为'dynamic'，这时会根据每一轮各类像素的比重
            自行计算相应的权重，每一类的权重为：每类的比例 * num_classes。class_weight取默认值None时，各类的权重1，
            即平时使用的交叉熵损失函数。
        ignore_index (int): label上忽略的值，label为ignore_index的像素不参与损失函数的计算。默认255。
        pooling_crop_size (list): 当backbone为MobileNetV3_large_x1_0_ssld时，需设置为训练过程中模型输入大小, 格式为[W, H]。
            在encoder模块中获取图像平均值时被用到，若为None，则直接求平均值；若为模型输入大小，则使用'pool'算子得到平均值。
            默认值为None。
    Raises:
        ValueError: use_bce_loss或use_dice_loss为真且num_calsses > 2。
        ValueError: backbone取值不在['Xception65', 'Xception41', 'MobileNetV2_x0.25',
            'MobileNetV2_x0.5', 'MobileNetV2_x1.0', 'MobileNetV2_x1.5', 'MobileNetV2_x2.0', 'MobileNetV3_large_x1_0_ssld']之内。
        ValueError: class_weight为list, 但长度不等于num_class。
                class_weight为str, 但class_weight.low()不等于dynamic。
        TypeError: class_weight不为None时，其类型不是list或str。
    """

    def __init__(self,
                 num_classes=2,
                 backbone='MobileNetV2_x1.0',
                 output_stride=16,
                 aspp_with_sep_conv=True,
                 decoder_use_sep_conv=True,
                 encoder_with_aspp=True,
                 enable_decoder=True,
                 use_bce_loss=False,
                 use_dice_loss=False,
                 class_weight=None,
                 ignore_index=255,
                 pooling_crop_size=None):
        self.init_params = locals()
        super(DeepLabv3p, self).__init__('segmenter')
        # dice_loss或bce_loss只适用两类分割中
        if num_classes > 2 and (use_bce_loss or use_dice_loss):
            raise ValueError(
                "dice loss and bce loss is only applicable to binary classfication"
            )

        self.output_stride = output_stride

        if backbone not in [
                'Xception65', 'Xception41', 'MobileNetV2_x0.25',
                'MobileNetV2_x0.5', 'MobileNetV2_x1.0', 'MobileNetV2_x1.5',
                'MobileNetV2_x2.0', 'MobileNetV3_large_x1_0_ssld'
        ]:
            raise ValueError(
                "backbone: {} is set wrong. it should be one of "
                "('Xception65', 'Xception41', 'MobileNetV2_x0.25', 'MobileNetV2_x0.5',"
                " 'MobileNetV2_x1.0', 'MobileNetV2_x1.5', 'MobileNetV2_x2.0', 'MobileNetV3_large_x1_0_ssld')".
                format(backbone))

        if class_weight is not None:
            if isinstance(class_weight, list):
                if len(class_weight) != num_classes:
                    raise ValueError(
                        "Length of class_weight should be equal to number of classes"
                    )
            elif isinstance(class_weight, str):
                if class_weight.lower() != 'dynamic':
                    raise ValueError(
                        "if class_weight is string, must be dynamic!")
            else:
                raise TypeError(
                    'Expect class_weight is a list or string but receive {}'.
                    format(type(class_weight)))

        self.backbone = backbone
        self.num_classes = num_classes
        self.use_bce_loss = use_bce_loss
        self.use_dice_loss = use_dice_loss
        self.class_weight = class_weight
        self.ignore_index = ignore_index
        self.aspp_with_sep_conv = aspp_with_sep_conv
        self.decoder_use_sep_conv = decoder_use_sep_conv
        self.encoder_with_aspp = encoder_with_aspp
        self.enable_decoder = enable_decoder
        self.labels = None
        self.sync_bn = True
        self.fixed_input_shape = None
        self.pooling_stride = [1, 1]
        self.pooling_crop_size = pooling_crop_size
        self.aspp_with_se = False
        self.se_use_qsigmoid = False
        self.aspp_convs_filters = 256
        self.aspp_with_concat_projection = True
        self.add_image_level_feature = True
        self.use_sum_merge = False
        self.conv_filters = 256
        self.output_is_logits = False
        self.backbone_lr_mult_list = None
        if 'MobileNetV3' in backbone:
            self.output_stride = 32
            self.pooling_stride = (4, 5)
            self.aspp_with_se = True
            self.se_use_qsigmoid = True
            self.aspp_convs_filters = 128
            self.aspp_with_concat_projection = False
            self.add_image_level_feature = False
            self.use_sum_merge = True
            self.output_is_logits = True
            if self.output_is_logits:
                self.conv_filters = self.num_classes
            self.backbone_lr_mult_list = [0.15, 0.35, 0.65, 0.85, 1]

    def _get_backbone(self, backbone):
        def mobilenetv2(backbone):
            # backbone: xception结构配置
            # output_stride：下采样倍数
            # end_points: mobilenetv2的block数
            # decode_point: 从mobilenetv2中引出分支所在block数, 作为decoder输入
            if '0.25' in backbone:
                scale = 0.25
            elif '0.5' in backbone:
                scale = 0.5
            elif '1.0' in backbone:
                scale = 1.0
            elif '1.5' in backbone:
                scale = 1.5
            elif '2.0' in backbone:
                scale = 2.0
            end_points = 18
            decode_points = 4
            return paddlex.cv.nets.MobileNetV2(
                scale=scale,
                output_stride=self.output_stride,
                end_points=end_points,
                decode_points=decode_points)

        def xception(backbone):
            # decode_point: 从Xception中引出分支所在block数，作为decoder输入
            # end_point：Xception的block数
            if '65' in backbone:
                decode_points = 2
                end_points = 21
                layers = 65
            if '41' in backbone:
                decode_points = 2
                end_points = 13
                layers = 41
            if '71' in backbone:
                decode_points = 3
                end_points = 23
                layers = 71
            return paddlex.cv.nets.Xception(
                layers=layers,
                output_stride=self.output_stride,
                end_points=end_points,
                decode_points=decode_points)

        def mobilenetv3(backbone):
            scale = 1.0
            lr_mult_list = self.backbone_lr_mult_list
            return paddlex.cv.nets.MobileNetV3(
                scale=scale,
                model_name='large',
                output_stride=self.output_stride,
                lr_mult_list=lr_mult_list,
                for_seg=True)

        if 'Xception' in backbone:
            return xception(backbone)
        elif 'MobileNetV2' in backbone:
            return mobilenetv2(backbone)
        elif 'MobileNetV3' in backbone:
            return mobilenetv3(backbone)

    def build_net(self, mode='train'):
        model = paddlex.cv.nets.segmentation.DeepLabv3p(
            self.num_classes,
            mode=mode,
            backbone=self._get_backbone(self.backbone),
            output_stride=self.output_stride,
            aspp_with_sep_conv=self.aspp_with_sep_conv,
            decoder_use_sep_conv=self.decoder_use_sep_conv,
            encoder_with_aspp=self.encoder_with_aspp,
            enable_decoder=self.enable_decoder,
            use_bce_loss=self.use_bce_loss,
            use_dice_loss=self.use_dice_loss,
            class_weight=self.class_weight,
            ignore_index=self.ignore_index,
            fixed_input_shape=self.fixed_input_shape,
            pooling_stride=self.pooling_stride,
            pooling_crop_size=self.pooling_crop_size,
            aspp_with_se=self.aspp_with_se,
            se_use_qsigmoid=self.se_use_qsigmoid,
            aspp_convs_filters=self.aspp_convs_filters,
            aspp_with_concat_projection=self.aspp_with_concat_projection,
            add_image_level_feature=self.add_image_level_feature,
            use_sum_merge=self.use_sum_merge,
            conv_filters=self.conv_filters,
            output_is_logits=self.output_is_logits)
        inputs = model.generate_inputs()
        model_out = model.build_net(inputs)
        outputs = OrderedDict()
        if mode == 'train':
            self.optimizer.minimize(model_out)
            outputs['loss'] = model_out
        else:
            outputs['pred'] = model_out[0]
            outputs['logit'] = model_out[1]
        return inputs, outputs

    def default_optimizer(self,
                          learning_rate,
                          num_epochs,
                          num_steps_each_epoch,
                          lr_decay_power=0.9):
        decay_step = num_epochs * num_steps_each_epoch
        lr_decay = fluid.layers.polynomial_decay(
            learning_rate,
            decay_step,
            end_learning_rate=0,
            power=lr_decay_power)
        optimizer = fluid.optimizer.Momentum(
            lr_decay,
            momentum=0.9,
            regularization=fluid.regularizer.L2Decay(
                regularization_coeff=4e-05))
        return optimizer

    def train(self,
              num_epochs,
              train_dataset,
              train_batch_size=2,
              eval_dataset=None,
              save_interval_epochs=1,
              log_interval_steps=2,
              save_dir='output',
              pretrain_weights='IMAGENET',
              optimizer=None,
              learning_rate=0.01,
              lr_decay_power=0.9,
              use_vdl=False,
              sensitivities_file=None,
              eval_metric_loss=0.05,
              early_stop=False,
              early_stop_patience=5,
              resume_checkpoint=None):
        """训练。

        Args:
            num_epochs (int): 训练迭代轮数。
            train_dataset (paddlex.datasets): 训练数据读取器。
            train_batch_size (int): 训练数据batch大小。同时作为验证数据batch大小。默认为2。
            eval_dataset (paddlex.datasets): 评估数据读取器。
            save_interval_epochs (int): 模型保存间隔（单位：迭代轮数）。默认为1。
            log_interval_steps (int): 训练日志输出间隔（单位：迭代次数）。默认为2。
            save_dir (str): 模型保存路径。默认'output'。
            pretrain_weights (str): 若指定为路径时，则加载路径下预训练模型；若为字符串'IMAGENET'，
                则自动下载在ImageNet图片数据上预训练的模型权重；若为字符串'COCO'，
                则自动下载在COCO数据集上预训练的模型权重；若为字符串'CITYSCAPES'，
                则自动下载在CITYSCAPES数据集上预训练的模型权重；若为None，则不使用预训练模型。默认'IMAGENET。
            optimizer (paddle.fluid.optimizer): 优化器。当该参数为None时，使用默认的优化器：使用
                fluid.optimizer.Momentum优化方法，polynomial的学习率衰减策略。
            learning_rate (float): 默认优化器的初始学习率。默认0.01。
            lr_decay_power (float): 默认优化器学习率衰减指数。默认0.9。
            use_vdl (bool): 是否使用VisualDL进行可视化。默认False。
            sensitivities_file (str): 若指定为路径时，则加载路径下敏感度信息进行裁剪；若为字符串'DEFAULT'，
                则自动下载在Cityscapes图片数据上获得的敏感度信息进行裁剪；若为None，则不进行裁剪。默认为None。
            eval_metric_loss (float): 可容忍的精度损失。默认为0.05。
            early_stop (bool): 是否使用提前终止训练策略。默认值为False。
            early_stop_patience (int): 当使用提前终止训练策略时，如果验证集精度在`early_stop_patience`个epoch内
                连续下降或持平，则终止训练。默认值为5。
            resume_checkpoint (str): 恢复训练时指定上次训练保存的模型路径。若为None，则不会恢复训练。默认值为None。

        Raises:
            ValueError: 模型从inference model进行加载。
        """
        if not self.trainable:
            raise ValueError("Model is not trainable from load_model method.")

        self.labels = train_dataset.labels

        if optimizer is None:
            num_steps_each_epoch = train_dataset.num_samples // train_batch_size
            optimizer = self.default_optimizer(
                learning_rate=learning_rate,
                num_epochs=num_epochs,
                num_steps_each_epoch=num_steps_each_epoch,
                lr_decay_power=lr_decay_power)

        self.optimizer = optimizer
        # 构建训练、验证、预测网络
        self.build_program()
        # 初始化网络权重
        self.net_initialize(
            startup_prog=fluid.default_startup_program(),
            pretrain_weights=pretrain_weights,
            save_dir=save_dir,
            sensitivities_file=sensitivities_file,
            eval_metric_loss=eval_metric_loss,
            resume_checkpoint=resume_checkpoint)
        # 训练
        self.train_loop(
            num_epochs=num_epochs,
            train_dataset=train_dataset,
            train_batch_size=train_batch_size,
            eval_dataset=eval_dataset,
            save_interval_epochs=save_interval_epochs,
            log_interval_steps=log_interval_steps,
            save_dir=save_dir,
            use_vdl=use_vdl,
            early_stop=early_stop,
            early_stop_patience=early_stop_patience)

    def evaluate(self,
                 eval_dataset,
                 batch_size=1,
                 epoch_id=None,
                 return_details=False):
        """评估。

        Args:
            eval_dataset (paddlex.datasets): 评估数据读取器。
            batch_size (int): 评估时的batch大小。默认1。
            epoch_id (int): 当前评估模型所在的训练轮数。
            return_details (bool): 是否返回详细信息。默认False。

        Returns:
            dict: 当return_details为False时，返回dict。包含关键字：'miou'、'category_iou'、'macc'、
                'category_acc'和'kappa'，分别表示平均iou、各类别iou、平均准确率、各类别准确率和kappa系数。
            tuple (metrics, eval_details)：当return_details为True时，增加返回dict (eval_details)，
                包含关键字：'confusion_matrix'，表示评估的混淆矩阵。
        """
        arrange_transforms(
            model_type=self.model_type,
            class_name=self.__class__.__name__,
            transforms=eval_dataset.transforms,
            mode='eval')
        total_steps = math.ceil(eval_dataset.num_samples * 1.0 / batch_size)
        conf_mat = ConfusionMatrix(self.num_classes, streaming=True)
        data_generator = eval_dataset.generator(
            batch_size=batch_size, drop_last=False)
        if not hasattr(self, 'parallel_test_prog'):
            with fluid.scope_guard(self.scope):
                self.parallel_test_prog = fluid.CompiledProgram(
                    self.test_prog).with_data_parallel(
                        share_vars_from=self.parallel_train_prog)
        logging.info(
            "Start to evaluating(total_samples={}, total_steps={})...".format(
                eval_dataset.num_samples, total_steps))
        for step, data in tqdm.tqdm(
                enumerate(data_generator()), total=total_steps):
            images = np.array([d[0] for d in data])
            im_info = [d[1] for d in data]
            labels = [d[2] for d in data]

            num_samples = images.shape[0]
            if num_samples < batch_size:
                num_pad_samples = batch_size - num_samples
                pad_images = np.tile(images[0:1], (num_pad_samples, 1, 1, 1))
                images = np.concatenate([images, pad_images])
            feed_data = {'image': images}
            with fluid.scope_guard(self.scope):
                outputs = self.exe.run(
                    self.parallel_test_prog,
                    feed=feed_data,
                    fetch_list=list(self.test_outputs.values()),
                    return_numpy=True)
            pred = outputs[0]
            if num_samples < batch_size:
                pred = pred[0:num_samples]

            for i in range(num_samples):
                one_pred = np.squeeze(pred[i]).astype('uint8')
                one_label = labels[i]
                for info in im_info[i][::-1]:
                    if info[0] == 'resize':
                        w, h = info[1][1], info[1][0]
                        one_pred = cv2.resize(one_pred, (w, h),
                                              cv2.INTER_NEAREST)
                    elif info[0] == 'padding':
                        w, h = info[1][1], info[1][0]
                        one_pred = one_pred[0:h, 0:w]
                one_pred = one_pred.astype('int64')
                one_pred = one_pred[np.newaxis, :, :, np.newaxis]
                one_label = one_label[np.newaxis, np.newaxis, :, :]
                mask = one_label != self.ignore_index
                conf_mat.calculate(pred=one_pred, label=one_label, ignore=mask)
            _, iou = conf_mat.mean_iou()
            logging.debug("[EVAL] Epoch={}, Step={}/{}, iou={}".format(
                epoch_id, step + 1, total_steps, iou))

        category_iou, miou = conf_mat.mean_iou()
        category_acc, macc = conf_mat.accuracy()

        metrics = OrderedDict(
            zip(['miou', 'category_iou', 'macc', 'category_acc', 'kappa'],
                [miou, category_iou, macc, category_acc, conf_mat.kappa()]))
        if return_details:
            eval_details = {
                'confusion_matrix': conf_mat.confusion_matrix.tolist()
            }
            return metrics, eval_details
        return metrics

    @staticmethod
    def _preprocess(images, transforms, model_type, class_name, thread_num=1):
        arrange_transforms(
            model_type=model_type,
            class_name=class_name,
            transforms=transforms,
            mode='test')
        pool = ThreadPool(thread_num)
        batch_data = pool.map(transforms, images)
        pool.close()
        pool.join()
        padding_batch = generate_minibatch(batch_data)
        im = np.array(
            [data[0] for data in padding_batch],
            dtype=padding_batch[0][0].dtype)
        im_info = [data[1] for data in padding_batch]
        return im, im_info

    @staticmethod
    def _postprocess(results, im_info):
        pred_list = list()
        logit_list = list()
        for i, (pred, logit) in enumerate(zip(results[0], results[1])):
            pred = pred.astype('uint8')
            pred = np.squeeze(pred).astype('uint8')
            logit = np.transpose(logit, (1, 2, 0))
            for info in im_info[i][::-1]:
                if info[0] == 'resize':
                    w, h = info[1][1], info[1][0]
                    pred = cv2.resize(pred, (w, h), cv2.INTER_NEAREST)
                    logit = cv2.resize(logit, (w, h), cv2.INTER_LINEAR)
                elif info[0] == 'padding':
                    w, h = info[1][1], info[1][0]
                    pred = pred[0:h, 0:w]
                    logit = logit[0:h, 0:w, :]
            pred_list.append(pred)
            logit_list.append(logit)

        preds = list()
        for pred, logit in zip(pred_list, logit_list):
            preds.append({'label_map': pred, 'score_map': logit})
        return preds

    def predict(self, img_file, transforms=None):
        """预测。
        Args:
            img_file(str|np.ndarray): 预测图像路径，或者是解码后的排列格式为（H, W, C）且类型为float32且为BGR格式的数组。
            transforms(paddlex.cv.transforms): 数据预处理操作。

        Returns:
            dict: 包含关键字'label_map'和'score_map', 'label_map'存储预测结果灰度图，
                像素值表示对应的类别，'score_map'存储各类别的概率，shape=(h, w, num_classes)
        """

        if transforms is None and not hasattr(self, 'test_transforms'):
            raise Exception("transforms need to be defined, now is None.")
        if isinstance(img_file, (str, np.ndarray)):
            images = [img_file]
        else:
            raise Exception("img_file must be str/np.ndarray")

        if transforms is None:
            transforms = self.test_transforms
        im, im_info = DeepLabv3p._preprocess(
            images, transforms, self.model_type, self.__class__.__name__)

        with fluid.scope_guard(self.scope):
            result = self.exe.run(self.test_prog,
                                  feed={'image': im},
                                  fetch_list=list(self.test_outputs.values()),
                                  use_program_cache=True)

        preds = DeepLabv3p._postprocess(result, im_info)
        return preds[0]

    def batch_predict(self, img_file_list, transforms=None, thread_num=2):
        """预测。
        Args:
            img_file_list(list|tuple): 对列表（或元组）中的图像同时进行预测，列表中的元素可以是图像路径
                也可以是解码后的排列格式为（H，W，C）且类型为float32且为BGR格式的数组。
            transforms(paddlex.cv.transforms): 数据预处理操作。
            thread_num (int): 并发执行各图像预处理时的线程数。

        Returns:
            list: 每个元素都为列表，表示各图像的预测结果。各图像的预测结果用字典表示，包含关键字'label_map'和'score_map', 'label_map'存储预测结果灰度图，
                像素值表示对应的类别，'score_map'存储各类别的概率，shape=(h, w, num_classes)
        """

        if transforms is None and not hasattr(self, 'test_transforms'):
            raise Exception("transforms need to be defined, now is None.")
        if not isinstance(img_file_list, (list, tuple)):
            raise Exception("im_file must be list/tuple")
        if transforms is None:
            transforms = self.test_transforms
        im, im_info = DeepLabv3p._preprocess(
            img_file_list, transforms, self.model_type,
            self.__class__.__name__, thread_num)

        with fluid.scope_guard(self.scope):
            result = self.exe.run(self.test_prog,
                                  feed={'image': im},
                                  fetch_list=list(self.test_outputs.values()),
                                  use_program_cache=True)

        preds = DeepLabv3p._postprocess(result, im_info)
        return preds


from typing import Any, Dict, List, Optional

import attr
import numpy as np

import habitat_sim.agent
import habitat_sim.bindings as hsim
from habitat_sim import errors, utils

@attr.s(auto_attribs=True)
class GreedyGeodesicFollower(object):
    r"""Greedily fits actions to follow the geodesic shortest path

    Args:
        pathfinder (hsim.PathFinder): Instance of the pathfinder that has the correct navmesh already loaded
        agent (habitat_sim.agent.Agent): Agent to fit actions for.  This agent's current configuration is used
            to specify the actions.  The fitted actions will also correspond to keys in the agents action_space.
            `None` is used to signify that the goal location has been reached
        goal_radius (Optional[float]): Specifies how close the agent must get to the goal in order for it to be considered
            reached.  If `None`, 0.75 times the agents step size is used.
    """

    pathfinder: hsim.PathFinder
    agent: habitat_sim.agent.Agent
    goal_radius: Optional[float] = attr.ib(default=None)
    action_mapping: Dict[hsim.GreedyFollowerCodes, Any] = attr.ib(
        init=False, factory=dict, repr=False
    )
    impl: hsim.GreedyGeodesicFollowerImpl = attr.ib(
        init=False, default=None, repr=False
    )
    forward_spec: habitat_sim.agent.ActuationSpec = attr.ib(
        init=False, default=None, repr=False
    )
    left_spec: habitat_sim.agent.ActuationSpec = attr.ib(
        init=False, default=None, repr=False
    )
    right_spec: habitat_sim.agent.ActuationSpec = attr.ib(
        init=False, default=None, repr=False
    )

    def __attrs_post_init__(self):
        self.action_mapping[hsim.GreedyFollowerCodes.STOP] = None

        key, spec = self._find_action("move_forward")
        self.forward_spec = spec
        self.action_mapping[hsim.GreedyFollowerCodes.FORWARD] = key

        key, spec = self._find_action("turn_left")
        self.left_spec = spec
        self.action_mapping[hsim.GreedyFollowerCodes.LEFT] = key

        key, spec = self._find_action("turn_right")
        self.right_spec = spec
        self.action_mapping[hsim.GreedyFollowerCodes.RIGHT] = key

        if self.goal_radius is None:
            self.goal_radius = 0.75 * self.forward_spec.amount

        self.impl = hsim.GreedyGeodesicFollowerImpl(
            self.pathfinder,
            self._move_forward,
            self._turn_left,
            self._turn_right,
            self.goal_radius,
            self.forward_spec.amount,
            np.deg2rad(self.left_spec.amount),
        )

    def _find_action(self, name):
        candidates = list(
            filter(
                lambda v: v[1].name == name,
                self.agent.agent_config.action_space.items(),
            )
        )

        assert (
            len(candidates) == 1
        ), f"Could not find an action spec corresponding to {name}"

        return candidates[0][0], candidates[0][1].actuation

    def _move_forward(self, obj: hsim.SceneNode):
        self.agent.controls(obj, "move_forward", self.forward_spec, True)

    def _turn_left(self, obj: hsim.SceneNode):
        self.agent.controls(obj, "turn_left", self.left_spec, True)

    def _turn_right(self, obj: hsim.SceneNode):
        self.agent.controls(obj, "turn_right", self.right_spec, True)

    def next_action_along(self, goal_pos: np.ndarray) -> Any:
        r"""Find the next action to greedily follow the geodesic shortest path from the agent's current position
        to get to the goal

        Args:
            goal_pos (np.array): The position of the goal

        Returns:
            Any: the action to take
        """
        state = self.agent.state
        next_act = self.impl.next_action_along(
            state.position, utils.quat_to_coeffs(state.rotation), goal_pos
        )

        if next_act == hsim.GreedyFollowerCodes.ERROR:
            raise errors.GreedyFollowerError()
        else:
            return self.action_mapping[next_act]

    def find_path(self, goal_pos: np.ndarray) -> List[Any]:
        r"""Finds the sequence actions that greedily follow the geodesic shortest path
        from the agent's current position to get to the goal.  This is roughly equivilent to just
        calling `next_action_along` until it returns `None`, but is faster

        Args:
            goal_pos (np.array): The position of the goal

        Returns:
            List[Any]: The list of actions to take.  Ends with `None`
        """
        state = self.agent.state
        path = self.impl.find_path(
            state.position, utils.quat_to_coeffs(state.rotation), goal_pos
        )

        if len(path) == 0:
            raise errors.GreedyFollowerError()

        path = list(map(lambda v: self.action_mapping[v], path))

        return path


from functools import partial
from typing import Dict, Set

import numpy as np
from omegaconf import ListConfig
from scipy.spatial.transform import Rotation as R

from calvin_env.utils.utils import timeit

class Tasks:
    def __init__(self, tasks):
        """
        A task is defined as a specific change between the start_info and end_info dictionaries.
        Use config file in conf/tasks/ to define tasks using the base task functions defined in this class
        """
        # register task functions from config file
        self.tasks = {name: partial(getattr(self, args[0]), *args[1:]) for name, args in dict(tasks).items()}
        # dictionary mapping from task name to task id
        self.task_to_id = {name: i for i, name in enumerate(self.tasks.keys())}
        # dictionary mapping from task id to task name
        self.id_to_task = {i: name for i, name in enumerate(self.tasks.keys())}

    def get_task_info(self, start_info: Dict, end_info: Dict) -> Set:
        """
        start_info: dict with scene info and robot info
        end_info: dict with scene info and robot info
        returns set with achieved tasks
        """
        # call functions that are registered in self.tasks
        return {
            task_name
            for task_name, function in self.tasks.items()
            if function(start_info=start_info, end_info=end_info)
        }

    def get_task_info_for_set(self, start_info: Dict, end_info: Dict, task_filter: Set) -> Set:
        """
        start_info: dict with scene info and robot info
        end_info: dict with scene info and robot info
        task_filter: set with task names to check
        returns set with achieved tasks
        """
        # call functions that are registered in self.tasks
        return {
            task_name
            for task_name, function in self.tasks.items()
            if task_name in task_filter and function(start_info=start_info, end_info=end_info)
        }

    @property
    def num_tasks(self):
        return len(self.tasks)

    @staticmethod
    def rotate_object(
            obj_name, z_degrees, x_y_threshold=30, z_threshold=180, movement_threshold=0.1, start_info=None,
            end_info=None
    ):
        """
        Returns True if the object with obj_name was rotated more than z_degrees degrees around the z-axis while not
        being rotated more than x_y_threshold degrees around the x or y axis.
        z_degrees is negative for clockwise rotations and positive for counter-clockwise rotations.
        """
        obj_start_info = start_info["scene_info"]["movable_objects"][obj_name]
        obj_end_info = end_info["scene_info"]["movable_objects"][obj_name]
        start_orn = R.from_quat(obj_start_info["current_orn"])
        end_orn = R.from_quat(obj_end_info["current_orn"])
        rotation = end_orn * start_orn.inv()
        x, y, z = rotation.as_euler("xyz", degrees=True)

        start_pos = np.array(obj_start_info["current_pos"])
        end_pos = np.array(obj_end_info["current_pos"])
        pos_diff = end_pos - start_pos
        if np.linalg.norm(pos_diff) > movement_threshold:
            return False

        end_contacts = set(c[2] for c in obj_end_info["contacts"])
        robot_uid = {start_info["robot_info"]["uid"]}
        if len(end_contacts - robot_uid) == 0:
            return False

        if z_degrees > 0:
            return z_degrees < z < z_threshold and abs(x) < x_y_threshold and abs(y) < x_y_threshold
        else:
            return z_degrees > z > -z_threshold and abs(x) < x_y_threshold and abs(y) < x_y_threshold

    @staticmethod
    def push_object(obj_name, x_direction, y_direction, start_info, end_info):
        """
        Returns True if the object with 'obj_name' was moved more than 'x_direction' meters in x direction
        (or 'y_direction' meters in y direction analogously).
        Note that currently x and y pushes are mutually exclusive, meaning that one of the arguments has to be 0.
        The sign matters, e.g. pushing an object to the right when facing the table coincides with a movement in
        positive x-direction.
        """
        assert x_direction * y_direction == 0 and x_direction + y_direction != 0
        obj_start_info = start_info["scene_info"]["movable_objects"][obj_name]
        obj_end_info = end_info["scene_info"]["movable_objects"][obj_name]
        start_pos = np.array(obj_start_info["current_pos"])
        end_pos = np.array(obj_end_info["current_pos"])
        pos_diff = end_pos - start_pos

        start_contacts = set(c[2] for c in obj_start_info["contacts"])
        end_contacts = set(c[2] for c in obj_end_info["contacts"])
        robot_uid = {start_info["robot_info"]["uid"]}

        # computing set difference to check if object had surface contact (excluding robot) at both times
        surface_contact = len(start_contacts - robot_uid) > 0 and (start_contacts - robot_uid) == (
            end_contacts - robot_uid
        )
        if not surface_contact:
            return False

        if x_direction > 0:
            return pos_diff[0] > x_direction
        elif x_direction < 0:
            return pos_diff[0] < x_direction

        if y_direction > 0:
            return pos_diff[1] > y_direction
        elif y_direction < 0:
            return pos_diff[1] < y_direction

    @staticmethod
    def lift_object(obj_name, z_direction, surface_body=None, surface_link=None, start_info=None, end_info=None):
        """
        Returns True if the object with 'obj_name' was grasped by the robot and lifted more than 'z_direction' meters.
        """
        assert z_direction > 0
        obj_start_info = start_info["scene_info"]["movable_objects"][obj_name]
        obj_end_info = end_info["scene_info"]["movable_objects"][obj_name]

        start_pos = np.array(obj_start_info["current_pos"])
        end_pos = np.array(obj_end_info["current_pos"])
        pos_diff = end_pos - start_pos
        z_diff = pos_diff[2]

        robot_uid = start_info["robot_info"]["uid"]
        start_contacts = set(c[2] for c in obj_start_info["contacts"])
        end_contacts = set(c[2] for c in obj_end_info["contacts"])

        surface_criterion = True
        if surface_body and surface_link is None:
            surface_uid = start_info["scene_info"]["fixed_objects"][surface_body]["uid"]
            surface_criterion = surface_uid in start_contacts
        elif surface_body and surface_link:
            surface_uid = start_info["scene_info"]["fixed_objects"][surface_body]["uid"]
            surface_link_id = start_info["scene_info"]["fixed_objects"][surface_body]["links"][surface_link]
            start_contacts_links = set((c[2], c[4]) for c in obj_start_info["contacts"])
            surface_criterion = (surface_uid, surface_link_id) in start_contacts_links

        return (
            z_diff > z_direction
            and robot_uid not in start_contacts
            and robot_uid in end_contacts
            and len(end_contacts) == 1
            and surface_criterion
        )

    @staticmethod
    def place_object(dest_body, dest_link=None, start_info=None, end_info=None):
        """
        Returns True if the object that the robot has currently lifted is placed on the body 'dest_body'
        (on 'dest_link' if provided).
        The robot may not touch the object after placing.
        """
        robot_uid = start_info["robot_info"]["uid"]

        robot_contacts_start = set(c[2] for c in start_info["robot_info"]["contacts"])
        robot_contacts_end = set(c[2] for c in end_info["robot_info"]["contacts"])
        if not len(robot_contacts_start) == 1:
            return False
        obj_uid = list(robot_contacts_start)[0]

        if obj_uid in robot_contacts_end:
            return False
        _obj_name = [k for k, v in start_info["scene_info"]["movable_objects"].items() if v["uid"] == obj_uid]
        if not len(_obj_name) == 1:
            return False
        obj_name = _obj_name[0]

        dest_uid = end_info["scene_info"]["fixed_objects"][dest_body]["uid"]

        object_contacts_start = set(c[2] for c in start_info["scene_info"]["movable_objects"][obj_name]["contacts"])
        if dest_link is None:
            object_contacts_end = set(c[2] for c in end_info["scene_info"]["movable_objects"][obj_name]["contacts"])
            return (
                robot_uid in object_contacts_start
                and len(object_contacts_start) == 1
                and dest_uid in object_contacts_end
            )
        else:
            dest_link_id = end_info["scene_info"]["fixed_objects"][dest_body]["links"][dest_link]
            end_contacts_links = set(
                (c[2], c[4]) for c in end_info["scene_info"]["movable_objects"][obj_name]["contacts"]
            )
            return (
                robot_uid in object_contacts_start
                and len(object_contacts_start) == 1
                and (dest_uid, dest_link_id) in end_contacts_links
            )

    @staticmethod
    def push_object_into(obj_name, src_body, src_link, dest_body, dest_link, start_info=None, end_info=None):
        """
        obj_name is either a list of object names or a string
        Returns True if the object / any of the objects changes contact from src_body to dest_body.
        The robot may neither touch the object at start nor end.
        """
        if isinstance(obj_name, (list, ListConfig)):
            return any(
                Tasks.push_object_into(ob, src_body, src_link, dest_body, dest_link, start_info, end_info)
                for ob in obj_name
            )
        robot_uid = start_info["robot_info"]["uid"]

        src_uid = start_info["scene_info"]["fixed_objects"][src_body]["uid"]
        src_link_id = start_info["scene_info"]["fixed_objects"][src_body]["links"][src_link]
        dest_uid = end_info["scene_info"]["fixed_objects"][dest_body]["uid"]
        dest_link_id = end_info["scene_info"]["fixed_objects"][dest_body]["links"][dest_link]

        start_contacts = set((c[2], c[4]) for c in start_info["scene_info"]["movable_objects"][obj_name]["contacts"])
        end_contacts = set((c[2], c[4]) for c in end_info["scene_info"]["movable_objects"][obj_name]["contacts"])
        return (
            robot_uid not in start_contacts | end_contacts
            and len(start_contacts) == 1
            and (src_uid, src_link_id) in start_contacts
            and (dest_uid, dest_link_id) in end_contacts
        )

    @staticmethod
    def move_door_abs(joint_name, start_threshold, end_threshold, start_info, end_info):
        """
        Returns True if the joint specified by 'obj_name' and 'joint_name' (e.g. a door or drawer)
        is moved from at least 'start_threshold' to 'end_threshold'.
        """
        start_joint_state = start_info["scene_info"]["doors"][joint_name]["current_state"][0]
        end_joint_state = end_info["scene_info"]["doors"][joint_name]["current_state"][0]

        if start_threshold < end_threshold:
            return start_joint_state < start_threshold < end_threshold < end_joint_state
        elif start_threshold > end_threshold:
            return start_joint_state > start_threshold > end_threshold > end_joint_state
        else:
            raise ValueError

    @staticmethod
    def move_door_rel(joint_name, threshold, start_info, end_info):
        """
        Returns True if the joint specified by 'obj_name' and 'joint_name' (e.g. a door or drawer)
        is moved from at least 'start_threshold' to 'end_threshold'.
        """
        start_joint_state = start_info["scene_info"]["doors"][joint_name]["current_state"]
        end_joint_state = end_info["scene_info"]["doors"][joint_name]["current_state"]

        return (
            0 < threshold < end_joint_state - start_joint_state or 0 > threshold > end_joint_state - start_joint_state
        )

    @staticmethod
    def toggle_light(light_name, start_state, end_state, start_info, end_info):
        return (
            start_info["scene_info"]["lights"][light_name]["logical_state"] == start_state
            and end_info["scene_info"]["lights"][light_name]["logical_state"] == end_state
        )

    @staticmethod
    def stack_objects(max_vel=1, start_info=None, end_info=None):
        obj_uids = set(obj["uid"] for obj in start_info["scene_info"]["movable_objects"].values())

        for obj_name in start_info["scene_info"]["movable_objects"]:
            obj_start_info = start_info["scene_info"]["movable_objects"][obj_name]
            obj_end_info = end_info["scene_info"]["movable_objects"][obj_name]
            obj_start_contacts = set(c[2] for c in obj_start_info["contacts"])
            obj_end_contacts = set(c[2] for c in obj_end_info["contacts"])

            if (
                not len(obj_uids & obj_start_contacts)
                and len(obj_uids & obj_end_contacts)
                and not len(obj_end_contacts - obj_uids)
            ):
                # object velocity may not exceed max_vel for successful stack
                if np.all(np.abs(obj_end_info["current_lin_vel"]) < max_vel) and np.all(
                    np.abs(obj_end_info["current_ang_vel"]) < max_vel
                ):
                    return True
        return False

    @staticmethod
    def unstack_objects(max_vel=1, start_info=None, end_info=None):
        obj_uids = set(obj["uid"] for obj in start_info["scene_info"]["movable_objects"].values())

        for obj_name in start_info["scene_info"]["movable_objects"]:
            obj_start_info = start_info["scene_info"]["movable_objects"][obj_name]
            obj_end_info = end_info["scene_info"]["movable_objects"][obj_name]
            obj_start_contacts = set(c[2] for c in obj_start_info["contacts"])
            obj_end_contacts = set(c[2] for c in obj_end_info["contacts"])

            if (
                len(obj_uids & obj_start_contacts)
                and not len(obj_start_contacts - obj_uids)
                and not len(obj_uids & obj_end_contacts)
            ):
                # object velocity may not exceed max_vel for successful stack
                if np.all(np.abs(obj_start_info["current_lin_vel"]) < max_vel) and np.all(
                    np.abs(obj_start_info["current_ang_vel"]) < max_vel
                ):
                    return True
        return False

    @staticmethod
    def slider_drawer_light_led(start_info, end_info):
        slider_left = Tasks.move_door_rel('base__slide', 0.15, start_info, end_info)
        drawer = Tasks.move_door_rel('base__drawer', 0.12, start_info, end_info)
        light = Tasks.toggle_light('lightbulb', 0, 1, start_info, end_info)
        led = Tasks.toggle_light('led', 0, 1, start_info, end_info)

        return (slider_left and drawer and light and led)

    @staticmethod
    def cleanup(blocks, start_info, end_info):
        block_in_drawer = []
        for block in blocks:
            block_in_drawer.append(0.30 <= end_info["scene_info"]["movable_objects"][block]['current_pos'][2] <= 0.40) # hack
        drawer = end_info["scene_info"]["doors"]["base__drawer"]["current_state"] < 0.02
        light = Tasks.toggle_light('lightbulb', 1, 0, start_info, end_info)
        led = Tasks.toggle_light('led', 1, 0, start_info, end_info)
        return (all(block_in_drawer) and drawer and light and led)

    @staticmethod
    def cleanup_articulated(start_info, end_info):
        drawer = end_info["scene_info"]["doors"]["base__drawer"]["current_state"] < 0.02
        light = Tasks.toggle_light('lightbulb', 1, 0, start_info, end_info)
        led = Tasks.toggle_light('led', 1, 0, start_info, end_info)
        slider_right = Tasks.move_door_rel('base__slide', -0.15, start_info, end_info)
        return (drawer and light and led and slider_right)


import numpy as np
import app.maskrcnn_utils as rcnn_utils
import app.maskrcnn_config as rcnn_config

config = rcnn_config.Config()

def mold_inputs(images):
    """Takes a list of images and modifies them to the format expected
    as an input to the neural network.
    images: List of image matrices [height,width,depth]. Images can have
        different sizes.

    Returns 3 Numpy matrices:
    molded_images: [N, h, w, 3]. Images resized and normalized.
    image_metas: [N, length of meta data]. Details about each image.
    windows: [N, (y1, x1, y2, x2)]. The portion of the image that has the
        original image (padding excluded).
    """
    molded_images = []
    image_metas = []
    windows = []
    for image in images:
        # Resize image
        # TODO: move resizing to mold_image()
        molded_image, window, scale, padding, crop = rcnn_utils.resize_image(
            image,
            min_dim=config.IMAGE_MIN_DIM,
            min_scale=config.IMAGE_MIN_SCALE,
            max_dim=config.IMAGE_MAX_DIM,
            mode=config.IMAGE_RESIZE_MODE)
        molded_image = rcnn_utils.mold_image(molded_image, config)
        # Build image_meta
        image_meta = rcnn_utils.compose_image_meta(
            0, image.shape, molded_image.shape, window, scale,
            np.zeros([config.NUM_CLASSES], dtype=np.int32))
        # Append
        molded_images.append(molded_image)
        windows.append(window)
        image_metas.append(image_meta)
    # Pack into arrays
    molded_images = np.stack(molded_images)
    image_metas = np.stack(image_metas)
    windows = np.stack(windows)
    return molded_images, image_metas, windows

def get_anchors(image_shape):
    """Returns anchor pyramid for the given image size."""
    backbone_shapes = rcnn_utils.compute_backbone_shapes(config, image_shape)
    # Cache anchors and reuse if image shape is the same
    # if not hasattr("_anchor_cache"):
    _anchor_cache = {}
    if not tuple(image_shape) in _anchor_cache:
        # Generate Anchors
        a = rcnn_utils.generate_pyramid_anchors(
            config.RPN_ANCHOR_SCALES,
            config.RPN_ANCHOR_RATIOS,
            backbone_shapes,
            config.BACKBONE_STRIDES,
            config.RPN_ANCHOR_STRIDE)
        # Keep a copy of the latest anchors in pixel coordinates because
        # it's used in inspect_model notebooks.
        # TODO: Remove this after the notebook are refactored to not use it
        anchors = a
        # Normalize coordinates
        _anchor_cache[tuple(image_shape)] = rcnn_utils.norm_boxes(a, image_shape[:2])
    return _anchor_cache[tuple(image_shape)]

def detect(images):
    """Runs the detection pipeline.

    images: List of images, potentially of different sizes.

    Returns a list of dicts, one dict per image. The dict contains:
    rois: [N, (y1, x1, y2, x2)] detection bounding boxes
    class_ids: [N] int class IDs
    scores: [N] float probability scores for the class IDs
    masks: [H, W, N] instance binary masks
    """

    # assert mode == "inference", "Create model in inference mode."
    assert len(images) == config.BATCH_SIZE, "len(images) must be equal to BATCH_SIZE"

    # Mold inputs to format expected by the neural network
    molded_images, image_metas, windows = mold_inputs(images)

    # Validate image sizes
    # All images in a batch MUST be of the same size
    image_shape = molded_images[0].shape
    for g in molded_images[1:]:
        assert g.shape == image_shape, \
            "After resizing, all images must have the same size. Check IMAGE_RESIZE_MODE and image sizes."

    # Anchors
    anchors = get_anchors(image_shape)
    # Duplicate across the batch dimension because Keras requires it
    # TODO: can this be optimized to avoid duplicating the anchors?
    anchors = np.broadcast_to(anchors, (config.BATCH_SIZE,) + anchors.shape)

    # # Run object detection
    # detections, _, _, mrcnn_mask, _, _, _ = \
    #     self.keras_model.predict([molded_images, image_metas, anchors], verbose=0)
    # # Process detections
    # results = []
    # for i, image in enumerate(images):
    #     final_rois, final_class_ids, final_scores, final_masks = \
    #         self.unmold_detections(detections[i], mrcnn_mask[i],
    #                                image.shape, molded_images[i].shape,
    #                                windows[i])
    #     results.append({
    #         "rois": final_rois,
    #         "class_ids": final_class_ids,
    #         "scores": final_scores,
    #         "masks": final_masks,
    #     })
    return molded_images, image_metas, anchors, windows


from mpi4py import MPI
from neuron import h
from Simulation import Simulation
from cells import AfferentFiber
import random as rnd
import time
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from tools import firings_tools as tlsf
import pickle
from tools import seed_handler as sh
sh.set_seed()

comm = MPI.COMM_WORLD
sizeComm = comm.Get_size()
rank = comm.Get_rank()

class ForwardSimulation(Simulation):
	""" Integration of a NeuralNetwork object over time given an input (ees or afferent input or both). """

	def __init__(self, parallelContext, neuralNetwork, afferentInput=None, eesObject=None, eesModulation=None, tStop = 100):
		""" Object initialization.

		Keyword arguments:
		parallelContext -- Neuron parallelContext object.
		neuralNetwork -- NeuralNetwork object.
		afferentInput -- Dictionary of lists for each type of fiber containing the
			fibers firing rate over time and the dt at wich the firing rate is updated.
			If no afferent input is desired use None (default = None).
		eesObject -- EES object connected to the NeuralNetwork, usefull for some plotting
			info and mandatory for eesModulation (Default = None).
		eesModulation -- possible dictionary with the following strucuture: {'modulation':
			dictionary containing a	signal of 0 and 1s used to activate/inactivate
			the stimulation for every muscle that we want to modulate (the dictionary
			keys have to be the muscle names used in the neural network structure), 'dt':
			modulation dt}. If no modulation of the EES is intended use None (default = None).
		tStop -- Time in ms at wich the simulation will stop (default = 100). In case
			the time is set to -1 the neuralNetwork will be integrated for all the duration
			of the afferentInput.
		"""

		Simulation.__init__(self,parallelContext)
		if rank==1:
			print "\nMPI execution: the cells are divided in the different hosts\n"
		self._nn = neuralNetwork
		self._Iaf = self._nn.get_primary_afferents_names()[0] if self._nn.get_primary_afferents_names() else []
		self._IIf = self._nn.get_secondary_afferents_names()[0] if self._nn.get_secondary_afferents_names() else []
		self._Mn = self._nn.get_motoneurons_names() if self._nn.get_motoneurons_names() else []

		self._set_integration_step(AfferentFiber.get_update_period())

		# Initialization of the afferent modulation
		if afferentInput == None:
			self._afferentModulation = False
			self._afferentInput = None
			if tStop>0: self._set_tstop(tStop)
			else : raise(Exception("If no afferents input are provided tStop has to be greater than 0."))
		else:
			self._afferentModulation = True
			self._afferentInput = afferentInput[0]
			self._dtUpdateAfferent = afferentInput[1]
			self._init_afferents_fr()
			key=[]
			key.append(self._afferentInput.keys()[0])
			key.append(self._afferentInput[key[0]].keys()[0])
			self._inputDuration = len(self._afferentInput[key[0]][key[1]])*self._dtUpdateAfferent
			if tStop == -1 or tStop>= self._inputDuration: self._set_tstop(self._inputDuration-self._dtUpdateAfferent)
			else : self._set_tstop(tStop)

		self._ees = eesObject
		# Initialization of the binary stim modulation
		if eesModulation == None or eesObject == None:
			self._eesBinaryModulation = False
			self._eesProportionalModulation = False
			self._eesParam = {'state':None, 'amp':None, 'modulation':None, 'dt':None}
		elif eesModulation['type']=="binary":
			self._eesBinaryModulation = True
			self._eesProportionalModulation = False
			current, percIf, percIIf, percMn = self._ees.get_amplitude()
			self._eesParam = {'state':{}, 'modulation':eesModulation['modulation'], 'dt':eesModulation['dt']}
			self._eesParam['amp'] = [percIf, percIIf, percMn]
			for muscle in eesModulation['modulation']:
				self._eesParam['state'][muscle] = 1
		elif eesModulation['type']=="proportional":
			self._eesBinaryModulation = False
			self._eesProportionalModulation = True
			self._eesParam = {'modulation':eesModulation['modulation'], 'dt':eesModulation['dt']}
			current, percIf, percIIf, percMn = self._ees.get_amplitude()
			self._eesParam['maxAmp'] = np.array([percIf, percIIf, percMn])

		#Initialization of the result dictionaries
		self._meanFr = None
		self._estimatedEMG = None
		self._nSpikes = None
		self._nActiveCells = None

	"""
	Redefinition of inherited methods
	"""
	def _initialize(self):
		Simulation._initialize(self)
		self._init_aff_fibers()
		self._timeUpdateAfferentsFr = 0
		self._timeUpdateEES = 0

	def _update(self):
		""" Update simulation parameters. """
		comm.Barrier()
		self._nn.update_afferents_ap(h.t)
		if self._afferentModulation:
			if h.t-self._timeUpdateAfferentsFr>= (self._dtUpdateAfferent-0.5*self._get_integration_step()):
				self._timeUpdateAfferentsFr = h.t
				self._set_afferents_fr(int(h.t/self._dtUpdateAfferent))
				self._nn.set_afferents_fr(self._afferentFr)

		if self._eesBinaryModulation:
			if h.t-self._timeUpdateEES>= (self._eesParam['dt']-0.5*self._get_integration_step()):
				ind = int(h.t/self._eesParam['dt'])
				for muscle in self._eesParam['modulation']:
					if self._eesParam['state'][muscle] != self._eesParam['modulation'][muscle][ind]:
						if self._eesParam['state'][muscle] == 0: self._ees.set_amplitude(self._eesParam['amp'],[muscle])
						else: self._ees.set_amplitude([0,0,0],[muscle])
						self._eesParam['state'][muscle] = self._eesParam['modulation'][muscle][ind]

		if self._eesProportionalModulation:
			if h.t-self._timeUpdateEES>= (self._eesParam['dt']-0.5*self._get_integration_step()):
				ind = int(h.t/self._eesParam['dt'])
				for muscle in self._eesParam['modulation']:
					amp = list(self._eesParam['maxAmp']*self._eesParam['modulation'][muscle][ind])
					self._ees.set_amplitude(amp,[muscle])

	def _end_integration(self,result_name):
		""" Print the total simulation time and extract the results. """
		Simulation._end_integration(self)
		self._extract_results(result_name)
		print "It is called"

	"""
	Specific Methods of this class
	"""
	def _init_aff_fibers(self):
		""" Return the percentage of afferent action potentials erased by the stimulation. """
		for muscleName in self._nn.cells:
			for cellName in self._nn.cells[muscleName]:
				if cellName in self._nn.get_afferents_names():
					for fiber in self._nn.cells[muscleName][cellName]:
						fiber.initialise()

	def _init_afferents_fr(self):
		""" Initialize the dictionary necessary to update the afferent fibers. """
		self._afferentFr = {}
		for muscle in self._afferentInput:
			self._afferentFr[muscle]={}
			for cellType in self._afferentInput[muscle]:
				if cellType in self._nn.get_afferents_names():
					self._afferentFr[muscle][cellType]= 0.
				else: raise(Exception("Wrong afferent input structure!"))

	def _set_afferents_fr(self,i):
		""" Set the desired firing rate in the _afferentFr dictionary. """
		for muscle in self._afferentInput:
			for cellType in self._afferentInput[muscle]:
				self._afferentFr[muscle][cellType] = self._afferentInput[muscle][cellType][i]

	def _extract_results(self,result_name,samplingRate = 1000.):
		""" Extract the simulation results. """
		if rank==0: print "Extracting the results... ",
		self._firings = {}
		self._meanFr = {}
		self._estimatedEMG = {}
		self._nSpikes = {}
		self._nActiveCells = {}
		for muscle in self._nn.actionPotentials:
			self._firings[muscle]={}
			self._meanFr[muscle]={}
			self._estimatedEMG[muscle]={}
			self._nSpikes[muscle]={}
			self._nActiveCells[muscle]={}
			for cell in self._nn.actionPotentials[muscle]:
				self._firings[muscle][cell] = tlsf.exctract_firings(self._nn.actionPotentials[muscle][cell],self._get_tstop(),samplingRate)
				if rank==0: self._nActiveCells[muscle][cell] = np.count_nonzero(np.sum(self._firings[muscle][cell],axis=1))
				self._nSpikes[muscle][cell] = np.sum(self._firings[muscle][cell])
				self._meanFr[muscle][cell] = tlsf.compute_mean_firing_rate(self._firings[muscle][cell],samplingRate)
				if cell in self._nn.get_motoneurons_names():
					self._estimatedEMG[muscle][cell] = tlsf.synth_rat_emg(self._firings[muscle][cell],samplingRate)
		if rank==0:
			resultsFolder = "../../results/"
			with open(resultsFolder+result_name + ".p", 'w') as pickle_file:
				pickle.dump(self._firings, pickle_file)
				pickle.dump(self._meanFr, pickle_file)
				pickle.dump(self._estimatedEMG, pickle_file)
				pickle.dump(self._nSpikes, pickle_file)
				pickle.dump(self._nActiveCells, pickle_file)

			fig, ax = plt.subplots(3, 2,figsize=(24,10),sharex='col',sharey='col')

			extMuscle = "GM"
			flexMuscle = "TA"
			ax[0,0].plot(self._meanFr[extMuscle]['Mn'],color='b')
			ax[0,0].set_title('Extensor MN')
			ax[0,1].plot(self._meanFr[flexMuscle]['Mn'],color='b')
			ax[0,1].set_title('Flexor MN')
			ax[1,0].plot(self._meanFr[extMuscle]['IaInt'],color='b')
			ax[1,0].set_title('Extensor IaInt')
			ax[1,1].plot(self._meanFr[flexMuscle]['IaInt'],color='b')
			ax[1,1].set_title('Flexor IaInt')
			ax[2,0].plot(self._meanFr[extMuscle]['RORa'],color='b')
			ax[2,0].set_title('Extensor RORa')
			ax[2,1].plot(self._meanFr[flexMuscle]['RORa'],color='b')
			ax[2,1].set_title('Flexor RORa')

			fileName = time.strftime(result_name+".pdf")
			pp = PdfPages(self._resultsFolder+fileName)
			pp.savefig(fig)
			pp.close()

			print "...completed."

	def get_estimated_emg(self,muscleName):
		emg = [self._estimatedEMG[muscleName][mnName] for mnName in self._Mn]
		emg = np.sum(emg,axis=0)
		return emg

	def get_mn_spikes_profile(self,muscleName):
		spikesProfile = [self._firings[muscleName][mnName] for mnName in self._Mn]
		spikesProfile = np.sum(spikesProfile,axis=0)
		spikesProfile = np.sum(spikesProfile,axis=0)
		return spikesProfile

	def _get_perc_aff_ap_erased(self,muscleName,cellName):
		""" Return the percentage of afferent action potentials erased by the stimulation. """
		if cellName in self._nn.get_afferents_names():
			percErasedAp = []
			meanPercErasedAp = None
			for fiber in self._nn.cells[muscleName][cellName]:
				sent,arrived,collisions,perc = fiber.get_stats()
				percErasedAp.append(perc)
			percErasedAp = comm.gather(percErasedAp,root=0)

			if rank==0:
				percErasedAp = sum(percErasedAp,[])
				meanPercErasedAp = np.array(percErasedAp).mean()

			meanPercErasedAp = comm.bcast(meanPercErasedAp,root=0)
			percErasedAp = comm.bcast(percErasedAp,root=0)
			return meanPercErasedAp,percErasedAp
		else: raise(Exception("The selected cell is not and afferent fiber!"))

	def plot(self,flexorMuscle,extensorMuscle,name="",showFig=True):
		""" Plot and save the simulation results.

		Keyword arguments:
		flexorMuscle -- flexor muscle to plot
		extensorMuscle -- extensor muscle to plot
		name -- string to add at predefined file name of the saved pdf.
		"""
		meanPercErasedApFlexIaf,percErasedApFlexIaf = self._get_perc_aff_ap_erased(flexorMuscle,self._Iaf)
		meanPercErasedApExtIaf,percErasedApExtIaf = self._get_perc_aff_ap_erased(extensorMuscle,self._Iaf)
		meanPercErasedApFlexIIf,percErasedApFlexIIf = self._get_perc_aff_ap_erased(flexorMuscle,self._IIf)
		meanPercErasedApExtIIf,percErasedApExtIIf = self._get_perc_aff_ap_erased(extensorMuscle,self._IIf)
		meanPerEraserApIaf = np.mean([meanPercErasedApFlexIaf,meanPercErasedApExtIaf])
		meanPerEraserApIIf = np.mean([meanPercErasedApFlexIIf,meanPercErasedApExtIIf])

		#should be on rank 0?
		if not self._ees == None:
			percFiberActEes = self._ees.get_amplitude()
			title = 'EES _ {0:.0f}uA _ {1:.0f}Hz _ Delay _ {2:.0f}ms '.format(percFiberActEes[0],self._ees.get_frequency(),\
				self._nn.cells[flexorMuscle][self._Iaf][0].get_delay())
			plt.suptitle(title+"\n Iaf = {0:.0f}%, IIf = {1:.0f}%, Mn = {2:.0f}%, PercErasedApIaf = {3:.0f}%, prasedApIIf = {4:.0f}% ".format(\
				100*percFiberActEes[1],100*percFiberActEes[2],100*percFiberActEes[3],meanPerEraserApIaf,meanPerEraserApIIf))
		else:
			title = ' No EES _ Delay _ {2:.0f} ms '.format(self._nn.cells[flexorMuscle][self._Iaf][0].get_delay())
			plt.suptitle(title+"\n PercErasedApIaf = {3:.0f}%, prasedApIIf = {4:.0f}% ".format(meanPerEraserApIaf,meanPerEraserApIIf))

		if rank == 0:
			if not self._nn.recordMotoneurons and not self._nn.recordMotoneurons:
				raise(Exception("To plot the results it is necessary to have the NeuralNetwork recordMotoneurons and recordAfferents flags set to True"))

			flexIaAfferentModel = self._meanFr[flexorMuscle][self._Iaf]
			flexIIAfferentModel = self._meanFr[flexorMuscle]['IIf']
			flexRAf_footAfferentModel = self._meanFr[flexorMuscle]['II_RAf_foot']
			flexSAIf_footAfferentModel = self._meanFr[flexorMuscle]['II_SAIf_foot']
			flexRAfAfferentModel = self._meanFr[flexorMuscle]['II_RAf']
			flexSAIfAfferentModel = self._meanFr[flexorMuscle]['II_SAIf']

			extIaAfferntModel = self._meanFr[extensorMuscle][self._Iaf]
			extIIAfferntModel = self._meanFr[extensorMuscle]['IIf']
			extRAf_footAfferntModel = self._meanFr[extensorMuscle]['II_RAf_foot']
			extSAIf_footAfferntModel = self._meanFr[extensorMuscle]['II_SAIf_foot']
			extRAfAfferntModel = self._meanFr[extensorMuscle]['II_RAf']
			extSAIfAfferntModel = self._meanFr[extensorMuscle]['II_SAIf']

			flexMnAll = [self._meanFr[flexorMuscle][mnName] for mnName in self._Mn]
			extMnAll = [self._meanFr[extensorMuscle][mnName] for mnName in self._Mn]
			flexMn = np.mean(flexMnAll,axis=0)
			extMn = np.mean(extMnAll,axis=0)

			if self._afferentModulation:
				flexIaAfferentInput = self._afferentInput[flexorMuscle][self._Iaf]
				flexIIAfferentInput = self._afferentInput[flexorMuscle]['IIf']
				flexRAf_footAfferentInput = self._afferentInput[flexorMuscle]['II_RAf_foot']
				flexSAIf_footAfferentInput = self._afferentInput[flexorMuscle]['II_SAIf_foot']
				flexRAfAfferentInput = self._afferentInput[flexorMuscle]['II_RAf']
				flexSAIfAfferentInput = self._afferentInput[flexorMuscle]['II_SAIf']

				extIaAfferentInput = self._afferentInput[extensorMuscle][self._Iaf]
				extIIAfferentInput = self._afferentInput[extensorMuscle]['IIf']
				extRAf_footAfferentInput = self._afferentInput[extensorMuscle]['II_RAf_foot']
				extSAIf_footAfferentInput = self._afferentInput[extensorMuscle]['II_SAIf_foot']
				extRAfAfferentInput = self._afferentInput[extensorMuscle]['II_RAf']
				extSAIfAfferentInput = self._afferentInput[extensorMuscle]['II_SAIf']

			tStop = self._get_tstop()
			info,temp = self._nn.get_mn_info()
			strInfo = []
			for line in info: strInfo.append(" ".join(line))

			fig1, ax1 = plt.subplots(2, 4,figsize=(24,10),sharex='col',sharey='col')
			""" Ia afferents fr subplots """
			ax1[0,0].plot(flexIaAfferentModel,color='b',label='model')
			if self._afferentModulation:
				ax1[0,0].plot(np.arange(0,tStop,self._dtUpdateAfferent),flexIaAfferentInput[:np.ceil(float(tStop)/self._dtUpdateAfferent).astype(int)],color='r',label='input')
			ax1[0,0].set_title('Ia afferents mean firing rate')
			ax1[0,0].legend(loc='upper right')
			ax1[0,0].set_ylabel("Ia firing rate (Hz) - flex")
			ax1[1,0].plot(extIaAfferntModel,color='b',label='model')
			if self._afferentModulation:
				ax1[1,0].plot(np.arange(0,tStop,self._dtUpdateAfferent),extIaAfferentInput[:np.ceil(float(tStop)/self._dtUpdateAfferent).astype(int)],color='r',label='input')
			ax1[1,0].legend(loc='upper right')
			ax1[1,0].set_ylabel("Ia firing rate (Hz) - ext")
			ax1[1,0].set_xlabel("time (ms)")

			""" afferents percAperased subplots """
			ax1[0,1].hist(percErasedApFlexIaf, bins=range(0,101,2), facecolor='green', alpha=0.75)
			ax1[0,1].set_title("Histogram of the Ap perc erased in Iaf")
			ax1[0,1].set_ylabel("Frquency - flex")
			ax1[0,1].axis([0, 100, 0, 60])
			ax1[1,1].hist(percErasedApExtIaf, bins=range(0,101,2), facecolor='green', alpha=0.75)
			ax1[1,1].set_ylabel("Frquency - ext")
			ax1[1,1].set_xlabel("Percentage %")
			ax1[1,1].axis([0, 100, 0, 60])

			""" Motoneurons subplots """
			ax1[0,2].plot(flexMn)
			ax1[0,2].set_title('Motoneurons mean firing rate')
			ax1[0,2].set_ylabel("Mn firing rate (Hz) - flex")
			ax1[1,2].plot(extMn)
			ax1[1,2].set_ylabel("Mn firing rate (Hz) - ext")
			ax1[1,2].set_xlabel("time (ms)")

			""" EMG plot """
			ax1[0,3].plot(self.get_estimated_emg(flexorMuscle),color='b',label='model')
			ax1[0,3].set_ylabel("amplutide (a.u.) - flex")
			ax1[0,3].set_xlabel("time (ms)")
			ax1[0,3].set_title('Estimated EMG')
			ax1[1,3].plot(self.get_estimated_emg(extensorMuscle),color='b',label='model')
			ax1[1,3].set_ylabel("amplutide (a.u.) - ext")
			ax1[1,3].set_xlabel("time (ms)")

			title = title.replace(" ","")
			if showFig:	plt.show()

			fig2, ax2 = plt.subplots(2, 5,figsize=(24,10),sharex='col',sharey='col')
			""" II afferents fr subplots """
			ax2[0,0].plot(flexIIAfferentModel,color='b',label='model')
			if self._afferentModulation:
				ax2[0,0].plot(np.arange(0,tStop,self._dtUpdateAfferent),flexIIAfferentInput[:np.ceil(float(tStop)/self._dtUpdateAfferent).astype(int)],color='r',label='input')
			ax2[0,0].set_title('II afferents mean firing rate')
			ax2[0,0].legend(loc='upper right')
			ax2[0,0].set_ylabel("II firing rate (Hz) - flex")
			ax2[1,0].plot(extIIAfferntModel,color='b',label='model')
			if self._afferentModulation:
				ax2[1,0].plot(np.arange(0,tStop,self._dtUpdateAfferent),extIIAfferentInput[:np.ceil(float(tStop)/self._dtUpdateAfferent).astype(int)],color='r',label='input')
			ax2[1,0].legend(loc='upper right')
			ax2[1,0].set_ylabel("II firing rate (Hz) - ext")
			ax2[1,0].set_xlabel("time (ms)")

			""" Cutaneous RA afferents fr subplots """
			#ax2[0,1].plot(flexRAfAfferentModel,color='b',label='model')
			#if self._afferentModulation:
			#	ax2[0,1].plot(np.arange(0,tStop,self._dtUpdateAfferent),flexRAfAfferentInput[:np.ceil(float(tStop)/self._dtUpdateAfferent).astype(int)],color='r',label='input')
			#ax2[0,1].set_title('RA afferents mean firing rate')
			#ax2[0,1].legend(loc='upper right')
			#ax2[0,1].set_ylabel("RA firing rate (Hz) - flex")
			#ax2[1,1].plot(extRAfAfferntModel,color='b',label='model')
			#if self._afferentModulation:
			#	ax2[1,1].plot(np.arange(0,tStop,self._dtUpdateAfferent),extRAfAfferentInput[:np.ceil(float(tStop)/self._dtUpdateAfferent).astype(int)],color='r',label='input')
			#ax2[1,1].legend(loc='upper right')
			#ax2[1,1].set_ylabel("RA firing rate (Hz) - ext")
			#ax2[1,1].set_xlabel("time (ms)")

			ax2[0,1].plot(flexRAfAfferentModel,color='b',label='model')
			if self._afferentModulation:
				ax2[0,1].plot(np.arange(0,tStop,self._dtUpdateAfferent),flexRAfAfferentInput[:np.ceil(float(tStop)/self._dtUpdateAfferent).astype(int)],color='r',label='input')
			ax2[0,1].set_title('RA afferents mean firing rate')
			ax2[0,1].legend(loc='upper right')
			ax2[0,1].set_ylabel("RA firing rate (Hz) - flex")
			ax2[1,1].plot(extRAfAfferntModel,color='b',label='model')
			if self._afferentModulation:
				ax2[1,1].plot(np.arange(0,tStop,self._dtUpdateAfferent),extRAfAfferentInput[:np.ceil(float(tStop)/self._dtUpdateAfferent).astype(int)],color='r',label='input')
			ax2[1,1].legend(loc='upper right')
			ax2[1,1].set_ylabel("RA firing rate (Hz) - ext")
			ax2[1,1].set_xlabel("time (ms)")

			""" Cutaneous RA foot afferents fr subplots """
			ax2[0,2].plot(flexRAf_footAfferentModel,color='b',label='model')
			if self._afferentModulation:
				ax2[0,2].plot(np.arange(0,tStop,self._dtUpdateAfferent),flexRAf_footAfferentInput[:np.ceil(float(tStop)/self._dtUpdateAfferent).astype(int)],color='r',label='input')
			ax2[0,2].set_title('RA foot afferents mean firing rate')
			ax2[0,2].legend(loc='upper right')
			ax2[0,2].set_ylabel("RA foot firing rate (Hz) - flex")
			ax2[1,2].plot(extRAf_footAfferntModel,color='b',label='model')
			if self._afferentModulation:
				ax2[1,2].plot(np.arange(0,tStop,self._dtUpdateAfferent),extRAf_footAfferentInput[:np.ceil(float(tStop)/self._dtUpdateAfferent).astype(int)],color='r',label='input')
			ax2[1,2].legend(loc='upper right')
			ax2[1,2].set_ylabel("RA foot firing rate (Hz) - ext")
			ax2[1,2].set_xlabel("time (ms)")

			""" Cutaneous SAI afferents fr subplots """
			ax2[0,3].plot(flexSAIfAfferentModel,color='b',label='model')
			if self._afferentModulation:
				ax2[0,3].plot(np.arange(0,tStop,self._dtUpdateAfferent),flexSAIfAfferentInput[:np.ceil(float(tStop)/self._dtUpdateAfferent).astype(int)],color='r',label='input')
			ax2[0,3].set_title('SAI afferents mean firing rate')
			ax2[0,3].legend(loc='upper right')
			ax2[0,3].set_ylabel("SAI firing rate (Hz) - flex")
			ax2[1,3].plot(extSAIfAfferntModel,color='b',label='model')
			if self._afferentModulation:
				ax2[1,3].plot(np.arange(0,tStop,self._dtUpdateAfferent),extSAIfAfferentInput[:np.ceil(float(tStop)/self._dtUpdateAfferent).astype(int)],color='r',label='input')
			ax2[1,3].legend(loc='upper right')
			ax2[1,3].set_ylabel("SAI firing rate (Hz) - ext")
			ax2[1,3].set_xlabel("time (ms)")

			""" Cutaneous SAI foot afferents fr subplots """
			ax2[0,4].plot(flexSAIf_footAfferentModel,color='b',label='model')
			if self._afferentModulation:
				ax2[0,4].plot(np.arange(0,tStop,self._dtUpdateAfferent),flexSAIf_footAfferentInput[:np.ceil(float(tStop)/self._dtUpdateAfferent).astype(int)],color='r',label='input')
			ax2[0,4].set_title('SAI foot afferents mean firing rate')
			ax2[0,4].legend(loc='upper right')
			ax2[0,4].set_ylabel("SAI foot firing rate (Hz) - flex")
			ax2[1,4].plot(extSAIf_footAfferntModel,color='b',label='model')
			if self._afferentModulation:
				ax2[1,4].plot(np.arange(0,tStop,self._dtUpdateAfferent),extSAIf_footAfferentInput[:np.ceil(float(tStop)/self._dtUpdateAfferent).astype(int)],color='r',label='input')
			ax2[1,4].legend(loc='upper right')
			ax2[1,4].set_ylabel("SAI foot firing rate (Hz) - ext")
			ax2[1,4].set_xlabel("time (ms)")

			fig3 = plt.figure()
			ax3 = fig3.add_subplot(111)
			ax3.text(0.5, 0.5,"\n".join(strInfo), horizontalalignment='center',verticalalignment='center',transform=ax3.transAxes)
			ax3.xaxis.set_visible(False)
			ax3.yaxis.set_visible(False)

			fileName = time.strftime("%Y_%m_%d_FS_"+title+name+".pdf")
			pp = PdfPages(self._resultsFolder+fileName)
			pp.savefig(fig1)
			pp.savefig(fig2)
			pp.savefig(fig3)
			pp.close()


import sys
from Detector.MtcnnDetector import MtcnnDetector
from Detector.detector import Detector
from Detector.fcn_detector import FcnDetector
from Train_Model.mtcnn_model import P_Net, R_Net, O_Net
import cv2
import numpy as np
import tensorflow as tf

def load_classification(model_path = 'checkpoints/kcnn_aug'):
    model = tf.keras.Sequential()
    model.add(
        tf.keras.layers.Conv2D(32, kernel_size=(3, 3), activation=tf.nn.relu, input_shape=(img_rows, img_cols, 1)))
    model.add(tf.keras.layers.Conv2D(64, (3, 3), activation=tf.nn.relu))
    model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))
    model.add(tf.keras.layers.Dropout(0.25))
    model.add(tf.keras.layers.Flatten())
    model.add(tf.keras.layers.Dense(128, activation=tf.nn.relu))
    model.add(tf.keras.layers.Dropout(0.5))
    model.add(tf.keras.layers.Dense(8, activation=tf.nn.softmax))

    model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

    model.load_weights(model_path)

    return model

def cropImg(im, x1, y1, x2, y2):
    # img = im.copy()
    # box = [x1, y1, x2, y2]
    # drawBoxes(img, box)
    # cv2.imshow("original", img)
    # cv2.waitKey(0)
    tmp = im[y1:y2, x1:x2].copy()
    # cv2.imshow("result", tmp)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    return tmp

# Configure
model_path = ['MTCNN_Model/PNet/PNet-30', 'MTCNN_Model/RNet/RNet-22', 'MTCNN_Model/ONet/ONet-22']
min_face_size = 24
thresh = [0.9, 0.8, 0.9]

class_names = ['0SingleOne', '1SingleTwo', '2SingleFour', '3SingleSix',
               '4SingleEight', '5SingleNine', '6SingleBad', '7SingleGood']

img_rows, img_cols = 28, 28

model = load_classification('CNN_Model/kcnn_aug')

stride = 2
slide_window = False
shuffle = False

detectors = [None, None, None]

PNet = FcnDetector(P_Net, model_path[0])
detectors[0] = PNet
RNet = Detector(R_Net, 24, 1, model_path[1])
detectors[1] = RNet
ONet = Detector(O_Net, 48, 1, model_path[2])
detectors[2] = ONet

mtcnn_detector = MtcnnDetector(detectors=detectors, min_face_size=min_face_size,
                               stride=stride, threshold=thresh, slide_window=slide_window)

cap = cv2.VideoCapture(0)
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 240)
MAX_X = 320
MAX_Y = 240
shot_photo = False
Real_time = True

img_list = []
sample_amount = 5
count = 0
if (shot_photo):
    while count < sample_amount:
        __, frame = cap.read()
        cv2.imshow("camera", frame)
        if cv2.waitKey(1) & 0xFF == ord('s'):
            img = frame
            img_list.append(img)
            cv2.imshow("taken", img)
            cv2.waitKey(0)
            count += 1
    cap.release()
    cv2.destroyAllWindows()

corpbbox = None

while (shot_photo):
    # fps = video_capture.get(cv2.CAP_PROP_FPS)
    for frame in img_list:
        ret = True
        if ret:
            image = np.array(frame)
            boxes_c,_ = mtcnn_detector.detect(image)

            for i in range(boxes_c.shape[0]):
                bbox = boxes_c[i, :4]
                score = boxes_c[i, 4]
                corpbbox = [int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])]
                # if score > thresh:
                cv2.rectangle(frame, (corpbbox[0], corpbbox[1]),
                              (corpbbox[2], corpbbox[3]), (255, 0, 0), 1)
                cv2.putText(frame, '{:.3f}'.format(score), (corpbbox[0], corpbbox[1] - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                            (0, 0, 255), 2)

            # time end
            cv2.imshow("", frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                continue
        else:
            print('img not find')
            break
    break

while (Real_time):
    # fps = video_capture.get(cv2.CAP_PROP_FPS)
    t1 = cv2.getTickCount()
    ret, frame = cap.read()
    if ret:
        image = np.array(frame)
        boxes_c,_ = mtcnn_detector.detect(image)

        t2 = cv2.getTickCount()
        t = (t2 - t1) / cv2.getTickFrequency()
        fps = 1.0 / t
        for i in range(boxes_c.shape[0]):
            bbox = boxes_c[i, :4]
            score = boxes_c[i, 4]
            corpbbox = [int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])]
            w = bbox[2] - bbox[0]
            h = bbox[3] - bbox[1]
            x1 = int(bbox[0] - 0.1 * w)
            y1 = int(bbox[1] - 0.1 * h)
            x2 = int(bbox[2] + 0.1 * w)
            y2 = int(bbox[3] + 0.1 * h)
            if x1 < 0:
                x1 = 0
            elif x1 > MAX_X:
                x1 = MAX_X
            if x2 < 0:
                x2 = 0
            elif x2 > MAX_X:
                x2 = MAX_X
            if y1 < 0:
                y1 = 0
            elif y1 > MAX_Y:
                y1 = MAX_Y
            if y2 < 0:
                y2 = 0
            elif y2 > MAX_Y:
                y2 = MAX_Y
            if x1>0x2 or y1>=y2:
                im = cropImg(frame, corpbbox[0], corpbbox[1], corpbbox[2], corpbbox[3])
            else:
                im = cropImg(frame, x1, y1, x2, y2)
            print((x1, y1), (x2, y2))

            im = cv2.resize(im, (28, 28))
            im = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
            im = np.expand_dims(im, axis=2)
            im = np.array([im])
            # print(im.shape)
            prediction = model.predict(im)
            predicted_label = np.argmax(prediction)

            cv2.putText(frame, class_names[predicted_label], (corpbbox[0], corpbbox[1] - 2), cv2.FONT_HERSHEY_SIMPLEX,
                        1.5, (0, 255, 0), 2)
            # if score > thresh:
            cv2.rectangle(frame, (corpbbox[0], corpbbox[1]),
                          (corpbbox[2], corpbbox[3]), (255, 0, 0), 1)
            cv2.putText(frame, '{:.3f}'.format(score), (corpbbox[0], corpbbox[1] - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                        (0, 0, 255), 2)
        cv2.putText(frame, '{:.4f}'.format(t) + " " + '{:.3f}'.format(fps), (10, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                    (255, 0, 255), 2)

        # time end
        cv2.imshow("", frame)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    else:
        print('device not find')
        break

cap.release()
cv2.destroyAllWindows()


import random

import numpy as np

from audiomentations.core.transforms_interface import BaseWaveformTransform
from audiomentations.core.utils import (
    calculate_rms,
)

class TanhDistortion(BaseWaveformTransform):
    """
    Apply tanh (hyperbolic tangent) distortion to the audio. This technique is sometimes
    used for adding distortion to guitar recordings. The tanh() function can give a rounded
    "soft clipping" kind of distortion, and the distortion amount is proportional to the
    loudness of the input and the pre-gain. Tanh is symmetric, so the positive and
    negative parts of the signal are squashed in the same way. This transform can be
    useful as data augmentation because it adds harmonics. In other words, it changes
    the timbre of the sound.

    See this page for examples: http://gdsp.hf.ntnu.no/lessons/3/17/
    """

    supports_multichannel = True

    def __init__(
        self, min_distortion: float = 0.01, max_distortion: float = 0.7, p: float = 0.5
    ):
        """
        :param min_distortion: Minimum amount of distortion (between 0 and 1)
        :param max_distortion: Maximum amount of distortion (between 0 and 1)
        :param p: The probability of applying this transform
        """
        super().__init__(p)
        assert 0 <= min_distortion <= 1
        assert 0 <= max_distortion <= 1
        assert min_distortion <= max_distortion
        self.min_distortion = min_distortion
        self.max_distortion = max_distortion

    def randomize_parameters(self, samples, sample_rate):
        super().randomize_parameters(samples, sample_rate)
        if self.parameters["should_apply"]:
            self.parameters["distortion_amount"] = random.uniform(
                self.min_distortion, self.max_distortion
            )

    def apply(self, samples, sample_rate):
        # Find out how much to pre-gain the audio to get a given amount of distortion
        percentile = 100 - 99 * self.parameters["distortion_amount"]
        threshold = np.percentile(abs(samples), percentile)
        gain_factor = 0.5 / (threshold + 1e-6)

        # Distort the audio
        distorted_samples = np.tanh(gain_factor * samples)

        # Scale the output so its loudness matches the input
        rms_before = calculate_rms(samples)
        if rms_before > 1e-9:
            rms_after = calculate_rms(distorted_samples)
            post_gain = rms_before / rms_after
            distorted_samples = post_gain * distorted_samples

        return distorted_samples