# ---------------------------------------------------------------------------
# Unified Panoptic Segmentation Network
#
# Modifications Copyright (c) 2019 Uber Technologies, Inc.
# ---------------------------------------------------------------------------
# Based on:
# --------------------------------------------------------
# Deformable Convolutional Networks
# Copyright (c) 2017 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
# and
# --------------------------------------------------------
# mx-maskrcnn
#
# 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.
# --------------------------------------------------------

import numpy as np
import numpy.random as npr
from .generate_anchors import generate_anchors
from .anchors import anchors_cython
from upsnet.config.config import config
from upsnet.bbox.bbox_transform import bbox_transform, bbox_overlaps
from upsnet.rpn.generate_anchors import get_field_of_anchors, compute_targets, unmap
from upsnet.bbox.bbox_transform import bbox_overlaps
import pickle

def assign_anchor(feat_shape, gt_boxes, im_info, feat_stride=16,
                  scales=(8, 16, 32), ratios=(0.5, 1, 2), allowed_border=0):
    """
    assign ground truth boxes to anchor positions
    :param feat_shape: infer output shape
    :param gt_boxes: assign ground truth
    :param im_info: filter out anchors overlapped with edges
    :param feat_stride: anchor position step
    :param scales: used to generate anchors, affects num_anchors (per location)
    :param ratios: aspect ratios of generated anchors
    :param allowed_border: filter out anchors with edge overlap > allowed_border
    :return: dict of label
    'label': of shape (batch_size, 1) <- (batch_size, num_anchors, feat_height, feat_width)
    'bbox_target': of shape (batch_size, num_anchors * 4, feat_height, feat_width)
    'bbox_inside_weight': *todo* mark the assigned anchors
    'bbox_outside_weight': used to normalize the bbox_loss, all weights sums to RPN_POSITIVE_WEIGHT
    """
    def _unmap(data, count, inds, fill=0):
        """" unmap a subset inds of data into original data of size count """
        if len(data.shape) == 1:
            ret = np.empty((count,), dtype=np.float32)
            ret.fill(fill)
            ret[inds] = data
        else:
            ret = np.empty((count,) + data.shape[1:], dtype=np.float32)
            ret.fill(fill)
            ret[inds, :] = data
        return ret

    DEBUG = False
    im_info = im_info[0]
    scales = np.array(scales, dtype=np.float32)
    base_anchors = generate_anchors(base_size=feat_stride, ratios=list(ratios), scales=scales)
    num_anchors = base_anchors.shape[0]
    feat_height, feat_width = feat_shape[-2:]

    # 1. generate proposals from bbox deltas and shifted anchors
    shift_x = np.arange(0, feat_width) * feat_stride
    shift_y = np.arange(0, feat_height) * feat_stride
    shift_x, shift_y = np.meshgrid(shift_x, shift_y)
    shifts = np.vstack((shift_x.ravel(), shift_y.ravel(), shift_x.ravel(), shift_y.ravel())).transpose()
    # add A anchors (1, A, 4) to
    # cell K shifts (K, 1, 4) to get
    # shift anchors (K, A, 4)
    # reshape to (K*A, 4) shifted anchors
    A = num_anchors
    K = shifts.shape[0]
    all_anchors = base_anchors.reshape((1, A, 4)) + shifts.reshape((1, K, 4)).transpose((1, 0, 2))
    all_anchors = all_anchors.reshape((K * A, 4))
    total_anchors = int(K * A)

    # only keep anchors inside the image
    inds_inside = np.where((all_anchors[:, 0] >= -allowed_border) &
                           (all_anchors[:, 1] >= -allowed_border) &
                           (all_anchors[:, 2] < im_info[1] + allowed_border) &
                           (all_anchors[:, 3] < im_info[0] + allowed_border))[0]

    # keep only inside anchors
    anchors = all_anchors[inds_inside, :]

    # label: 1 is positive, 0 is negative, -1 is dont care
    labels = np.empty((len(inds_inside),), dtype=np.float32)
    labels.fill(-1)

    if gt_boxes.size > 0:
        # overlap between the anchors and the gt boxes
        # overlaps (ex, gt)
        overlaps = bbox_overlaps(anchors.astype(np.float), gt_boxes.astype(np.float))
        argmax_overlaps = overlaps.argmax(axis=1)
        max_overlaps = overlaps[np.arange(len(inds_inside)), argmax_overlaps]
        gt_argmax_overlaps = overlaps.argmax(axis=0)
        gt_max_overlaps = overlaps[gt_argmax_overlaps, np.arange(overlaps.shape[1])]
        gt_argmax_overlaps = np.where(overlaps == gt_max_overlaps)[0]

        if not config.train.rpn_clobber_positives:
            # assign bg labels first so that positive labels can clobber them
            labels[max_overlaps < config.train.rpn_negative_overlap] = 0

        # fg label: for each gt, anchor with highest overlap
        labels[gt_argmax_overlaps] = 1

        # fg label: above threshold IoU
        labels[max_overlaps >= config.train.rpn_positive_overlap] = 1

        if config.train.rpn_clobber_positives:
            # assign bg labels last so that negative labels can clobber positives
            labels[max_overlaps < config.train.rpn_negative_overlap] = 0
    else:
        labels[:] = 0

    # subsample positive labels if we have too many
    num_fg = int(config.train.rpn_fg_fraction * config.train.rpn_batch_size)
    fg_inds = np.where(labels == 1)[0]
    if len(fg_inds) > num_fg:
        disable_inds = npr.choice(fg_inds, size=(len(fg_inds) - num_fg), replace=False)
        # if DEBUG:
        #     disable_inds = fg_inds[:(len(fg_inds) - num_fg)]
        labels[disable_inds] = -1

    # subsample negative labels if we have too many
    num_bg = config.train.rpn_batch_size - np.sum(labels == 1)
    bg_inds = np.where(labels == 0)[0]
    if len(bg_inds) > num_bg:
        disable_inds = npr.choice(bg_inds, size=(len(bg_inds) - num_bg), replace=False)
        labels[disable_inds] = -1

    bbox_targets = np.zeros((len(inds_inside), 4), dtype=np.float32)
    if gt_boxes.size > 0:
        bbox_targets[:] = bbox_transform(anchors, gt_boxes[argmax_overlaps, :4])

    bbox_weights = np.zeros((len(inds_inside), 4), dtype=np.float32)
    bbox_weights[labels == 1, :] = np.array(config.train.rpn_bbox_weights)

    # map up to original set of anchors
    labels = _unmap(labels, total_anchors, inds_inside, fill=-1)
    bbox_targets = _unmap(bbox_targets, total_anchors, inds_inside, fill=0)
    bbox_weights = _unmap(bbox_weights, total_anchors, inds_inside, fill=0)

    labels = labels.reshape((1, feat_height, feat_width, A)).transpose(0, 3, 1, 2)
    labels = labels.reshape((1, A * feat_height * feat_width))
    bbox_targets = bbox_targets.reshape((1, feat_height, feat_width, A * 4)).transpose(0, 3, 1, 2)
    bbox_weights = bbox_weights.reshape((1, feat_height, feat_width, A * 4)).transpose((0, 3, 1, 2))

    label = {'label': labels,
             'bbox_target': bbox_targets,
             'bbox_weight': bbox_weights}
    return label

def assign_pyramid_anchor(gt_boxes, im_info, feat_strides=(64, 32, 16, 8, 4),
                          scales=(8,), ratios=(0.5, 1, 2), allowed_border=0):
    """
    assign ground truth boxes to anchor positions
    :param feat_shape: infer output shape
    :param gt_boxes: assign ground truth
    :param im_info: filter out anchors overlapped with edges
    :param feat_stride: anchor position step
    :param scales: used to generate anchors, affects num_anchors (per location)
    :param ratios: aspect ratios of generated anchors
    :param allowed_border: filter out anchors with edge overlap > allowed_border
    :return: tuple
    labels: of shape (batch_size, 1) <- (batch_size, num_anchors, feat_height, feat_width)
    bbox_targets: of shape (batch_size, num_anchors * 4, feat_height, feat_width)
    bbox_weights: mark the assigned anchors
    """
    def _unmap(data, count, inds, fill=0):
        """" unmap a subset inds of data into original data of size count """
        if len(data.shape) == 1:
            ret = np.empty((count,), dtype=np.float32)
            ret.fill(fill)
            ret[inds] = data
        else:
            ret = np.empty((count,) + data.shape[1:], dtype=np.float32)
            ret.fill(fill)
            ret[inds, :] = data
        return ret

    DEBUG = False
    im_info = im_info[0]
    # clean up boxes
    nonneg = np.where(gt_boxes[:, 4] != -1)[0]
    gt_boxes = gt_boxes[nonneg]
    scales = np.array(scales, dtype=np.float32)

    anchors_list = []
    anchors_num_list = []
    inds_inside_list = []
    feat_infos = []
    A_list = []
    for i in range(len(feat_strides)):
        base_anchors = generate_anchors(base_size=feat_strides[i], ratios=list(ratios), scales=scales)
        num_anchors = base_anchors.shape[0]
        # feat_height, feat_width = feat_shape[i][-2:]
        feat_height, feat_width, s = im_info[0], im_info[1], feat_strides[i]
        s = s // 4
        feat_height, feat_width = int(np.ceil(feat_height / 2)) // 2, int(np.ceil(feat_width / 2)) // 2,
        while s > 1:
            feat_height, feat_width = int(np.ceil(feat_height / 2)), int(np.ceil(feat_width / 2))
            s = s // 2
        feat_stride = feat_strides[i]
        feat_infos.append([feat_height, feat_width])

        A = num_anchors
        A_list.append(A)
        K = feat_height * feat_width

        # shift_x = np.arange(0, feat_width) * feat_stride
        # shift_y = np.arange(0, feat_height) * feat_stride
        # shift_x, shift_y = np.meshgrid(shift_x, shift_y)
        # shifts = np.vstack((shift_x.ravel(), shift_y.ravel(), shift_x.ravel(), shift_y.ravel())).transpose()
        # all_anchors = base_anchors.reshape((1, A, 4)) + shifts.reshape((1, K, 4)).transpose((1, 0, 2))
        all_anchors = anchors_cython(feat_height, feat_width, feat_stride, base_anchors)
        all_anchors = all_anchors.reshape((K * A, 4))

        total_anchors = int(K * A)
        anchors_num_list.append(total_anchors)
        # only keep anchors inside the image
        inds_inside = np.where((all_anchors[:, 0] >= -allowed_border) &
                               (all_anchors[:, 1] >= -allowed_border) &
                               (all_anchors[:, 2] < im_info[1] + allowed_border) &
                               (all_anchors[:, 3] < im_info[0] + allowed_border))[0]
        if DEBUG:
            print('total_anchors', total_anchors)
            print('inds_inside', len(inds_inside))

        # keep only inside anchors
        anchors = all_anchors[inds_inside, :]

        anchors_list.append(anchors)
        inds_inside_list.append(inds_inside)

    # Concat anchors from each level
    anchors = np.concatenate(anchors_list)
    for i in range(1, len(inds_inside_list)):
        inds_inside_list[i] = inds_inside_list[i] + sum(anchors_num_list[:i])
    inds_inside = np.concatenate(inds_inside_list)
    total_anchors = sum(anchors_num_list)

    # label: 1 is positive, 0 is negative, -1 is dont care
    labels = np.empty((len(inds_inside),), dtype=np.float32)
    labels.fill(-1)

    if gt_boxes.size > 0:
        # overlap between the anchors and the gt boxes
        # overlaps (ex, gt)
        overlaps = bbox_overlaps(anchors.astype(np.float), gt_boxes.astype(np.float))
        argmax_overlaps = overlaps.argmax(axis=1)
        max_overlaps = overlaps[np.arange(len(inds_inside)), argmax_overlaps]
        gt_argmax_overlaps = overlaps.argmax(axis=0)
        gt_max_overlaps = overlaps[gt_argmax_overlaps, np.arange(overlaps.shape[1])]
        gt_argmax_overlaps = np.where(overlaps == gt_max_overlaps)[0]

        if not config.train.rpn_clobber_positives:
            # assign bg labels first so that positive labels can clobber them
            labels[max_overlaps < config.train.rpn_negative_overlap] = 0

        # fg label: for each gt, anchor with highest overlap
        labels[gt_argmax_overlaps] = 1

        # fg label: above threshold IoU
        labels[max_overlaps >= config.train.rpn_positive_overlap] = 1

        if config.train.rpn_clobber_positives:
            # assign bg labels last so that negative labels can clobber positives
            labels[max_overlaps < config.train.rpn_negative_overlap] = 0
    else:
        labels[:] = 0

    # subsample positive labels if we have too many
    num_fg = int(config.train.rpn_fg_fraction * config.train.rpn_batch_size)
    fg_inds = np.where(labels == 1)[0]
    if len(fg_inds) > num_fg:
        disable_inds = npr.choice(fg_inds, size=(len(fg_inds) - num_fg), replace=False)
        if DEBUG:
            disable_inds = fg_inds[:(len(fg_inds) - num_fg)]
        labels[disable_inds] = -1

    # subsample negative labels if we have too many
    num_bg = config.train.rpn_batch_size - np.sum(labels == 1)
    bg_inds = np.where(labels == 0)[0]
    if len(bg_inds) > num_bg:
        disable_inds = npr.choice(bg_inds, size=(len(bg_inds) - num_bg), replace=False)
        if DEBUG:
            disable_inds = bg_inds[:(len(bg_inds) - num_bg)]
        labels[disable_inds] = -1

    bbox_targets = np.zeros((len(inds_inside), 4), dtype=np.float32)
    if gt_boxes.size > 0:
        bbox_targets[:] = bbox_transform(anchors, gt_boxes[argmax_overlaps, :4])

    bbox_weights = np.zeros((len(inds_inside), 4), dtype=np.float32)
    bbox_weights[labels == 1, :] = np.array(config.train.rpn_bbox_weights)

    if DEBUG:
        _sums = bbox_targets[labels == 1, :].sum(axis=0)
        _squared_sums = (bbox_targets[labels == 1, :] ** 2).sum(axis=0)
        _counts = np.sum(labels == 1)
        means = _sums / (_counts + 1e-14)
        stds = np.sqrt(_squared_sums / _counts - means ** 2)
        print('means', means)
        print('stdevs', stds)
    # map up to original set of anchors
    labels = _unmap(labels, total_anchors, inds_inside, fill=-1)
    bbox_targets = _unmap(bbox_targets, total_anchors, inds_inside, fill=0)
    bbox_weights = _unmap(bbox_weights, total_anchors, inds_inside, fill=0)

    if DEBUG:
        if gt_boxes.size > 0:
            print('rpn: max max_overlaps', np.max(max_overlaps))
        print('rpn: num_positives', np.sum(labels == 1))
        print('rpn: num_negatives', np.sum(labels == 0))
        _fg_sum = np.sum(labels == 1)
        _bg_sum = np.sum(labels == 0)
        _count = 1
        print('rpn: num_positive avg', _fg_sum / _count)
        print('rpn: num_negative avg', _bg_sum / _count)

    # resahpe
    label_list = list()
    bbox_target_list = list()
    bbox_weight_list = list()
    anchors_num_range = [0] + anchors_num_list
    for i in range(len(feat_strides)):
        feat_height, feat_width = feat_infos[i]
        A = A_list[i]
        label = labels[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
        bbox_target = bbox_targets[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
        bbox_weight = bbox_weights[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]

        label = label.reshape((1, feat_height, feat_width, A)).transpose(0, 3, 1, 2)
        label = label.reshape((1, A * feat_height * feat_width))
        bbox_target = bbox_target.reshape((1, feat_height*feat_width, A * 4)).transpose(0, 2, 1)
        bbox_weight = bbox_weight.reshape((1, feat_height*feat_width, A * 4)).transpose((0, 2, 1))

        label_list.append(label)
        bbox_target_list.append(bbox_target)
        bbox_weight_list.append(bbox_weight)

    label_concat = np.concatenate(label_list, axis=1)
    bbox_target_concat = np.concatenate(bbox_target_list, axis=2)
    bbox_weight_concat = np.concatenate(bbox_weight_list, axis=2)

    label = {'label': label_concat,
             'bbox_target': bbox_target_concat,
             'bbox_weight': bbox_weight_concat}
    return label

def add_rpn_blobs(blobs, im_scales, roidb):
    """Add blobs needed training RPN-only and end-to-end Faster R-CNN models."""
    if config.network.has_fpn:
        # RPN applied to many feature levels, as in the FPN paper
        foas = []
        for field_stride in config.network.rpn_feat_stride:
            anchor_sizes = (config.network.anchor_scales[0] * field_stride,)
            anchor_aspect_ratios = config.network.anchor_ratios
            foa = get_field_of_anchors(
                field_stride, anchor_sizes, anchor_aspect_ratios
            )
            foas.append(foa)
        all_anchors = np.concatenate([f.field_of_anchors for f in foas])
    else:
        foa = get_field_of_anchors(
            config.network.rpn_feat_stride, np.array(
                config.network.anchor_scales) * config.network.rpn_feat_stride, config.network.anchor_ratios
        )
        all_anchors = foa.field_of_anchors

    for im_i, entry in enumerate(roidb):
        scale = im_scales[im_i]
        im_height = np.round(entry['height'] * scale)
        im_width = np.round(entry['width'] * scale)
        gt_inds = np.where(
            (entry['gt_classes'] > 0) & (entry['is_crowd'] == 0)
        )[0]
        crowd_gt_inds = np.where(
            (entry['gt_classes'] > 0) & (entry['is_crowd'] == 1)
        )[0]
        gt_rois = entry['boxes'][gt_inds, :] * scale
        crowd_gt_rois = entry['boxes'][crowd_gt_inds, :] * scale
        # TODO(rbg): gt_boxes is poorly named;
        # should be something like 'gt_rois_info'
        gt_boxes = np.zeros((len(gt_inds), 6), dtype=np.float32)
        gt_boxes[:, 0] = im_i  # batch inds
        gt_boxes[:, 1:5] = gt_rois
        gt_boxes[:, 5] = entry['gt_classes'][gt_inds]
        im_info = np.array([[im_height, im_width, scale]], dtype=np.float32)
        blobs['im_info'].append(im_info)

        # Add RPN targets
        if config.network.has_fpn:
            # RPN applied to many feature levels, as in the FPN paper
            rpn_blobs = _get_rpn_blobs(
                im_height, im_width, foas, all_anchors, gt_rois)
            for i, lvl in enumerate(config.network.rpn_feat_stride):
                for k, v in rpn_blobs[i].items():
                    blobs[k + '_fpn' + str(lvl)].append(v)
        else:
            # Classical RPN, applied to a single feature level
            rpn_blobs = _get_rpn_blobs(
                im_height, im_width, [foa], all_anchors, gt_rois)
            for k, v in rpn_blobs.items():
                blobs[k].append(v)

    for k, v in blobs.items():
        if isinstance(v, list) and len(v) > 0:
            blobs[k] = np.concatenate(v)

    # valid_keys = [
    #     'has_visible_keypoints', 'boxes', 'segms', 'seg_areas', 'gt_classes',
    #     'gt_overlaps', 'is_crowd', 'box_to_gt_ind_map', 'gt_keypoints'
    # ]
    valid_keys = [
        'boxes', 'segms', 'seg_areas', 'gt_classes',
        'gt_overlaps', 'is_crowd', 'box_to_gt_ind_map', 'image'
    ]
    minimal_roidb = [{} for _ in range(len(roidb))]
    for i, e in enumerate(roidb):
        for k in valid_keys:
            if k in e:
                minimal_roidb[i][k] = e[k]
    blobs['roidb'] = minimal_roidb

    # Always return valid=True, since RPN minibatches are valid by design
    return True

def _get_rpn_blobs(im_height, im_width, foas, all_anchors, gt_boxes):
    total_anchors = all_anchors.shape[0]
    straddle_thresh = config.train.rpn_straddle_thresh

    if straddle_thresh >= 0:
        # Only keep anchors inside the image by a margin of straddle_thresh
        # Set TRAIN.RPN_STRADDLE_THRESH to -1 (or a large value) to keep all
        # anchors
        inds_inside = np.where(
            (all_anchors[:, 0] >= -straddle_thresh) &
            (all_anchors[:, 1] >= -straddle_thresh) &
            (all_anchors[:, 2] < im_width + straddle_thresh) &
            (all_anchors[:, 3] < im_height + straddle_thresh)
        )[0]
        # keep only inside anchors
        anchors = all_anchors[inds_inside, :]
    else:
        inds_inside = np.arange(all_anchors.shape[0])
        anchors = all_anchors
    num_inside = len(inds_inside)

    # Compute anchor labels:
    # label=1 is positive, 0 is negative, -1 is don't care (ignore)
    labels = np.empty((num_inside,), dtype=np.int32)
    labels.fill(-1)
    if len(gt_boxes) > 0:
        # Compute overlaps between the anchors and the gt boxes overlaps
        anchor_by_gt_overlap = bbox_overlaps(anchors, gt_boxes)
        # Map from anchor to gt box that has highest overlap
        anchor_to_gt_argmax = anchor_by_gt_overlap.argmax(axis=1)
        # For each anchor, amount of overlap with most overlapping gt box
        anchor_to_gt_max = anchor_by_gt_overlap[np.arange(num_inside),
                                                anchor_to_gt_argmax]

        # Map from gt box to an anchor that has highest overlap
        gt_to_anchor_argmax = anchor_by_gt_overlap.argmax(axis=0)
        # For each gt box, amount of overlap with most overlapping anchor
        gt_to_anchor_max = anchor_by_gt_overlap[
            gt_to_anchor_argmax,
            np.arange(anchor_by_gt_overlap.shape[1])
        ]
        # Find all anchors that share the max overlap amount
        # (this includes many ties)
        anchors_with_max_overlap = np.where(
            anchor_by_gt_overlap == gt_to_anchor_max
        )[0]

        # Fg label: for each gt use anchors with highest overlap
        # (including ties)
        labels[anchors_with_max_overlap] = 1
        # Fg label: above threshold IOU
        labels[anchor_to_gt_max >= config.train.rpn_positive_overlap] = 1

    # subsample positive labels if we have too many
    num_fg = int(config.train.rpn_fg_fraction * config.train.rpn_batch_size)
    fg_inds = np.where(labels == 1)[0]
    if len(fg_inds) > num_fg:
        disable_inds = np.random.choice(
            fg_inds, size=(len(fg_inds) - num_fg), replace=False
        )
        # print('assign_anchor debug use first 128 fg')
        # labels[fg_inds[-(len(fg_inds) - num_fg):]] = -1
        labels[disable_inds] = -1
    fg_inds = np.where(labels == 1)[0]

    # subsample negative labels if we have too many
    # (samples with replacement, but since the set of bg inds is large most
    # samples will not have repeats)
    num_bg = config.train.rpn_batch_size - np.sum(labels == 1)
    bg_inds = np.where(anchor_to_gt_max < config.train.rpn_negative_overlap)[0]
    if len(bg_inds) > num_bg:
        # enable_inds = bg_inds[np.random.randint(len(bg_inds), size=num_bg)]
        enable_inds = bg_inds[np.random.choice(len(bg_inds), num_bg, replace=False)]
        # print('assign_anchor debug use first 128 bg')
        # labels[bg_inds[:num_bg]] = 0
        labels[enable_inds] = 0
    bg_inds = np.where(labels == 0)[0]

    bbox_targets = np.zeros((num_inside, 4), dtype=np.float32)
    bbox_targets[fg_inds, :] = compute_targets(
        anchors[fg_inds, :], gt_boxes[anchor_to_gt_argmax[fg_inds], :]
    )

    # Bbox regression loss has the form:
    #   loss(x) = weight_outside * L(weight_inside * x)
    # Inside weights allow us to set zero loss on an element-wise basis
    # Bbox regression is only trained on positive examples so we set their
    # weights to 1.0 (or otherwise if config is different) and 0 otherwise
    bbox_inside_weights = np.zeros((num_inside, 4), dtype=np.float32)
    bbox_inside_weights[labels == 1, :] = (1.0, 1.0, 1.0, 1.0)

    # The bbox regression loss only averages by the number of images in the
    # mini-batch, whereas we need to average by the total number of example
    # anchors selected
    # Outside weights are used to scale each element-wise loss so the final
    # average over the mini-batch is correct
    bbox_outside_weights = np.zeros((num_inside, 4), dtype=np.float32)
    # uniform weighting of examples (given non-uniform sampling)
    num_examples = np.sum(labels >= 0)
    bbox_outside_weights[labels == 1, :] = 1.0 / num_examples
    bbox_outside_weights[labels == 0, :] = 1.0 / num_examples

    # Map up to original set of anchors
    labels = unmap(labels, total_anchors, inds_inside, fill=-1)
    bbox_targets = unmap(
        bbox_targets, total_anchors, inds_inside, fill=0
    )
    bbox_inside_weights = unmap(
        bbox_inside_weights, total_anchors, inds_inside, fill=0
    )
    bbox_outside_weights = unmap(
        bbox_outside_weights, total_anchors, inds_inside, fill=0
    )

    # Split the generated labels, etc. into labels per each field of anchors
    blobs_out = []
    start_idx = 0
    for foa in foas:
        H = foa.field_size
        W = foa.field_size
        A = foa.num_cell_anchors
        end_idx = start_idx + H * W * A
        _labels = labels[start_idx:end_idx]
        _bbox_targets = bbox_targets[start_idx:end_idx, :]
        _bbox_inside_weights = bbox_inside_weights[start_idx:end_idx, :]
        _bbox_outside_weights = bbox_outside_weights[start_idx:end_idx, :]
        start_idx = end_idx

        # labels output with shape (1, A, height, width)
        _labels = _labels.reshape((1, H, W, A)).transpose(0, 3, 1, 2)
        # bbox_targets output with shape (1, 4 * A, height, width)
        _bbox_targets = _bbox_targets.reshape(
            (1, H, W, A * 4)).transpose(0, 3, 1, 2)
        # bbox_inside_weights output with shape (1, 4 * A, height, width)
        _bbox_inside_weights = _bbox_inside_weights.reshape(
            (1, H, W, A * 4)).transpose(0, 3, 1, 2)
        # bbox_outside_weights output with shape (1, 4 * A, height, width)
        _bbox_outside_weights = _bbox_outside_weights.reshape(
            (1, H, W, A * 4)).transpose(0, 3, 1, 2)
        blobs_out.append(
            dict(
                rpn_labels_int32_wide=_labels,
                rpn_bbox_targets_wide=_bbox_targets,
                rpn_bbox_inside_weights_wide=_bbox_inside_weights,
                rpn_bbox_outside_weights_wide=_bbox_outside_weights
            )
        )
    return blobs_out[0] if len(blobs_out) == 1 else blobs_out


from .solver.neuron_solver import neuron_solver
from .solver.aux import snap_to_nearest
import dolfin
from ..mesh.mesh_utils import EMIEntityMap, load_h5_mesh
from .solver.probing import get_geom_centers, Probe
import numpy as np
from copy import copy, deepcopy
from pathlib import Path
import time

_default_problem_parameters = {
    'neurons':  # List or dict. If list, length must be the same of number of neurons in mesh. If dict, the same
                # params are used for all neurons in the mesh
                {
                 'cond_int': 7,           # float: Intracellular conductivity in mS/cm^2
                 'Cm': 1.0,               # float: Membrane capacitance in uF/um^2
                 'models': {},            # dict: Models for neuron domains. Default: {'dendrite': 'pas', 'soma': 'hh',
                                          #                                             'axon': 'hh'}
                 'model_args': {},        # dict of tuples: Overwrite model default arguments.
                                          # E.g. ('dendrite', 'g_L'), ('soma', g_Na)'
                 'stimulation': {'type': 'syn',           # str: Stimulation type ('syn', 'step', 'pulse')
                                 'start_time': 0.01,      # float: Start of stimulation in ms
                                 'stop_time': 1.0,        # float: Stop of stimulation in ms (it type is 'pulse')
                                 'syn_weight': 10.0,      # float: Synaptic weight in in mS/cm^2 (if 'type' is 'syn')
                                 'stim_current': 10,      # float: Stimulation current in nA (if 'type' is 'step'/'pulse')
                                 'position': [0, 0, 150], # array: 3D position of stimulation point in um.
                                 'length': 20,            # float: length of stimulated portion in um.
                                 'radius': 5,             # float: radius of stimulated area
                 }
    },
    'ext':     {
                 'cond_ext': 3,           # float: Extracellular conductivity: mS/cm^2
                 'insulated_bcs': []      # list: Insulated BC for bounding box. It can be: 'max_x', 'min_x', etc.
    },
    'probe':   {
                 'stimulated_sites': None,  # List or tuple: Stimulated electrodes (e.g. [contact_3, contact_22]
                 'type': None,            # str: Stimulation type ('step', 'pulse')
                 'start_time': 0.1,       # float: Start of stimulation in ms
                 'stop_time': 1.0,        # float: Stop of stimulation in ms (it type is 'pulse')
                 'current': 0             # float: Stimulation current in mA. If list, each value correspond to a stimulated
                                          #        site
    },
    'solver':  {
                 'dt_fem': 0.01,          # float: dt for fem solver in ms
                 'dt_ode': 0.01,          # float: dt for ode solver in ms
                 'sim_duration': 5,      # float: duration od the simulation in ms
    }
}

def get_default_emi_params():
    return deepcopy(_default_problem_parameters)

def simulate_emi(mesh_folder, problem_params=None, u_probe_locations=None,
                 i_probe_locations=None, v_probe_locations=None,
                 save_simulation_output=True, save_format='pvd', save_folder=None, verbose=False):
    '''
    Simulates the 3d-3d EMI solution given a neuronmi-generated mesh (containing at least one neuron and optionally
    one probe).
    The simulation results (currents and potentials) are saved in the 'emi_sim' folder inside the mesh_folder.

    Parameters
    ----------
    mesh_folder: str or Path
        The path to the neuronmi-generated mesh containing .h5 and .json
    problem_params: dict
        Dictionary with simulation parameters. To retrieve default parameters, run **neuronmi.get_default_emi_params()**
    u_probe_locations: np.array
        Array of 3D points in um to measure potential
    i_probe_locations: np.array
        Array of 3D points in um to measure currents (if not on neuron surface the point is snapped
        to closest neuron point)
    v_probe_locations: np.array
        Array of 3D points in um to measure membrane potentials (if not on neuron surface the point is snapped
        to closest neuron point)
    save_simulation_output: bool
        If True, simulation output (u: potential, i: current, v: membrane potential) is saved
    save_format: str
        'pvd' or 'xdmf'
    save_folder: str
        Path to save folder. If None, an 'emi_simulation' folder is created in the mesh folder.
    '''
    scale_factor = 1e-4
    mesh_folder = Path(mesh_folder)

    mesh_h5 = [f for f in mesh_folder.iterdir() if f.suffix == '.h5']
    mesh_json = [f for f in mesh_folder.iterdir() if f.suffix == '.json']

    if len(mesh_h5) != 1:
        raise ValueError('No or more than one .h5 mesh file found in %s' % mesh_folder)
    else:
        mesh_h5_path = mesh_h5[0]

    if len(mesh_json) != 1:
        raise ValueError('No or more than one .json mesh file found in %s' % mesh_folder)
    else:
        mesh_json_path = mesh_json[0]

    with mesh_json_path.open() as json_fp:
        emi_map = EMIEntityMap(json_fp=json_fp)

    if problem_params is None:
        problem_params = _default_problem_parameters

    if save_folder is None:
        save_folder = mesh_folder / 'emi_simulation'
    else:
        save_folder = Path(save_folder)

    if save_simulation_output:
        if save_format == 'pvd':
            i_out = dolfin.File(str(save_folder / 'i.pvd'))
            u_out = dolfin.File(str(save_folder / 'u.pvd'))
            v_out = dolfin.File(str(save_folder / 'v.pvd'))
        elif save_format == 'xdmf':
            i_out = dolfin.XDMFFile(str(save_folder / 'i.xdmf'))
            u_out = dolfin.XDMFFile(str(save_folder / 'u.xdmf'))
            v_out = dolfin.XDMFFile(str(save_folder / 'v.xdmf'))
        else:
            raise AttributeError("'save_format' can be 'pvd' or 'xdmf'")

    dolfin.parameters['allow_extrapolation'] = True

    t_start = time.time()
    u_record = []
    i_record = []
    v_record = []
    i_proxy = None
    v_proxy = None

    for (t, u, i, v) in neuron_solver(mesh_h5_path, emi_map, problem_params, scale_factor, verbose):
        if save_simulation_output:
            if save_format == 'pvd':
                i_out << i, t
                u_out << u, t
                v_out << v, t
            elif save_format == 'xdmf':
                i_out.write(i, float(t))
                u_out.write(u, float(t))
                v_out.write(v, float(t))

        if i_proxy is None:
            i_proxy = snap_to_nearest(i)

        if v_proxy is None:
            v_proxy = snap_to_nearest(v)

        if u_probe_locations is not None:
            u_probe_t = np.zeros(len(u_probe_locations))
            for idx, p in enumerate(u_probe_locations):
                u_probe_t[idx] = u(np.array(p) * scale_factor)
            u_record.append(u_probe_t)

        #TODO now I_proxy is in mAcm-2 --> multiply by facet area to obtain current!
        if i_probe_locations is not None:
            i_probe_t = np.zeros(len(i_probe_locations))
            for idx, p in enumerate(i_probe_locations):
                i_probe_t[idx] = i_proxy(np.array(p) * scale_factor)
            i_record.append(i_probe_t)

        if v_probe_locations is not None:
            v_probe_t = np.zeros(len(v_probe_locations))
            for idx, p in enumerate(v_probe_locations):
                v_probe_t[idx] = v_proxy(np.array(p) * scale_factor)
            v_record.append(v_probe_t)

    if u_probe_locations is not None:
        u_record = np.array(u_record)
    if i_probe_locations is not None:
        i_record = np.array(i_record)
    if v_probe_locations is not None:
        v_record = np.array(v_record)

    if save_simulation_output:
        print 'Results saved in ' + str(save_folder)
    print 'Elapsed time: ' + str(time.time() - t_start)

    return u_record, i_record, v_record


import numpy as np
import mido, os, pickle, yaml, hgtk, argparse, math
from tqdm import tqdm
from argparse import Namespace  # for type

from Audio import Audio_Prep, Mel_Generate
from yin import pitch_calc
from Arg_Parser import Recursive_Parse

def Decompose(syllable):
    onset, nucleus, coda = hgtk.letter.decompose(syllable)
    coda += '_'

    return onset, nucleus, coda

def Pattern_Generate(
    hyper_paramters: Namespace,
    dataset_path: str
    ):
    min_Duration, max_Duration = math.inf, -math.inf
    min_Note, max_Note = math.inf, -math.inf

    paths = []
    for root, _, files in os.walk(dataset_path):
        for file in sorted(files):
            if os.path.splitext(file)[1] != '.wav':
                continue
            wav_Path = os.path.join(root, file).replace('\\', '/')
            midi_Path = wav_Path.replace('vox.wav', 'midi.mid')
            paths.append((wav_Path, midi_Path))

    for index, (wav_Path, midi_Path) in enumerate(paths):
        mid = mido.MidiFile(midi_Path, charset='CP949')
        music = []
        current_Lyric = ''
        current_Note = None
        current_Time = 0.0

        # Note on 쉼표
        # From Lyric to message before note on: real note
        for message in list(mid):
            if message.type == 'note_on':
                if message.time < 0.1:
                    current_Time += message.time
                    if current_Lyric in ['J', 'H', None]:
                        music.append((current_Time, '<X>', 0))
                    else:
                        music.append((current_Time, current_Lyric, current_Note))
                else:
                    if not current_Lyric in ['J', 'H', None]:
                        music.append((current_Time, current_Lyric, current_Note))
                    else:
                        message.time += current_Time
                    music.append((message.time, '<X>', 0))
                current_Time = 0.0
                current_Lyric = ''
                current_Note = None
            elif message.type == 'lyrics':
                current_Lyric = message.text.strip()
                current_Time += message.time
            elif message.type == 'note_off':
                current_Note = message.note
                current_Time += message.time
                if current_Lyric == 'H':    # it is temp.
                    break
            else:
                current_Time += message.time

        if current_Lyric in ['J', 'H']:
            if music[-1][1] == '<X>':
                music[-1] = (music[-1][0] + current_Time, music[-1][1], music[-1][2])
            else:
                music.append((current_Time, '<X>', 0))
        else:
            music.append((current_Time, current_Lyric, current_Note))
        music = music[1:]

        audio = Audio_Prep(wav_Path, hyper_paramters.Sound.Sample_Rate)
        if music[0][1] == '<X>':
            audio = audio[int(music[0][0] * hyper_paramters.Sound.Sample_Rate):]
            music = music[1:]
        if music[-1][1] == '<X>':
            audio = audio[:-int(music[-1][0] * hyper_paramters.Sound.Sample_Rate)]
            music = music[:-1]

        previous_Used = 0
        absolute_Position = 0
        mel_Based = []
        for x in music:
            duration = int(x[0] * hyper_paramters.Sound.Sample_Rate) + previous_Used
            previous_Used = duration % hyper_paramters.Sound.Frame_Shift
            duration = duration // hyper_paramters.Sound.Frame_Shift
            if x[1] == '<X>':
                mel_Based.append((absolute_Position, duration, x[1], x[2]))
                absolute_Position += duration
            else:
                lyrics = Decompose(x[1])
                mel_Based.append((absolute_Position, 2, lyrics[0], x[2]))   # Onset
                absolute_Position += 2
                mel_Based.append((absolute_Position, duration - 4, lyrics[1], x[2]))   # Onset
                absolute_Position += duration - 4
                mel_Based.append((absolute_Position, 2, lyrics[2], x[2]))   # Onset
                absolute_Position += 2
        music = mel_Based

        mel = Mel_Generate(
            audio,
            sample_rate= hyper_paramters.Sound.Sample_Rate,
            num_mel= hyper_paramters.Sound.Mel_Dim,
            num_frequency= hyper_paramters.Sound.Spectrogram_Dim,
            window_length= hyper_paramters.Sound.Frame_Length,
            hop_length= hyper_paramters.Sound.Frame_Shift,
            mel_fmin= hyper_paramters.Sound.Mel_F_Min,
            mel_fmax= hyper_paramters.Sound.Mel_F_Max,
            max_abs_value= hyper_paramters.Sound.Max_Abs_Mel
            )[:absolute_Position]

        pitch = pitch_calc(
            sig= audio,
            sr= hyper_paramters.Sound.Sample_Rate,
            w_len= hyper_paramters.Sound.Frame_Length,
            w_step= hyper_paramters.Sound.Frame_Shift,
            f0_min= hyper_paramters.Sound.F0_Min,
            f0_max= hyper_paramters.Sound.F0_Max,
            confidence_threshold= hyper_paramters.Sound.Confidence_Threshold,
            gaussian_smoothing_sigma = hyper_paramters.Sound.Gaussian_Smoothing_Sigma
            )[:absolute_Position] / hyper_paramters.Sound.F0_Max

        silence = np.where(np.mean(mel, axis=1) < -3.5, np.zeros_like(np.mean(mel, axis=1)), np.ones_like(np.mean(mel, axis=1)))

        pattern_Index = 0
        for start_Index in tqdm(range(len(music)), desc= os.path.basename(wav_Path)):
            for end_Index in range(start_Index + 1, len(music), 5):
                music_Sample = music[start_Index:end_Index]
                sample_Length = music_Sample[-1][0] + music_Sample[-1][1] - music_Sample[0][0]
                if sample_Length < hyper_paramters.Min_Duration:
                    continue
                elif sample_Length > hyper_paramters.Max_Duration:
                    break

                audio_Sample = audio[music_Sample[0][0] * hyper_paramters.Sound.Frame_Shift:(music_Sample[-1][0] + music_Sample[-1][1]) * hyper_paramters.Sound.Frame_Shift]
                mel_Sample = mel[music_Sample[0][0]:music_Sample[-1][0] + music_Sample[-1][1]]
                silence_Sample = silence[music_Sample[0][0]:music_Sample[-1][0] + music_Sample[-1][1]]
                pitch_Sample = pitch[music_Sample[0][0]:music_Sample[-1][0] + music_Sample[-1][1]]

                _, duration_Sample, text_Sample, Note_Sample = zip(*music_Sample)

                pattern = {
                    'Audio': audio_Sample.astype(np.float32),
                    'Mel': mel_Sample.astype(np.float32),
                    'Silence': silence_Sample.astype(np.uint8),
                    'Pitch': pitch_Sample.astype(np.float32),
                    'Duration': duration_Sample,
                    'Text': text_Sample,
                    'Note': Note_Sample,
                    'Singer': 'Female_0',
                    'Dataset': 'NAMS',
                    }

                pattern_Path = os.path.join(
                    hyper_paramters.Train.Train_Pattern.Path if np.random.rand() > 0.001 else hyper_paramters.Train.Eval_Pattern.Path,
                    'NAMS',
                    '{:03d}'.format(index),
                    'NAMS.S_{:03d}.P_{:05d}.pickle'.format(index, pattern_Index)
                    ).replace('\\', '/')
                os.makedirs(os.path.dirname(pattern_Path), exist_ok= True)
                pickle.dump(
                    pattern,
                    open(pattern_Path, 'wb'),
                    protocol=4
                    )
                pattern_Index += 1

                min_Duration, max_Duration = min(sample_Length, min_Duration), max(sample_Length, max_Duration)
        min_Note, max_Note = min(list(zip(*music))[3] + (min_Note,)), max(list(zip(*music))[3] + (max_Note,))

    print('Duration range: {} - {}'.format(min_Duration, max_Duration))
    print('Note range: {} - {}'.format(min_Note, max_Note))

def Token_Dict_Generate(hyper_parameters: Namespace):
    tokens = \
        list(hgtk.letter.CHO) + \
        list(hgtk.letter.JOONG) + \
        ['{}_'.format(x) for x in hgtk.letter.JONG]

    os.makedirs(os.path.dirname(hyper_parameters.Token_Path), exist_ok= True)
    yaml.dump(
        {token: index for index, token in enumerate(['<S>', '<E>', '<X>'] + sorted(tokens))},
        open(hyper_parameters.Token_Path, 'w')
        )

def Metadata_Generate(
    hyper_parameters: Namespace,
    eval: bool= False
    ):
    pattern_Path = hyper_parameters.Train.Eval_Pattern.Path if eval else hyper_parameters.Train.Train_Pattern.Path
    metadata_File = hyper_parameters.Train.Eval_Pattern.Metadata_File if eval else hyper_parameters.Train.Train_Pattern.Metadata_File

    new_Metadata_Dict = {
        'Spectrogram_Dim': hyper_parameters.Sound.Spectrogram_Dim,
        'Mel_Dim': hyper_parameters.Sound.Mel_Dim,
        'Frame_Shift': hyper_parameters.Sound.Frame_Shift,
        'Frame_Length': hyper_parameters.Sound.Frame_Length,
        'Sample_Rate': hyper_parameters.Sound.Sample_Rate,
        'Max_Abs_Mel': hyper_parameters.Sound.Max_Abs_Mel,
        'Mel_F_Min': hyper_parameters.Sound.Mel_F_Min,
        'Mel_F_Max': hyper_parameters.Sound.Mel_F_Max,
        'File_List': [],
        'Audio_Length_Dict': {},
        'Mel_Length_Dict': {},
        'Music_Length_Dict': {},
        }

    files_TQDM = tqdm(
        total= sum([len(files) for root, _, files in os.walk(pattern_Path)]),
        desc= 'Eval_Pattern' if eval else 'Train_Pattern'
        )

    for root, _, files in os.walk(pattern_Path):
        for file in files:
            with open(os.path.join(root, file).replace("\\", "/"), "rb") as f:
                pattern_Dict = pickle.load(f)
            file = os.path.join(root, file).replace("\\", "/").replace(pattern_Path, '').lstrip('/')
            try:
                if not all([
                    key in pattern_Dict.keys()
                    for key in ('Audio', 'Mel', 'Silence', 'Pitch', 'Duration', 'Text', 'Note', 'Singer', 'Dataset')
                    ]):
                    continue
                new_Metadata_Dict['Audio_Length_Dict'][file] = pattern_Dict['Audio'].shape[0]
                new_Metadata_Dict['Mel_Length_Dict'][file] = pattern_Dict['Mel'].shape[0]
                new_Metadata_Dict['Music_Length_Dict'][file] = len(pattern_Dict['Duration'])
                new_Metadata_Dict['File_List'].append(file)
            except:
                print('File \'{}\' is not correct pattern file. This file is ignored.'.format(file))
            files_TQDM.update(1)

    with open(os.path.join(pattern_Path, metadata_File.upper()).replace("\\", "/"), 'wb') as f:
        pickle.dump(new_Metadata_Dict, f, protocol= 4)

    print('Metadata generate done.')

if __name__ == "__main__":
    argParser = argparse.ArgumentParser()
    argParser.add_argument("-d", "--dataset_path", required= True)
    argParser.add_argument("-hp", "--hyper_paramters", required= True)
    args = argParser.parse_args()

    hp = Recursive_Parse(yaml.load(
        open(args.hyper_paramters, encoding='utf-8'),
        Loader=yaml.Loader
        ))

    Token_Dict_Generate(hyper_parameters= hp)
    Pattern_Generate(hyper_paramters= hp, dataset_path= args.dataset_path)
    Metadata_Generate(hp, False)
    Metadata_Generate(hp, True)

    # python Pattern_Generator.py -hp Hyper_Parameters.yaml -d E:/Kor_Music_Confidential


import numpy as np

from sitator.util import PBCCalculator

# From https://github.com/tqdm/tqdm/issues/506#issuecomment-373126698
import sys
try:
    ipy_str = str(type(get_ipython()))
    if 'zmqshell' in ipy_str:
        from tqdm import tqdm_notebook as tqdm
    if 'terminal' in ipy_str:
        from tqdm import tqdm
except:
    if sys.stderr.isatty():
        from tqdm import tqdm
    else:
        def tqdm(iterable, **kwargs):
            return iterable

import importlib
import tempfile

import helpers
from sitator import SiteNetwork, SiteTrajectory

from functools import wraps
def analysis_result(func):
    @property
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        if not self._has_run:
            raise ValueError("This LandmarkAnalysis hasn't been run yet.")
        return func(self, *args, **kwargs)
    return wrapper

class LandmarkAnalysis(object):
    """Track a mobile species through a fixed lattice using landmark vectors."""

    def __init__(self,
                 clustering_algorithm = 'dotprod',
                 clustering_params = {},
                 cutoff_midpoint = 1.5,
                 cutoff_steepness = 30,
                 minimum_site_occupancy = 0.01,
                 peak_evening = 'none',
                 weighted_site_positions = True,
                 check_for_zero_landmarks = True,
                 static_movement_threshold = 1.0,
                 dynamic_lattice_mapping = False,
                 relaxed_lattice_checks = False,
                 max_mobile_per_site = 1,
                 force_no_memmap = False,
                 verbose = True):
        """
        :param double cutoff_center: The midpoint for the logistic function used
            as the landmark cutoff function. (unitless)
        :param double cutoff_steepness: Steepness of the logistic cutoff function.
        :param double minimum_site_occupancy = 0.1: Minimum occupancy (% of time occupied)
            for a site to qualify as such.
        :param dict clustering_params: Parameters for the chosen clustering_algorithm
        :param str peak_evening: Whether and what kind of peak "evening" to apply;
            that is, processing that makes all large peaks in the landmark vector
            more similar in magnitude. This can help in site clustering.

            Valid options: 'none', 'clip'
        :param bool weighted_site_positions: When computing site positions, whether
            to weight the average by assignment confidence.
        :param bool check_for_zero_landmarks: Whether to check for and raise exceptions
            when all-zero landmark vectors are computed.
        :param float static_movement_threshold: (Angstrom) the maximum allowed
            distance between an instantanous static atom position and it's ideal position.
        :param bool dynamic_lattice_mapping: Whether to dynamically decide each
            frame which static atom represents each average lattice position;
            this allows the LandmarkAnalysis to deal with, say, a rare exchage of
            two static atoms that does not change the structure of the lattice.

            It does NOT allow LandmarkAnalysis to deal with lattices whose structures
            actually change over the course of the trajectory.

            In certain cases this is better delt with by MergeSitesByDynamics.
        :param int max_mobile_per_site: The maximum number of mobile atoms that can
            be assigned to a single site without throwing an error. Regardless of the
            value, assignments of more than one mobile atom to a single site will
            be recorded and reported.

            Setting this to 2 can be necessary for very diffusive, liquid-like
            materials at high temperatures.

            Statistics related to this are reported in self.avg_mobile_per_site
            and self.n_multiple_assignments.
        :param bool force_no_memmap: if True, landmark vectors will be stored only in memory.
            Only useful if access to landmark vectors after the analysis has run is desired.
        :param bool verbose: If `True`, progress bars and messages will be printed to stdout.
        """

        self._cutoff_midpoint = cutoff_midpoint
        self._cutoff_steepness = cutoff_steepness
        self._minimum_site_occupancy = minimum_site_occupancy

        self._cluster_algo = clustering_algorithm
        self._clustering_params = clustering_params

        if not peak_evening in ['none', 'clip']:
          raise ValueError("Invalid value `%s` for peak_evening" % peak_evening)
        self._peak_evening = peak_evening

        self.verbose = verbose
        self.check_for_zero_landmarks = check_for_zero_landmarks
        self.weighted_site_positions = weighted_site_positions
        self.dynamic_lattice_mapping = dynamic_lattice_mapping
        self.relaxed_lattice_checks = relaxed_lattice_checks

        self._landmark_vectors = None
        self._landmark_dimension = None

        self.static_movement_threshold = static_movement_threshold
        self.max_mobile_per_site = max_mobile_per_site

        self.force_no_memmap = force_no_memmap

        self._has_run = False

    @property
    def cutoff(self):
      return self._cutoff

    @analysis_result
    def landmark_vectors(self):
        view = self._landmark_vectors[:]
        view.flags.writeable = False
        return view

    @analysis_result
    def landmark_dimension(self):
        return self._landmark_dimension

    def run(self, sn, frames):
        """Run the landmark analysis.

        The input SiteNetwork is a network of predicted sites; it's sites will
        be used as the "basis" for the landmark vectors.

        Takes a SiteNetwork and returns a SiteTrajectory.
        """
        assert isinstance(sn, SiteNetwork)

        if self._has_run:
            raise ValueError("Cannot rerun LandmarkAnalysis!")

        if frames.shape[1:] != (sn.n_total, 3):
            raise ValueError("Wrong shape %s for frames." % frames.shape)

        if sn.vertices is None:
            raise ValueError("Input SiteNetwork must have vertices")

        n_frames = len(frames)

        if self.verbose:
            print "--- Running Landmark Analysis ---"

        # Create PBCCalculator
        self._pbcc = PBCCalculator(sn.structure.cell)

        # -- Step 1: Compute site-to-vertex distances
        self._landmark_dimension = sn.n_sites

        longest_vert_set = np.max([len(v) for v in sn.vertices])
        verts_np = np.array([v + [-1] * (longest_vert_set - len(v)) for v in sn.vertices])
        site_vert_dists = np.empty(shape = verts_np.shape, dtype = np.float)
        site_vert_dists.fill(np.nan)

        for i, polyhedron in enumerate(sn.vertices):
            verts_poses = sn.static_structure.get_positions()[polyhedron]
            dists = self._pbcc.distances(sn.centers[i], verts_poses)
            site_vert_dists[i, :len(polyhedron)] = dists

        # -- Step 2: Compute landmark vectors
        if self.verbose: print "  - computing landmark vectors -"
        # Compute landmark vectors

        # The dimension of one landmark vector is the number of Voronoi regions
        shape = (n_frames * sn.n_mobile, self._landmark_dimension)

        with tempfile.NamedTemporaryFile() as mmap_backing:
            if self.force_no_memmap:
                self._landmark_vectors = np.empty(shape = shape, dtype = np.float)
            else:
                self._landmark_vectors = np.memmap(mmap_backing.name,
                                                   mode = 'w+',
                                                   dtype = np.float,
                                                   shape = shape)

            helpers._fill_landmark_vectors(self, sn, verts_np, site_vert_dists,
                                            frames, check_for_zeros = self.check_for_zero_landmarks,
                                            tqdm = tqdm)

            # -- Step 3: Cluster landmark vectors
            if self.verbose: print "  - clustering landmark vectors -"
            #  - Preprocess -
            self._do_peak_evening()

            #  - Cluster -
            cluster_func = importlib.import_module("..cluster." + self._cluster_algo, package = __name__).do_landmark_clustering

            cluster_counts, lmk_lbls, lmk_confs = \
                cluster_func(self._landmark_vectors,
                             clustering_params = self._clustering_params,
                             min_samples = self._minimum_site_occupancy / float(sn.n_mobile),
                             verbose = self.verbose)

        if self.verbose:
            print "    Failed to assign %i%% of mobile particle positions to sites." % (100.0 * np.sum(lmk_lbls < 0) / float(len(lmk_lbls)))

        # reshape lables and confidences
        lmk_lbls.shape = (n_frames, sn.n_mobile)
        lmk_confs.shape = (n_frames, sn.n_mobile)

        n_sites = len(cluster_counts)

        if n_sites < sn.n_mobile:
            raise ValueError("There are %i mobile particles, but only identified %i sites. Check clustering_params." % (sn.n_mobile, n_sites))

        if self.verbose:
            print "    Identified %i sites with assignment counts %s" % (n_sites, cluster_counts)

        # Check that multiple particles are never assigned to one site at the
        # same time, cause that would be wrong.
        n_more_than_ones = 0
        avg_mobile_per_site = 0
        divisor = 0
        for frame_i, site_frame in enumerate(lmk_lbls):
            _, counts = np.unique(site_frame[site_frame >= 0], return_counts = True)
            count_msk = counts > self.max_mobile_per_site
            if np.any(count_msk):
                raise ValueError("%i mobile particles were assigned to only %i site(s) (%s) at frame %i." % (np.sum(counts[count_msk]), np.sum(count_msk), np.where(count_msk)[0], frame_i))
            n_more_than_ones += np.sum(counts > 1)
            avg_mobile_per_site += np.sum(counts)
            divisor += len(counts)

        self.n_multiple_assignments = n_more_than_ones
        self.avg_mobile_per_site = avg_mobile_per_site / float(divisor)

        # -- Do output
        # - Compute site centers
        site_centers = np.empty(shape = (n_sites, 3), dtype = frames.dtype)

        for site in xrange(n_sites):
            mask = lmk_lbls == site
            pts = frames[:, sn.mobile_mask][mask]
            if self.weighted_site_positions:
                site_centers[site] = self._pbcc.average(pts, weights = lmk_confs[mask])
            else:
                site_centers[site] = self._pbcc.average(pts)

        # Build output obejcts
        out_sn = sn.copy()

        out_sn.centers = site_centers
        assert out_sn.vertices is None

        out_st = SiteTrajectory(out_sn, lmk_lbls, lmk_confs)
        out_st.set_real_traj(frames)
        self._has_run = True

        return out_st

    # -------- "private" methods --------

    def _do_peak_evening(self):
      if self._peak_evening == 'none':
          return
      elif self._peak_evening == 'clip':
          lvec_peaks = np.max(self._landmark_vectors, axis = 1)
          # Clip all peaks to the lowest "normal" (stdev.) peak
          lvec_clip = np.mean(lvec_peaks) - np.std(lvec_peaks)
          # Do the clipping
          self._landmark_vectors[self._landmark_vectors > lvec_clip] = lvec_clip


import datetime
import os
import pygame
import numpy
import torch

#sys.path.insert(0, "muzero-general/Tanks_new")
#sys.path.insert(0, "muzero-general/games")
from .settings import Settings
from Tanks_new.game_func import GameFunc
from .abstract_game import AbstractGame

class MuZeroConfig:
    def __init__(self):
        # More information is available here: https://github.com/werner-duvaud/muzero-general/wiki/Hyperparameter-Optimization

        self.seed = 0  # Seed for numpy, torch and the game
        self.max_num_gpus = 1  # Fix the maximum number of GPUs to use. It's usually faster to use a single GPU (set it to 1) if it has enough memory. None will use every GPUs available

        ### Game
        self.observation_shape = (3, 96, 96)  # Dimensions of the game observation, must be 3D (channel, height, width). For a 1D array, please reshape it to (1, 1, length of array)
        self.action_space = list(range(4))  # Fixed list of all possible actions. You should only edit the length
        self.players = list(range(1))  # List of players. You should only edit the length
        self.stacked_observations = 0  # Number of previous observations and previous actions to add to the current observation

        # Evaluate
        self.muzero_player = 0  # Turn Muzero begins to play (0: MuZero plays first, 1: MuZero plays second)
        self.opponent = None  # Hard coded agent that MuZero faces to assess his progress in multiplayer games. It doesn't influence training. None, "random" or "expert" if implemented in the Game class

        ### Self-Play
        self.num_workers = 1  # Number of simultaneous threads/workers self-playing to feed the replay buffer
        self.selfplay_on_gpu = False
        self.max_moves = 200  # Maximum number of moves if game is not finished before
        self.num_simulations = 20  # Number of future moves self-simulated
        self.discount = 1  # Chronological discount of the reward
        self.temperature_threshold = None  # Number of moves before dropping the temperature given by visit_softmax_temperature_fn to 0 (ie selecting the best action). If None, visit_softmax_temperature_fn is used every time

        # Root prior exploration noise
        self.root_dirichlet_alpha = 0.1
        self.root_exploration_fraction = 0.25

        # UCB formula
        self.pb_c_base = 19652
        self.pb_c_init = 1.25

        ### Network
        self.network = "resnet"  # "resnet" / "fullyconnected"
        self.support_size = 10  # Value and reward are scaled (with almost sqrt) and encoded on a vector with a range of -support_size to support_size. Choose it so that support_size <= sqrt(max(abs(discounted reward)))

        # Residual Network
        self.downsample = "CNN"  # Downsample observations before representation network, False / "CNN" (lighter) / "resnet" (See paper appendix Network Architecture)
        self.blocks = 1  # Number of blocks in the ResNet
        self.channels = 16  # Number of channels in the ResNet
        self.reduced_channels_reward = 4  # Number of channels in reward head
        self.reduced_channels_value = 4  # Number of channels in value head
        self.reduced_channels_policy = 4  # Number of channels in policy head
        self.resnet_fc_reward_layers = [16]  # Define the hidden layers in the reward head of the dynamic network
        self.resnet_fc_value_layers = [16]  # Define the hidden layers in the value head of the prediction network
        self.resnet_fc_policy_layers = [16]  # Define the hidden layers in the policy head of the prediction network

        # Fully Connected Network
        self.encoding_size = 64
        self.fc_representation_layers = []  # Define the hidden layers in the representation network
        self.fc_dynamics_layers = [16]  # Define the hidden layers in the dynamics network
        self.fc_reward_layers = [16]  # Define the hidden layers in the reward network
        self.fc_value_layers = []  # Define the hidden layers in the value network
        self.fc_policy_layers = []  # Define the hidden layers in the policy network

        ### Training
        self.results_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "../results", os.path.basename(__file__)[:-3], datetime.datetime.now().strftime("%Y-%m-%d--%H-%M-%S"))  # Path to store the model weights and TensorBoard logs
        self.save_model = True  # Save the checkpoint in results_path as model.checkpoint
        self.training_steps = 100  # Total number of training steps (ie weights update according to a batch)
        self.batch_size = 2  # Number of parts of games to train on at each training step
        self.checkpoint_interval = 10  # Number of training steps before using the model for self-playing
        self.value_loss_weight = 0.25  # Scale the value loss to avoid overfitting of the value function, paper recommends 0.25 (See paper appendix Reanalyze)
        self.train_on_gpu = torch.cuda.is_available()  # Train on GPU if available

        self.optimizer = "SGD"  # "Adam" or "SGD". Paper uses SGD
        self.weight_decay = 1e-4  # L2 weights regularization
        self.momentum = 0.9  # Used only if optimizer is SGD

        # Exponential learning rate schedule
        self.lr_init = 0.003  # Initial learning rate
        self.lr_decay_rate = 1  # Set it to 1 to use a constant learning rate
        self.lr_decay_steps = 10000

        ### Replay Buffer
        self.replay_buffer_size = 3000  # Number of self-play games to keep in the replay buffer
        self.num_unroll_steps = 20  # Number of game moves to keep for every batch element
        self.td_steps = 20  # Number of steps in the future to take into account for calculating the target value
        self.PER = True  # Prioritized Replay (See paper appendix Training), select in priority the elements in the replay buffer which are unexpected for the network
        self.PER_alpha = 0.5  # How much prioritization is used, 0 corresponding to the uniform case, paper suggests 1

        # Reanalyze (See paper appendix Reanalyse)
        self.use_last_model_value = True  # Use the last model to provide a fresher, stable n-step value (See paper appendix Reanalyze)
        self.reanalyse_on_gpu = False

        ### Adjust the self play / training ratio to avoid over/underfitting
        self.self_play_delay = 0  # Number of seconds to wait after each played game
        self.training_delay = 0  # Number of seconds to wait after each training step
        self.ratio = None  # Desired training steps per self played step ratio. Equivalent to a synchronous version, training can take much longer. Set it to None to disable it

    def visit_softmax_temperature_fn(self, trained_steps):
        """
        Parameter to alter the visit count distribution to ensure that the action selection becomes greedier as training progresses.
        The smaller it is, the more likely the best action (ie with the highest visit count) is chosen.
        Returns:
            Positive float.
        """
        return 1

class Game(AbstractGame):

    def __init__(self, seed=None):
        pygame.init()
        settings = Settings()
        clock = pygame.time.Clock()
        screen = pygame.display.set_mode(settings.screen_sizes)
        pygame.display.set_caption('Environment')
        game = GameFunc(settings, screen)
        self.env = game

    def step(self, action):
        observation, reward, done = self.env.step(action)
        return observation, reward, done

    def to_play(self):
        return self.env.to_play()

    def legal_actions(self):
        return list(range(4))

    def reset(self):
        return self.env.reset()

    def render(self):
        self.env.render()


from flask import Flask, request, render_template

import tensorflow as tf
import tensorflow_hub as hub
import numpy as np
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import load_img, img_to_array

app = Flask(__name__)

MODEL_PATH ='Mushroom_Model.h5'
model = tf.keras.models.load_model(MODEL_PATH, custom_objects={'KerasLayer': hub.KerasLayer})

@app.route('/', methods=['GET'])
def home():
    return render_template('index.html')

@app.route('/', methods=['POST'])
def predict():
    imagefile = request.files['imagefile']
    image_path = "./images/" + imagefile.filename
    imagefile.save(image_path)

    image = load_img(image_path, target_size=(224, 224))
    image = img_to_array(image)
    image = image/255
    image = np.expand_dims(image, axis=0)

    preds = model.predict(image)

    preds=np.argmax(model.predict(image))

    if preds==0:
        classification = "PREDICTION: Edible"
    else:
        classification = "PREDICTION: Poisonous"

    return render_template('index.html', prediction = classification)

if __name__ == '__main__':
    app.run(debug=True, port=3000)


import click
import os
import random
from Bio import SeqIO
import numpy as np

def random_base(reference, amino):
    if amino:
        bases = {'A': 1, 'R': 2, 'N': 3, 'D': 4, 'C': 5, 'Q': 6, 'E': 7, 'G': 8, 'H': 9, 'I': 10,
                 'L': 11, 'K': 12, 'M': 13, 'F': 14, 'P': 15, 'S': 16, 'T': 17, 'W': 18, 'Y': 19, 'V': 20}
    else:
        bases = {'A': 1, 'C': 2, 'T': 3, 'G': 4}
    try:
        del bases[reference.upper()]
    except:
        pass

    rbases = [i for i in bases.keys()]
    random.shuffle(rbases)
    return rbases[0]

def num_to_base(reference, amino):
    if amino:
        bases = {1: 'A', 2: 'R', 3: 'N', 4: 'D', 5: 'C', 6: 'Q', 7: 'E', 8: 'G', 9: 'H', 10: 'I',
                 11: 'L', 12: 'K', 13: 'M', 14: 'F', 15: 'P', 16: 'S', 17: 'T', 18: 'W', 19: 'Y', 20: 'V'}
    else:
        bases = {1: 'A', 2: 'C', 3: 'T', 4: 'G'}
    return bases[reference]

def insert(sequence, positions, max_indel_size, amino):
    opts = [i for i in range(1, max_indel_size+1)]

    for pos in positions:
        random.shuffle(opts)
        num_indel = opts[0]
        indel = "".join([num_to_base(random.randint(1, 4), amino)
                         for _ in range(num_indel)])

        sequence[pos] = indel

    return sequence

def delete(sequence, positions, max_indel_size, amino):
    opts = [i for i in range(1, max_indel_size+1)]

    for ix, pos in enumerate(positions):
        random.shuffle(opts)
        num_indel = opts[0]
        try:
            del sequence[pos-ix:num_indel+pos-ix]
        except:
            pass
    return sequence

def mismatch(sequence, positions, amino):
    for pos in positions:
        reference = sequence[pos]
        sequence[pos] = random_base(reference, amino)
    return sequence

@click.command()
@click.option('--input-file', default='', help='input fasta file to mutate fasta file')
@click.option('--mutations', default=10, help='percentage of mutations in the sequence (default: 5%)')
@click.option('--insertions', default=30, help='percentage of insertions out of the total mutations (default: 30%)')
@click.option('--deletions', default=30, help='percentage of deletions out of the total mutations (default: 30%)')
@click.option('--mismatches', default=40, help='percentage of mismatches out of the total mutations (default: 40%)')
@click.option('--max-indel-size', default=5, help='maximum indel size (default: 5)')
@click.option('--prefix', default='0', help='prefix added to output file')
@click.option('--prot', default=False, is_flag=True, help='aminoacid sequences')
def mutate(input_file, insertions, deletions, mismatches, mutations, max_indel_size, prefix, prot):
    '''
        Mutate a nucleotide sequence:

        This script takes a fasta file as input an mutates the sequence according to the insertion, deletion, and mismatches rates.
        Output is a fasta file with the modified entries.

     '''

    if not input_file:
        os.system('genetools mutate --help')
        exit()

    fo = open(input_file+'.M'+str(mutations)+'.m'+str(mismatches)+'.i' +
              str(insertions)+'.d'+str(deletions)+'.prefix-'+prefix+'.mut.fasta', 'w')

    # mutations are 2x
    mutations = int(mutations/2)

    # load fasta file
    fasta_file = SeqIO.parse(input_file, 'fasta')
    for record in fasta_file:

        # get the positions to mutate
        _id = record.id
        _sequence = list(record.seq)
        positions = np.array(range(len(_sequence)))
        random.shuffle(positions)
        # print(positions)

        _mutation_rate = int(len(_sequence) * mutations / 100)
        positions_to_mutate = positions[:_mutation_rate]
        random.shuffle(positions_to_mutate)
        # print(positions_to_mutate)

        _mismatches_rate = int(mismatches * _mutation_rate / 100)
        mismatches_positions = positions_to_mutate[:_mismatches_rate]
        # print(mismatches_positions)

        _insertions_rate = int(insertions * _mutation_rate / 100)
        insertions_positions = positions_to_mutate[_mismatches_rate:
                                                   _mismatches_rate + _insertions_rate]
        # print(insertions_positions)

        deletions_positions = positions_to_mutate[_mismatches_rate +
                                                  _insertions_rate:]
        # print(deletions_positions)

        _sequence = mismatch(_sequence, mismatches_positions, prot)
        _sequence = insert(_sequence, insertions_positions,
                           max_indel_size, prot)
        _sequence = delete(_sequence, deletions_positions,
                           max_indel_size, prot)

        entry = "".join(['>mut-', _id, '\n', "".join(_sequence), '\n'])
        fo.write(entry)


#!/usr/bin/env python
#
# # Protool - Python class for manipulating protein structures
# Copyright (C) 2010 Jens Erik Nielsen
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Contact information:
# Email: Jens.Nielsen_at_gmail.com
# Normal mail:
# Jens Nielsen
# SBBS, Conway Institute
# University College Dublin
# Dublin 4, Ireland

from numpy import *

class access:

    def __init__(self,protein):
        """Store the protool instance and calculate boxes"""
        self.PI=protein
        self.cutoff=8.0
        #
        # Count number of atoms within 6A of side chain atoms. Only heavy atoms are counted
        #
        self.residues=self.PI.residues.keys()
        self.residues.sort()
        return

    #
    # ----
    #

    def get_access(self,residue):
        """Get the number of atoms close"""
        acc=0.0
        count=0
        for res2 in self.PI.residues.keys():
            if res2==residue:
                continue
            if self.PI.dist(residue+':CA',res2+':CA')<35.0:
                counted={}
                for atom1 in self.PI.residues[residue]:
                    #if self.PI.is_backbone(atom1) and atom1.split(':')[-1]!='CA':
                    #    continue
                    for atom2 in self.PI.residues[res2]:
                        if not counted.has_key(atom2):
                            if self.PI.dist(atom1,atom2)<self.cutoff:
                                acc=acc+1.0
                                counted[atom2]=1
        #
        #
        #numatoms=float(len(self.PI.residues[residue])-4) #Subtract backbone atoms
        #
        # Correct for compactness of aromatics
        #
        #if self.PI.resname(residue) in ['PHE','TYR','HIS']:
        #    numatoms=numatoms-3
        #if self.PI.resname(residue) in ['TRP']:
        #    numatoms=numatoms-6
        #if numatoms>0.0:
        #    acc=acc#/numatoms
        #else:
        #    acc=0
        return max(0.0,-(acc-150)/2.3)


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import csv
import re
import json
import numpy as np
import smtplib
import argparse

from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.image import MIMEImage

parser = argparse.ArgumentParser(
    description='summarize ARAMIS LAB\'s Nespresso orders and send bills accordingly',
    prog='naramispresso',
    epilog='see https://github.com/jguillon/naramispresso for updates, bugs reports and answers')

parser.add_argument('typeform_csv', help='csv file exported from Typeform')
parser.add_argument('--send', help='send bills', action='store_true')
parser.add_argument('--date', help='reception date') # '03/08/2018 entre 9h00 et 18h00'
args, unknown_args = parser.parse_known_args()

with open('parameters.json') as f:
    params = json.load(f)

if args.send:
    if not args.date:
        print('Please specify a reception date using the \'--date\' argument.')
        exit(1)

    with open('bill.mail.html', 'r') as html_file:
        mail_html = html_file.read()
        mail_html = mail_html.replace('{reception_date}', args.date)
        mail_html = mail_html.replace('{paypal_id}', params['paypal_id'])
        mail_html = mail_html.replace('{slack_url}', params['slack_url'])
        mail_html = mail_html.replace('{contact_mail}', params['contact_mail'])
        mail_html = mail_html.replace('{office_address}', params['office_address'])
        html_file.close()

with open(args.typeform_csv, 'rU') as csv_file:
    nespresso_report_reader = csv.reader(csv_file, delimiter=',', quotechar='"')

    # Extract prices on first row
    first_row = nespresso_report_reader.__next__()
    brands = first_row[1:-8]
    prices = [0 for i in range(len(brands))]
    overall_quantities = [0 for i in range(len(brands))]
    overall_total = 0
    for i, brand in enumerate(brands):
        m = re.search('[0-9]+,[0-9]+', brand)
        prices[i] = float(m.group(0).replace(',','.'))

    # Itarate over customers lines
    for row in nespresso_report_reader:

        ## Information extraction
        first_name = row[-8]
        last_name = row[-7]
        mail = row[-6]
        quantities = [int(i) if i else 0 for i in row[1:-8]]
        overall_quantities = [sum(x) for x in zip(overall_quantities, quantities)]
        total = sum([quantities[i] * prices[i] for i in range(len(brands))])
        overall_total = overall_total + total
        print('\n' + first_name + ' ' + last_name + ' <' + mail + '> :')

        ## HTML table creation
        command_details = ''
        for i, quantity in enumerate(quantities):
            if quantity > 0:
                print(str(quantity) + ' x ' + brands[i])
                command_details = command_details \
                    + '\t\t\t\t
' \
                    + '' + brands[i] + '	' \
                    + '' + str(quantity) + ' x ' + str(prices[i]) + '€' + '	' \
                    + '' + str(quantity * prices[i]) + '€' + '	' \
                    + '
\n'
        print("=> " + str(sum(quantities)) + ' BOXES (' + str(total) + "€)")

        if args.send:

            ## Mail content writing
            perso_mail_html = mail_html.replace('{first_name}',
                                                first_name)
            perso_mail_html = perso_mail_html.replace('{last_name}',
                                                      last_name)
            perso_mail_html = perso_mail_html.replace('{first_name}',
                                                      first_name)
            perso_mail_html = perso_mail_html.replace('{command_details}',
                                                      command_details)
            perso_mail_html = perso_mail_html.replace('{total}',
                                                      str(total) )
            perso_mail_text = 'Oops! Use a modern mail app!'

            ## Mail object building
            msg = MIMEMultipart('related')
            msg['Subject'] = 'Naramispresso : Confirmation de Commande'
            msg['From'] = 'Naramispresso'
            msg['To'] = mail

            msg_alt = MIMEMultipart('alternative')
            msg.attach(msg_alt)

            part1 = MIMEText(perso_mail_text, 'plain', 'utf-8')
            msg_alt.attach(part1)

            part2 = MIMEText(perso_mail_html, 'html', 'utf-8')
            msg_alt.attach(part2)

            ## Deal with images
            fp = open('img/naramispresso-logotext-bw.png', 'rb')
            image1 = MIMEImage(fp.read())
            fp.close()
            image1.add_header('Content-ID', '<naramispresso-logotext-bw>')
            msg.attach(image1)

            fp = open('img/naramispresso-mail-header.jpeg', 'rb')
            image2 = MIMEImage(fp.read())
            fp.close()
            image2.add_header('Content-ID', '<naramispresso-mail-header>')
            msg.attach(image2)
            try:
                print('Sending bill...')
                ## Mail sending
                smtp_server = smtplib.SMTP_SSL(params['mail_server']['address'],
                                               params['mail_server']['port'])
                smtp_server.ehlo()
                # smtp_server.starttls()
                smtp_server.login(params['mail_server']['login'],
                                  params['mail_server']['password'])
                smtp_server.sendmail(msg['From'], msg['To'], msg.as_string())
                smtp_server.sendmail(msg['From'], params['contact_mail'],
                                     msg.as_string())
                smtp_server.close()
                print('            ...OK')
            except Exception as e:
                print(e)
                print('Somehting went wrong when trying to send the mails.')

    csv_file.close()

    ## Print recap'
    print("\n============= TOTAL =============\n")
    for i, quantity in enumerate(overall_quantities):
        if quantity > 0:
            print(str(quantity) + ' x ' + brands[i])
    print("=> " + str(sum(overall_quantities)) + ' BOXES (' +
          str(overall_total) + "€)")


import ast
import numpy as np
from PIL import Image
import torch
import torch.nn as nn
from torchvision import transforms as T
import torch.nn.functional as F
from collections import OrderedDict
from timm.models import create_model
from model_service.pytorch_model_service import PTServingBaseService
class garbage_classify_service(PTServingBaseService):
    def __init__(self, model_name, model_path):
        # these three parameters are no need to modify
        self.model_name = model_name
        self.model_path = model_path
        self.signature_key = 'predict_images'

        self.input_size = 456  # the input image size of the model

        # add the input and output key of your pb model here,
        # these keys are defined when you save a pb file
        self.input_key_1 = 'input_img'
        self.output_key_1 = 'output_score'
        self.transforms = T.Compose([
            T.Resize((int(self.input_size / 0.934), int(self.input_size/ 0.934))),
            T.CenterCrop((self.input_size,self.input_size)),
            T.ToTensor(),
            T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
        ])
        self.model = create_model(
            'tf_efficientnet_b5_ns',
            pretrained=False,
            num_classes=43,
            drop_rate=0,
            drop_connect_rate=None,  # DEPRECATED, use drop_path
            drop_path_rate=None,
            drop_block_rate=None,
            global_pool=None,
            bn_tf=False,
            bn_momentum=None,
            bn_eps=None,
            checkpoint_path=None)

        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        net_weight = torch.load(model_path,map_location=torch.device('cpu'))
        # new_pre = {}
        # for k, v in net_weight.items():
        #     name = k[7:]
        #     new_pre[name] = v
        # #
        self.model.load_state_dict(net_weight)
        self.model.eval()

        self.label_id_name_dict = \
            {
                "0": "其他垃圾/一次性快餐盒",
                "1": "其他垃圾/污损塑料",
                "2": "其他垃圾/烟蒂",
                "3": "其他垃圾/牙签",
                "4": "其他垃圾/破碎花盆及碟碗",
                "5": "其他垃圾/竹筷",
                "6": "厨余垃圾/剩饭剩菜",
                "7": "厨余垃圾/大骨头",
                "8": "厨余垃圾/水果果皮",
                "9": "厨余垃圾/水果果肉",
                "10": "厨余垃圾/茶叶渣",
                "11": "厨余垃圾/菜叶菜根",
                "12": "厨余垃圾/蛋壳",
                "13": "厨余垃圾/鱼骨",
                "14": "可回收物/充电宝",
                "15": "可回收物/包",
                "16": "可回收物/化妆品瓶",
                "17": "可回收物/塑料玩具",
                "18": "可回收物/塑料碗盆",
                "19": "可回收物/塑料衣架",
                "20": "可回收物/快递纸袋",
                "21": "可回收物/插头电线",
                "22": "可回收物/旧衣服",
                "23": "可回收物/易拉罐",
                "24": "可回收物/枕头",
                "25": "可回收物/毛绒玩具",
                "26": "可回收物/洗发水瓶",
                "27": "可回收物/玻璃杯",
                "28": "可回收物/皮鞋",
                "29": "可回收物/砧板",
                "30": "可回收物/纸板箱",
                "31": "可回收物/调料瓶",
                "32": "可回收物/酒瓶",
                "33": "可回收物/金属食品罐",
                "34": "可回收物/锅",
                "35": "可回收物/食用油桶",
                "36": "可回收物/饮料瓶",
                "37": "有害垃圾/干电池",
                "38": "有害垃圾/软膏",
                "39": "有害垃圾/过期药物",
                "40": "可回收物/毛巾",
                "41": "可回收物/饮料盒",
                "42": "可回收物/纸袋"
            }
    #
    def preprocess_img(self, img):
        img = img.convert('RGB')
        img = self.transforms(img)
        return img

    def _preprocess(self, data):
        preprocessed_data = {}
        for k, v in data.items():
            for file_name, file_content in v.items():
                img = Image.open(file_content)
                img = self.preprocess_img(img)
                preprocessed_data[k] = img
        return preprocessed_data

    def _inference(self, data):
        """
        model inference function
        Here are a inference example of resnet, if you use another model, please modify this function
        """
        img = data[self.input_key_1]
        img = img.unsqueeze(0)
        #
        if torch.cuda.is_available():
            img = img.cuda()
        with torch.no_grad():
            output = self.model(img)
            #output = F.softmax(output.data, dim=1)
            if output is not None:
                _, pred = torch.max(output.data, 1)
                pred_label = pred.cpu().numpy().tolist()[0]
                result = {'result': self.label_id_name_dict[str(pred_label)]}
            else:
                result = {'result': 'predict score is None'}
        return result

    def _postprocess(self, data):
        return data


# Authors: Eric Larson <larson.eric.d@gmail.com>
#
# License: BSD (3-clause)

import numpy as np
from scipy.fftpack import fft, ifft, rfft, irfft

from .utils import sizeof_fmt, logger, get_config, warn, _explain_exception

# Support CUDA for FFTs; requires scikits.cuda and pycuda
_cuda_capable = False
_multiply_inplace_c128 = _halve_c128 = _real_c128 = None

def _get_cudafft():
    """Helper to deal with scikit-cuda namespace change"""
    try:
        from skcuda import fft
    except ImportError:
        try:
            from scikits.cuda import fft
        except ImportError:
            fft = None
    return fft

def get_cuda_memory():
    """Get the amount of free memory for CUDA operations

    Returns
    -------
    memory : str
        The amount of available memory as a human-readable string.
    """
    if not _cuda_capable:
        warn('CUDA not enabled, returning zero for memory')
        mem = 0
    else:
        from pycuda.driver import mem_get_info
        mem = mem_get_info()[0]
    return sizeof_fmt(mem)

def init_cuda(ignore_config=False):
    """Initialize CUDA functionality

    This function attempts to load the necessary interfaces
    (hardware connectivity) to run CUDA-based filtering. This
    function should only need to be run once per session.

    If the config var (set via mne.set_config or in ENV)
    MNE_USE_CUDA == 'true', this function will be executed when
    the first CUDA setup is performed. If this variable is not
    set, this function can be manually executed.
    """
    global _cuda_capable, _multiply_inplace_c128, _halve_c128, _real_c128
    if _cuda_capable:
        return
    if not ignore_config and (get_config('MNE_USE_CUDA', 'false').lower() !=
                              'true'):
        logger.info('CUDA not enabled in config, skipping initialization')
        return
    # Triage possible errors for informative messaging
    _cuda_capable = False
    try:
        from pycuda import gpuarray, driver  # noqa
        from pycuda.elementwise import ElementwiseKernel
    except ImportError:
        warn('module pycuda not found, CUDA not enabled')
        return
    try:
        # Initialize CUDA; happens with importing autoinit
        import pycuda.autoinit  # noqa
    except ImportError:
        warn('pycuda.autoinit could not be imported, likely a hardware error, '
             'CUDA not enabled%s' % _explain_exception())
        return
    # Make sure scikit-cuda is installed
    cudafft = _get_cudafft()
    if cudafft is None:
        warn('module scikit-cuda not found, CUDA not enabled')
        return

    # let's construct our own CUDA multiply in-place function
    _multiply_inplace_c128 = ElementwiseKernel(
        'pycuda::complex<double> *a, pycuda::complex<double> *b',
        'b[i] *= a[i]', 'multiply_inplace')
    _halve_c128 = ElementwiseKernel(
        'pycuda::complex<double> *a', 'a[i] /= 2.0', 'halve_value')
    _real_c128 = ElementwiseKernel(
        'pycuda::complex<double> *a', 'a[i] = real(a[i])', 'real_value')

    # Make sure we can use 64-bit FFTs
    try:
        cudafft.Plan(16, np.float64, np.complex128)  # will get auto-GC'ed
    except Exception:
        warn('Device does not appear to support 64-bit FFTs, CUDA not '
             'enabled%s' % _explain_exception())
        return
    _cuda_capable = True
    # Figure out limit for CUDA FFT calculations
    logger.info('Enabling CUDA with %s available memory' % get_cuda_memory())

###############################################################################
# Repeated FFT multiplication

def setup_cuda_fft_multiply_repeated(n_jobs, h_fft):
    """Set up repeated CUDA FFT multiplication with a given filter

    Parameters
    ----------
    n_jobs : int | str
        If n_jobs == 'cuda', the function will attempt to set up for CUDA
        FFT multiplication.
    h_fft : array
        The filtering function that will be used repeatedly.
        If n_jobs='cuda', this function will be shortened (since CUDA
        assumes FFTs of real signals are half the length of the signal)
        and turned into a gpuarray.

    Returns
    -------
    n_jobs : int
        Sets n_jobs = 1 if n_jobs == 'cuda' was passed in, otherwise
        original n_jobs is passed.
    cuda_dict : dict
        Dictionary with the following CUDA-related variables:
            use_cuda : bool
                Whether CUDA should be used.
            fft_plan : instance of FFTPlan
                FFT plan to use in calculating the FFT.
            ifft_plan : instance of FFTPlan
                FFT plan to use in calculating the IFFT.
            x_fft : instance of gpuarray
                Empty allocated GPU space for storing the result of the
                frequency-domain multiplication.
            x : instance of gpuarray
                Empty allocated GPU space for the data to filter.
    h_fft : array | instance of gpuarray
        This will either be a gpuarray (if CUDA enabled) or np.ndarray.
        If CUDA is enabled, h_fft will be modified appropriately for use
        with filter.fft_multiply().

    Notes
    -----
    This function is designed to be used with fft_multiply_repeated().
    """
    cuda_dict = dict(use_cuda=False, fft_plan=None, ifft_plan=None,
                     x_fft=None, x=None)
    n_fft = len(h_fft)
    cuda_fft_len = int((n_fft - (n_fft % 2)) / 2 + 1)
    if n_jobs == 'cuda':
        n_jobs = 1
        init_cuda()
        if _cuda_capable:
            from pycuda import gpuarray
            cudafft = _get_cudafft()
            # set up all arrays necessary for CUDA
            # try setting up for float64
            try:
                # do the IFFT normalization now so we don't have to later
                h_fft = gpuarray.to_gpu(h_fft[:cuda_fft_len]
                                        .astype('complex_') / len(h_fft))
                cuda_dict.update(
                    use_cuda=True,
                    fft_plan=cudafft.Plan(n_fft, np.float64, np.complex128),
                    ifft_plan=cudafft.Plan(n_fft, np.complex128, np.float64),
                    x_fft=gpuarray.empty(cuda_fft_len, np.complex128),
                    x=gpuarray.empty(int(n_fft), np.float64))
                logger.info('Using CUDA for FFT FIR filtering')
            except Exception as exp:
                logger.info('CUDA not used, could not instantiate memory '
                            '(arrays may be too large: "%s"), falling back to '
                            'n_jobs=1' % str(exp))
        else:
            logger.info('CUDA not used, CUDA could not be initialized, '
                        'falling back to n_jobs=1')
    return n_jobs, cuda_dict, h_fft

def fft_multiply_repeated(h_fft, x, cuda_dict=dict(use_cuda=False)):
    """Do FFT multiplication by a filter function (possibly using CUDA)

    Parameters
    ----------
    h_fft : 1-d array or gpuarray
        The filtering array to apply.
    x : 1-d array
        The array to filter.
    cuda_dict : dict
        Dictionary constructed using setup_cuda_multiply_repeated().

    Returns
    -------
    x : 1-d array
        Filtered version of x.
    """
    if not cuda_dict['use_cuda']:
        # do the fourier-domain operations
        x = np.real(ifft(h_fft * fft(x), overwrite_x=True)).ravel()
    else:
        cudafft = _get_cudafft()
        # do the fourier-domain operations, results in second param
        cuda_dict['x'].set(x.astype(np.float64))
        cudafft.fft(cuda_dict['x'], cuda_dict['x_fft'], cuda_dict['fft_plan'])
        _multiply_inplace_c128(h_fft, cuda_dict['x_fft'])
        # If we wanted to do it locally instead of using our own kernel:
        # cuda_seg_fft.set(cuda_seg_fft.get() * h_fft)
        cudafft.ifft(cuda_dict['x_fft'], cuda_dict['x'],
                     cuda_dict['ifft_plan'], False)
        x = np.array(cuda_dict['x'].get(), dtype=x.dtype, subok=True,
                     copy=False)
    return x

###############################################################################
# FFT Resampling

def setup_cuda_fft_resample(n_jobs, W, new_len):
    """Set up CUDA FFT resampling

    Parameters
    ----------
    n_jobs : int | str
        If n_jobs == 'cuda', the function will attempt to set up for CUDA
        FFT resampling.
    W : array
        The filtering function to be used during resampling.
        If n_jobs='cuda', this function will be shortened (since CUDA
        assumes FFTs of real signals are half the length of the signal)
        and turned into a gpuarray.
    new_len : int
        The size of the array following resampling.

    Returns
    -------
    n_jobs : int
        Sets n_jobs = 1 if n_jobs == 'cuda' was passed in, otherwise
        original n_jobs is passed.
    cuda_dict : dict
        Dictionary with the following CUDA-related variables:
            use_cuda : bool
                Whether CUDA should be used.
            fft_plan : instance of FFTPlan
                FFT plan to use in calculating the FFT.
            ifft_plan : instance of FFTPlan
                FFT plan to use in calculating the IFFT.
            x_fft : instance of gpuarray
                Empty allocated GPU space for storing the result of the
                frequency-domain multiplication.
            x : instance of gpuarray
                Empty allocated GPU space for the data to resample.
    W : array | instance of gpuarray
        This will either be a gpuarray (if CUDA enabled) or np.ndarray.
        If CUDA is enabled, W will be modified appropriately for use
        with filter.fft_multiply().

    Notes
    -----
    This function is designed to be used with fft_resample().
    """
    cuda_dict = dict(use_cuda=False, fft_plan=None, ifft_plan=None,
                     x_fft=None, x=None, y_fft=None, y=None)
    n_fft_x, n_fft_y = len(W), new_len
    cuda_fft_len_x = int((n_fft_x - (n_fft_x % 2)) // 2 + 1)
    cuda_fft_len_y = int((n_fft_y - (n_fft_y % 2)) // 2 + 1)
    if n_jobs == 'cuda':
        n_jobs = 1
        init_cuda()
        if _cuda_capable:
            # try setting up for float64
            from pycuda import gpuarray
            cudafft = _get_cudafft()
            try:
                # do the IFFT normalization now so we don't have to later
                W = gpuarray.to_gpu(W[:cuda_fft_len_x]
                                    .astype('complex_') / n_fft_y)
                cuda_dict.update(
                    use_cuda=True,
                    fft_plan=cudafft.Plan(n_fft_x, np.float64, np.complex128),
                    ifft_plan=cudafft.Plan(n_fft_y, np.complex128, np.float64),
                    x_fft=gpuarray.zeros(max(cuda_fft_len_x,
                                             cuda_fft_len_y), np.complex128),
                    x=gpuarray.empty(max(int(n_fft_x),
                                     int(n_fft_y)), np.float64))
                logger.info('Using CUDA for FFT resampling')
            except Exception:
                logger.info('CUDA not used, could not instantiate memory '
                            '(arrays may be too large), falling back to '
                            'n_jobs=1')
        else:
            logger.info('CUDA not used, CUDA could not be initialized, '
                        'falling back to n_jobs=1')
    return n_jobs, cuda_dict, W

def fft_resample(x, W, new_len, npads, to_removes,
                 cuda_dict=dict(use_cuda=False)):
    """Do FFT resampling with a filter function (possibly using CUDA)

    Parameters
    ----------
    x : 1-d array
        The array to resample. Will be converted to float64 if necessary.
    W : 1-d array or gpuarray
        The filtering function to apply.
    new_len : int
        The size of the output array (before removing padding).
    npads : tuple of int
        Amount of padding to apply to the start and end of the
        signal before resampling.
    to_removes : tuple of int
        Number of samples to remove after resampling.
    cuda_dict : dict
        Dictionary constructed using setup_cuda_multiply_repeated().

    Returns
    -------
    x : 1-d array
        Filtered version of x.
    """
    # add some padding at beginning and end to make this work a little cleaner
    if x.dtype != np.float64:
        x = x.astype(np.float64)
    x = _smart_pad(x, npads)
    old_len = len(x)
    shorter = new_len < old_len
    if not cuda_dict['use_cuda']:
        N = int(min(new_len, old_len))
        # The below is equivalent to this, but faster
        # sl_1 = slice((N + 1) // 2)
        # y_fft = np.zeros(new_len, np.complex128)
        # x_fft = fft(x).ravel() * W
        # y_fft[sl_1] = x_fft[sl_1]
        # sl_2 = slice(-(N - 1) // 2, None)
        # y_fft[sl_2] = x_fft[sl_2]
        # y = np.real(ifft(y_fft, overwrite_x=True)).ravel()
        x_fft = rfft(x).ravel()
        x_fft *= W[np.arange(1, len(x) + 1) // 2].real
        y_fft = np.zeros(new_len, np.float64)
        sl_1 = slice(N)
        y_fft[sl_1] = x_fft[sl_1]
        if min(new_len, old_len) % 2 == 0:
            if new_len > old_len:
                y_fft[N - 1] /= 2.
        y = irfft(y_fft, overwrite_x=True).ravel()
    else:
        cudafft = _get_cudafft()
        cuda_dict['x'].set(np.concatenate((x, np.zeros(max(new_len - old_len,
                                                           0), x.dtype))))
        # do the fourier-domain operations, results put in second param
        cudafft.fft(cuda_dict['x'], cuda_dict['x_fft'], cuda_dict['fft_plan'])
        _multiply_inplace_c128(W, cuda_dict['x_fft'])
        # This is not straightforward, but because x_fft and y_fft share
        # the same data (and only one half of the full DFT is stored), we
        # don't have to transfer the slice like we do in scipy. All we
        # need to worry about is the Nyquist component, either halving it
        # or taking just the real component...
        use_len = new_len if shorter else old_len
        func = _real_c128 if shorter else _halve_c128
        if use_len % 2 == 0:
            nyq = int((use_len - (use_len % 2)) // 2)
            func(cuda_dict['x_fft'], slice=slice(nyq, nyq + 1))
        cudafft.ifft(cuda_dict['x_fft'], cuda_dict['x'],
                     cuda_dict['ifft_plan'], scale=False)
        y = cuda_dict['x'].get()[:new_len if shorter else None]

    # now let's trim it back to the correct size (if there was padding)
    if (to_removes > 0).any():
        keep = np.ones((new_len), dtype='bool')
        keep[:to_removes[0]] = False
        keep[-to_removes[1]:] = False
        y = np.compress(keep, y)

    return y

###############################################################################
# Misc

# this has to go in mne.cuda instead of mne.filter to avoid import errors
def _smart_pad(x, n_pad):
    """Pad vector x
    """
    if (n_pad == 0).all():
        return x
    elif (n_pad < 0).any():
        raise RuntimeError('n_pad must be non-negative')
    # need to pad with zeros if len(x) <= npad
    l_z_pad = np.zeros(max(n_pad[0] - len(x) + 1, 0), dtype=x.dtype)
    r_z_pad = np.zeros(max(n_pad[0] - len(x) + 1, 0), dtype=x.dtype)
    return np.concatenate([l_z_pad, 2 * x[0] - x[n_pad[0]:0:-1], x,
                           2 * x[-1] - x[-2:-n_pad[1] - 2:-1], r_z_pad])

import numpy,from scipy,import pycuda,from pycuda