#!/usr/bin/python3
"""This module manages all the products of provisionspall"""

from api.v1.views import app_views
from flask import jsonify, request, make_response
from models.model import User, Store, Store_Address
from api.v1 import db
from provisionspall_web import UPLOAD_FOLDER, allowed_file
import os
from werkzeug.utils import secure_filename

@app_views.route('/user/store/<store_id>', strict_slashes=False, methods=['DELETE', 'OPTIONS'])
def delete_store(store_id):
    """Delete a store using the store id"""

    if request.method == 'DELETE':
        try:
            store_del = db.session.get(Store, store_id)
            db.session.delete(store_del)
            db.session.commit()
            return jsonify({'Message': 'successfully deleted the store'}), 200
        except Exception as e:
            print(e)
            db.session.rollback()
            return jsonify({'Error': 'Could not delete the store'}), 400
    elif request.method == 'OPTIONS':
        response = make_response()
        response.headers.add("Access-Control-Allow-Origin", "http://127.0.0.1:5000")
        response.headers.add("Access-Control-Allow-Headers", "Content-Type,Authorization")
        response.headers.add("Access-Control-Allow-Methods", "DELETE")
        return response

@app_views.route('/user/<user_id>/stores', strict_slashes=False, methods=['GET', 'POST', 'PUT', 'OPTIONS'])
def get_stores(user_id):
    """Get all the stores as requested"""

    if request.method == 'OPTIONS':
        response = make_response()
        response.headers.add("Access-Control-Allow-Methods", "GET,POST,PUT,OPTIONS")
        return response
    elif request.method == 'GET':
        # Get all store of a User
        try:
            user_stores = db.session.query(User).filter(User.id == user_id).one().stores
            all_stores = []
            for store in user_stores:
                store_details = {
                    'id': store.id,
                    'name': store.name,
                    'description': store.description,
                    'date_created': store.date_created,
                    'user_id': store.user_id,
                    'image': store.image,
                }
                all_stores.append(store_details)
            return jsonify(all_stores), 200

        except Exception:
            db.session.rollback()
            return jsonify({'Error': 'Could not get store data at this time'}), 400

    elif request.method == 'POST' or request.method == 'PUT':
        # Handle a post and get request
        try:
            # Handle file uploads
            file = request.files.get('file')
            filename = ''
            if file and allowed_file(file.filename):
                filename = secure_filename(file.filename)
                file.save(os.path.join(UPLOAD_FOLDER, filename))

            options = {
                'name': request.form.get('name'),
                'description': request.form.get('description'),
                'image': filename
            }
            if request.method == 'POST':
                address_options = {
                                'number': request.form.get('number'),
                                'street':request.form.get('street'),
                                'area': request.form.get('area'),
                                'city': request.form.get('city'),
                                'country': request.form.get('country'),
                                'longitude': request.form.get('longitude'),
                                'latitude': request.form.get('latitude')
                                }
                address = Store_Address(**address_options)
                store = Store(**options)
                store.user_id = user_id
                address.store_id = store.id
                store.addresses = [address]
                db.session.add(store)
                db.session.commit()
                return jsonify({'message': 'Store created successfully', 'store_id': store.id, 'store_address_id': address.id}), 200
            elif request.method == 'PUT':
                store_id = request.form.get('store_id')
                update_store = db.session.get(Store, store_id)
                if update_store:
                    print(options)
                    for key, value in options.items():
                        if value:
                            setattr(update_store, key, value)
                    db.session.commit()
                    return jsonify({'Message': 'Updated successfully', 'store_id': store_id}), 200
                else:
                    return jsonify({'Message': 'Not found'}), 404

        except Exception as e:
            print(e)
            db.session.rollback()
            return jsonify({'Error': 'Error create Store'}), 400

@app_views.route('/stores', strict_slashes=False, methods=['GET'])
def all_stores():
    """Find stores matching the parameters"""
    if request.method == 'GET':
        try:
            stores = db.session.query(Store)
            all_stores = []
            for store in stores:
                store_details = {
                    'id': store.id,
                    'name': store.name,
                    'description': store.description,
                    'date_created': store.date_created,
                    'user_id': store.user_id,
                    'image': store.image,
                }
                all_stores.append(store_details)
            return jsonify(all_stores), 200

        except Exception:
            db.session.rollback()
            return jsonify({'Error': 'Could not get all store data at this time'}), 404


# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

from __future__ import absolute_import, print_function, unicode_literals

"""Utility functions to handle test chunking."""

import os
import json
from abc import ABCMeta, abstractmethod
from collections import defaultdict, OrderedDict

import six
from manifestparser import TestManifest
from manifestparser.filters import chunk_by_runtime
from mozbuild.util import memoize
from moztest.resolve import TestResolver, TestManifestLoader, get_suite_definition

from taskgraph import GECKO

here = os.path.abspath(os.path.dirname(__file__))
resolver = TestResolver.from_environment(cwd=here, loader_cls=TestManifestLoader)

def guess_mozinfo_from_task(task):
    """Attempt to build a mozinfo dict from a task definition.

    This won't be perfect and many values used in the manifests will be missing. But
    it should cover most of the major ones and be "good enough" for chunking in the
    taskgraph.

    Args:
        task (dict): A task definition.

    Returns:
        A dict that can be used as a mozinfo replacement.
    """
    info = {
        'asan': 'asan' in task['build-attributes']['build_platform'],
        'ccov': 'ccov' in task['build-attributes']['build_platform'],
        'debug': task['build-attributes']['build_type'] == 'debug',
        'e10s': task['attributes']['e10s'],
        'tsan': 'tsan' in task['build-attributes']['build_platform'],
        'webrender': task.get('webrender', False),
    }
    for platform in ('android', 'linux', 'mac', 'win'):
        if platform in task['build-attributes']['build_platform']:
            info['os'] = platform
            break
    else:
        raise ValueError("{} is not a known platform!".format(
                         task['build-attributes']['build_platform']))
    return info

@memoize
def get_runtimes(platform, suite_name):
    if not suite_name:
        raise TypeError('suite_name should be a value.')

    base = os.path.join(GECKO, 'testing', 'runtimes', 'manifest-runtimes-{}.json')
    for key in ('android', 'windows'):
        if key in platform:
            path = base.format(key)
            break
    else:
        path = base.format('unix')

    with open(path, 'r') as fh:
        return json.load(fh).get(suite_name)

# This is a bit of a hack to avoid evaluating the entire list of WPT test
# objects in multiple places. It can be removed when the TestResolver becomes
# the source of truth for WPT groups.
wpt_group_translation = defaultdict(set)

def chunk_manifests(flavor, subsuite, platform, chunks, manifests):
    """Run the chunking algorithm.

    Args:
        platform (str): Platform used to find runtime info.
        chunks (int): Number of chunks to split manifests into.
        manifests(list): Manifests to chunk.

    Returns:
        A list of length `chunks` where each item contains a list of manifests
        that run in that chunk.
    """
    # Obtain the suite definition given the flavor and subsuite which often
    # do not perfectly map onto the actual suite name in taskgraph.
    # This value will be used to retrive runtime information for that suite.
    suite_name, _ = get_suite_definition(flavor, subsuite)
    runtimes = get_runtimes(platform, suite_name)

    if flavor != "web-platform-tests":
        return [
            c[1] for c in chunk_by_runtime(
                None,
                chunks,
                runtimes
            ).get_chunked_manifests(manifests)
        ]

    paths = {k: v for k, v in wpt_group_translation.items() if k in manifests}

    # Python2 does not support native dictionary sorting, so use an OrderedDict
    # instead, appending in order of highest to lowest runtime.
    runtimes = OrderedDict(sorted(runtimes.items(), key=lambda x: x[1], reverse=True))

    # Keep track of test paths for each chunk, and the runtime information.
    chunked_manifests = [[[], 0] for _ in range(chunks)]

    # Begin chunking the test paths in order from highest running time to lowest.
    # The keys of runtimes dictionary should match the keys of the test paths
    # dictionary.
    for key, rt in runtimes.items():
        # Sort the chunks from fastest to slowest, based on runtime info
        # at x[1], then the number of test paths.
        chunked_manifests.sort(key=lambda x: (x[1], len(x[0])))

        # Look up if there are any test paths under the key in the paths dict.
        test_paths = paths[key]

        if test_paths:
            # Add the full set of paths that live under the key and increase the
            # total runtime counter by the value reported in runtimes.
            chunked_manifests[0][0].extend(test_paths)
            # chunked_manifests[0][0].append(key)
            chunked_manifests[0][1] += rt
            # Remove the key and test_paths since it has been scheduled.
            paths.pop(key)
            # Same goes for the value in runtimes dict.
            runtimes.pop(key)

    # Sort again prior to the next step.
    chunked_manifests.sort(key=lambda x: (x[1], len(x[0])))

    # Spread out the remaining test paths that were not scheduled in the previous
    # step. Such paths did not have runtime information associated, likely due to
    # implementation status.
    for index, key in enumerate(paths.keys()):
        # Append both the key and value in case the value is empty.
        chunked_manifests[index % chunks][0].append(key)

    # One last sort by the runtime, then number of test paths.
    chunked_manifests.sort(key=lambda x: (x[1], len(x[0])))

    # Return just the chunked test paths.
    return [c[0] for c in chunked_manifests]

@six.add_metaclass(ABCMeta)
class BaseManifestLoader(object):
    @abstractmethod
    def get_manifests(self, flavor, subsuite, mozinfo):
        """Compute which manifests should run for the given flavor, subsuite and mozinfo.

        This function returns skipped manifests separately so that more balanced
        chunks can be achieved by only considering "active" manifests in the
        chunking algorithm.

        Args:
            flavor (str): The suite to run. Values are defined by the 'build_flavor' key
                in `moztest.resolve.TEST_SUITES`.
            subsuite (str): The subsuite to run or 'undefined' to denote no subsuite.
            mozinfo (frozenset): Set of data in the form of (<key>, <value>) used
                                 for filtering.

        Returns:
            A tuple of two manifest lists. The first is the set of active manifests (will
            run at least one test. The second is a list of skipped manifests (all tests are
            skipped).
        """
        pass

class DefaultLoader(BaseManifestLoader):
    """Load manifests using metadata from the TestResolver."""

    @classmethod
    def get_wpt_group(cls, test):
        """Get the group for a web-platform-test that matches those created by the
        WPT harness.

        Args:
            test (dict): The test object to compute the group for.

        Returns:
            str: Label representing the group name.
        """
        depth = 3

        path = os.path.dirname(test['name'])
        # Remove path elements beyond depth of 3, because WPT harness uses
        # --run-by-dir=3 by convention in Mozilla CI.
        # Note, this means /_mozilla tests retain an effective depth of 2 due to
        # the extra prefix element.
        while path.count('/') >= depth + 1:
            path = os.path.dirname(path)

        return path

    @memoize
    def get_tests(self, flavor, subsuite):
        return list(resolver.resolve_tests(flavor=flavor, subsuite=subsuite))

    @memoize
    def get_manifests(self, flavor, subsuite, mozinfo):
        mozinfo = dict(mozinfo)
        # Compute all tests for the given suite/subsuite.
        tests = self.get_tests(flavor, subsuite)

        if flavor == "web-platform-tests":
            manifests = set()
            for t in tests:
                group = self.get_wpt_group(t)
                wpt_group_translation[t['manifest']].add(group)
                manifests.add(t['manifest'])

            return {"active": list(manifests), "skipped": []}

        manifests = set(chunk_by_runtime.get_manifest(t) for t in tests)

        # Compute  the active tests.
        m = TestManifest()
        m.tests = tests
        tests = m.active_tests(disabled=False, exists=False, **mozinfo)
        active = set(chunk_by_runtime.get_manifest(t) for t in tests)
        skipped = manifests - active
        return {"active": list(active), "skipped": list(skipped)}

manifest_loaders = {
    'default': DefaultLoader(),
}


#coding=utf-8
import itertools
import logging
import os.path as osp
import tempfile
import torch
import copy
import random

import mmcv
import numpy as np
from mmcv.utils import print_log
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
from terminaltables import AsciiTable
from torch.utils.data import Dataset

from mmdet.datasets.builder import DATASETS
from mmdet.datasets.pipelines import Compose

@DATASETS.register_module()
class SSLCocoDataset(Dataset):
    CLASSES = ('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus',
               'train', 'truck', 'boat', 'traffic light', 'fire hydrant',
               'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog',
               'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe',
               'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
               'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat',
               'baseball glove', 'skateboard', 'surfboard', 'tennis racket',
               'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
               'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot',
               'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
               'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop',
               'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
               'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock',
               'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush')

    def __init__(self,
                 ann_file_l,
                 ann_file_u,
                 pipeline_l,
                 pipeline_u,
                 classes=None,
                 data_root=None,
                 img_prefix='',
                 seg_prefix=None,
                 proposal_file=None,
                 test_mode=False,
                 filter_empty_gt=True,
                 **kwargs
        ):
        self.ann_file_l = ann_file_l
        self.ann_file_u = ann_file_u
        self.pipeline_l = pipeline_l
        self.pipeline_u = pipeline_u

        self.data_root = data_root
        if isinstance(img_prefix, (list, tuple)):
            self.img_prefix_l, self.img_prefix_u = img_prefix
        else:
            self.img_prefix_l, self.img_prefix_u = img_prefix, img_prefix
        self.seg_prefix = seg_prefix
        self.proposal_file = proposal_file
        self.test_mode = test_mode
        self.filter_empty_gt = filter_empty_gt
        self.CLASSES = self.get_classes(classes)

        # join paths if data_root is specified
        assert self.data_root is None

        # load annotations (and proposals)
        self.coco_l = COCO(self.ann_file_l)
        self.coco_u = COCO(self.ann_file_u)
        self.data_infos_l = self.load_annotations(self.coco_l)
        self.data_infos_u = self.load_annotations(self.coco_u)
        assert len(self.data_infos_u) >= len(self.data_infos_l)
        self.data_ratio = len(self.data_infos_u) // len(self.data_infos_l)

        self.cat_ids = self.coco_l.get_cat_ids(cat_names=self.CLASSES)
        self.cat2label = {cat_id: i for i, cat_id in enumerate(self.cat_ids)}

        # filter data infos if classes are customized
        assert not self.custom_classes
        assert self.proposal_file is None
        self.proposals = None

        # filter images too small
        if not test_mode:
            valid_inds_l = self._filter_imgs(self.coco_l, self.data_infos_l)
            self.data_infos_l = [self.data_infos_l[i] for i in valid_inds_l]
            valid_inds_u = self._filter_imgs(self.coco_u, self.data_infos_u)
            self.data_infos_u = [self.data_infos_u[i] for i in valid_inds_u]

            print('\n')
            print('\tdata_infos_l: {} images'.format(len(self.data_infos_l)))
            print('\tdata_infos_u: {} images'.format(len(self.data_infos_u)))

        # set group flag for the sampler
        if not self.test_mode:
            #self.flag_l = self._set_group_flag(self.data_infos_l)
            #self.flag_u = self._set_group_flag(self.data_infos_u)
            self.flag = np.zeros(len(self.data_infos_l), dtype=np.uint8)

        # processing pipeline
        ## pipeline_l: list of original pipeline in format: [weak_aug, strong_aug_1, strong_aug_2, ...]
        assert isinstance(self.pipeline_l, (list, tuple))
        assert isinstance(self.pipeline_u, (list, tuple))
        assert len(self.pipeline_l) == len(self.pipeline_u)
        assert len(self.pipeline_l) in [2, 3]
        self.pipeline_weak_aug_l = Compose(self.pipeline_l[0][:-1])
        self.pipeline_weak_aug_u = Compose(self.pipeline_u[0][:-1])
        self.pipeline_strong_aug_l = Compose(self.pipeline_l[1][:-1])
        self.pipeline_strong_aug_u = Compose(self.pipeline_u[1][:-1])
        if len(self.pipeline_l) == 3:
            self.pipeline_strong_aug_2_l = Compose(self.pipeline_l[2][:-1])
            self.pipeline_strong_aug_2_u = Compose(self.pipeline_u[2][:-1])
        assert self.pipeline_l[-1][-1]['type'] == self.pipeline_u[-1][-1]['type'] == 'CollectList'
        self.pipeline_collect = Compose(pipeline_l[-1][-1:])

    def _set_group_flag(self, data_infos):
        """Set flag according to image aspect ratio.

        Images with aspect ratio greater than 1 will be set as group 1,
        otherwise group 0.
        """
        flag = np.zeros(len(data_infos), dtype=np.uint8)
        for i in range(len(data_infos)):
            img_info = data_infos[i]
            if img_info['width'] / img_info['height'] > 1:
                flag[i] = 1
        return flag

    @classmethod
    def get_classes(cls, classes=None):
        """Get class names of current dataset
        Args:
            classes (Sequence[str] | str | None): If classes is None, use
                default CLASSES defined by builtin dataset. If classes is a
                string, take it as a file name. The file contains the name of
                classes where each line contains one class name. If classes is
                a tuple or list, override the CLASSES defined by the dataset.
        """
        if classes is None:
            cls.custom_classes = False
            return cls.CLASSES

        cls.custom_classes = True
        if isinstance(classes, str):
            # take it as a file path
            class_names = mmcv.list_from_file(classes)
        elif isinstance(classes, (tuple, list)):
            class_names = classes
        else:
            raise ValueError(f'Unsupported type {type(classes)} of classes.')

        return class_names

    def pre_pipeline(self, results, is_labeled=True):
        """Prepare results dict for pipeline"""
        if is_labeled:
            results['img_prefix'] = self.img_prefix_l
        else:
            results['img_prefix'] = self.img_prefix_u
        results['seg_prefix'] = self.seg_prefix
        results['proposal_file'] = self.proposal_file
        results['bbox_fields'] = []
        results['mask_fields'] = []
        results['seg_fields'] = []

    def _rand_another(self, flag, idx):
        """Get another random index from the same group as the given index"""
        pool = np.where(flag == flag[idx])[0]
        return np.random.choice(pool)

    def get_cat_ids(self, idx):
        """Get COCO category ids by index.
        Args:
            idx (int): Index of data.

        Returns:
            list[int]: All categories in the image of specified index.
        """

        img_id_l = self.data_infos_l[idx]['id']
        ann_ids_l = self.coco_l.get_ann_ids(img_ids=[img_id_l])
        ann_info_l = self.coco_l.load_anns(ann_ids_l)
        return [ann['category_id'] for ann in ann_info_l]

    def __len__(self):
        """Total number of samples of data"""
        return len(self.data_infos_l)

    def __getitem__(self, idx):
        """Get training/test data after pipeline

        Args:
            idx (int): Index of data.

        Returns:
            dict: Training/test data (with annotation if `test_mode` is set
                True).
        """

        if self.test_mode:
            return self.prepare_test_img(idx)
        while True:
            data = self.prepare_train_img(idx)
            if data is None:
                idx = self._rand_another(idx)
                continue
            return data

    def prepare_train_img(self, idx):
        """Get training data and annotations after pipeline.
        Args:
            idx (int): Index of data.
        Returns:
            dict: Training data and annotation after pipeline with new keys
                introduced by pipeline.
        """
        idx_l = idx
        idx_u = idx * self.data_ratio + np.random.randint(self.data_ratio)
        while True:
            ## extract labeled images
            img_info_l = self.data_infos_l[idx_l]
            ann_info_l = self.get_ann_info(self.data_infos_l, self.coco_l, idx_l)
            results_l_w_aug = dict(img_info=img_info_l, ann_info=ann_info_l)
            results_l_s_aug = copy.deepcopy(results_l_w_aug)
            if hasattr(self, 'pipeline_strong_aug_2_l'):
                results_l_s_aug_2 = copy.deepcopy(results_l_w_aug)

            self.pre_pipeline(results_l_w_aug, True)
            self.pre_pipeline(results_l_s_aug, True)
            if hasattr(self, 'pipeline_strong_aug_2_l'):
                self.pre_pipeline(results_l_s_aug_2, True)

            data_l_w = self.pipeline_weak_aug_l(results_l_w_aug)
            data_l_s = self.pipeline_strong_aug_l(results_l_s_aug)
            if hasattr(self, 'pipeline_strong_aug_2_l'):
                data_l_s_2 = self.pipeline_strong_aug_2_l(results_l_s_aug_2)

            if (data_l_w is None) or (data_l_s is None):
                idx_l = self._rand_another(self.flag_l, idx_l)
                continue

            ## extract from unlabeled images
            img_info_u = self.data_infos_u[idx_u]
            ann_info_u = self.get_ann_info(self.data_infos_u, self.coco_u, idx_u)
            results_u_w_aug = dict(img_info=img_info_u, ann_info=ann_info_u)
            results_u_s_aug = copy.deepcopy(results_u_w_aug)
            if hasattr(self, 'pipeline_strong_aug_2_u'):
                results_u_s_aug_2 = copy.deepcopy(results_u_w_aug)

            self.pre_pipeline(results_u_w_aug, False)
            self.pre_pipeline(results_u_s_aug, False)
            if hasattr(self, 'pipeline_strong_aug_2_u'):
                self.pre_pipeline(results_u_s_aug_2, False)

            data_u_w = self.pipeline_weak_aug_u(results_u_w_aug)
            data_u_s = self.pipeline_strong_aug_u(results_u_s_aug)
            if hasattr(self, 'pipeline_strong_aug_2_u'):
                data_u_s_2 = self.pipeline_strong_aug_2_u(results_u_s_aug_2)

            if (data_u_w is None) or (data_u_s is None):
                idx_u = self._rand_another(self.flag_u, idx_u)
                continue

            if hasattr(self, 'pipeline_strong_aug_2_u'):
                return self.pipeline_collect([data_l_w, data_l_s, data_l_s_2, data_u_w, data_u_s, data_u_s_2])

            return self.pipeline_collect([data_l_w, data_l_s, data_u_w, data_u_s])

    def load_annotations(self, coco_obj):
        img_ids = coco_obj.get_img_ids()
        data_infos = []
        for img_id in img_ids:
            info = coco_obj.load_imgs([img_id])[0]
            info['filename'] = info['file_name']
            data_infos.append(info)
        return data_infos

    def get_ann_info(self, data_infos, coco_obj, idx):
        """Get COCO annotation by index.

        Args:
            idx (int): Index of data.

        Returns:
            dict: Annotation info of specified index.
        """
        img_id = data_infos[idx]['id']
        ann_ids = coco_obj.get_ann_ids(img_ids=[img_id])
        ann_info = coco_obj.load_anns(ann_ids)
        return self._parse_ann_info(data_infos[idx], ann_info)

    def _filter_imgs(self, coco_obj, data_infos, min_size=32):
        """Filter images too small or without ground truths."""
        ids_with_ann = set(_['image_id'] for _ in coco_obj.anns.values())

        valid_inds = []
        for i, img_info in enumerate(data_infos):
            if self.filter_empty_gt and img_info['id'] not in ids_with_ann:
                #print('ignore empty groundtruth image')
                continue
            if min(img_info['width'], img_info['height']) >= min_size:
                valid_inds.append(i)
        return valid_inds

    def _parse_ann_info(self, img_info, ann_info):
        """Parse bbox and mask annotation.

        Args:
            ann_info (list[dict]): Annotation info of an image.
            with_mask (bool): Whether to parse mask annotations.

        Returns:
            dict: A dict containing the following keys: bboxes, bboxes_ignore,
                labels, masks, seg_map. "masks" are raw annotations and not
                decoded into binary masks.
        """
        gt_bboxes = []
        gt_labels = []
        gt_bboxes_ignore = []
        gt_masks_ann = []
        gt_scores = []                              # for pseudo boxes
        for i, ann in enumerate(ann_info):
            if ann.get('ignore', False):
                continue

            x1, y1, w, h = ann['bbox']
            x2, y2 = x1 + w, y1 + h
            inter_w = max(0, min(x1 + w, img_info['width']) - max(x1, 0))
            inter_h = max(0, min(y1 + h, img_info['height']) - max(y1, 0))
            if inter_w * inter_h == 0:
                continue
            if ann['area'] <= 0 or w < 1 or h < 1:
                continue
            if ann['category_id'] not in self.cat_ids:
                continue
            bbox = [x1, y1, x1 + w, y1 + h]
            if ann.get('iscrowd', False):
                gt_bboxes_ignore.append(bbox)
            else:
                score = ann.get('score', 1.0)       # labeled标注score默认为1.0
                gt_bboxes.append(bbox)
                gt_labels.append(self.cat2label[ann['category_id']])
                gt_scores.append(score)
                if 'segmentation' in ann:
                    gt_masks_ann.append(ann['segmentation'])
                else:
                    gt_masks_ann.append([[x1, y1, x2, y1, x2, y2, x1, y2, x1, y1]])

        if gt_bboxes:
            gt_bboxes = np.array(gt_bboxes, dtype=np.float32)
            gt_labels = np.array(gt_labels, dtype=np.int64)
            gt_scores = np.array(gt_scores, dtype=np.float32)
        else:
            gt_bboxes = np.zeros((0, 4), dtype=np.float32)
            gt_labels = np.array([], dtype=np.int64)
            gt_scores = np.array([], dtype=np.float32)

        if gt_bboxes_ignore:
            gt_bboxes_ignore = np.array(gt_bboxes_ignore, dtype=np.float32)
        else:
            gt_bboxes_ignore = np.zeros((0, 4), dtype=np.float32)

        seg_map = img_info['filename'].replace('jpg', 'png')

        ann = dict(
            bboxes=gt_bboxes,
            labels=gt_labels,
            scores=gt_scores,
            bboxes_ignore=gt_bboxes_ignore,
            masks=gt_masks_ann,
            seg_map=seg_map)

        return ann


#Here we import all the necessary libraries, pandas, sklearn and numpy for data pre-processing, keras for neural networks
# and matplotlib in order to plot the obtained data

import pandas as pd
pd.set_option('display.float_format', lambda x: '%.4f' % x)
from keras.models import Sequential
from keras.layers import *
from keras.callbacks import EarlyStopping
from sklearn.metrics import mean_squared_error
from sklearn.metrics import mean_absolute_error
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
import numpy as np

#This first function is specific to the datasets that I used, it takes the data of a given station and creates a list of
# values. NB! Make sure that the file contains all the entries even if they're NaN

def create_list(data, station = 'RIJE99PA'):
    df = data.loc[(data['EG_GH_ID'] == station)]

    df = df.to_numpy()

    df = df[:, 6:38]
    a = np.shape(df)[0]

    val = []
    for i in range(int(a/24)):
        for j in range(31):
            y1 = df[i*24:i*24 + 24, j]
            val.append(y1)
    return np.concatenate((val), axis = None)

#This function takes a list as created using create_list and creates 2 arrays
# one of the arrays contain prediction values going back forecast steps
# the other is a list containing answers that are forecast steps ahead from the last training entry
def create_dataset(dataset, look_back=1, forecast = 1):
    X, Y = [], []
    for i in range(len(dataset)-look_back-forecast):            #creates the data set exapmles, with look_back previous values forecasting forecast hours ahead
        a = dataset[i:(i+look_back)]
        X.append(a)
        Y.append(dataset[i + look_back+forecast-1])

    y3 = np.zeros((len(dataset)-look_back-forecast, look_back + 1))             #gets rid of all the training examples that contain nan
    for j in range(len(dataset)-look_back-forecast):
        y3[j, 0:look_back] = X[j]
        y3[j,look_back] = Y[j]
    return np.array(y3[:,0:look_back]), np.array(y3[:,look_back])

#The first part of this function takes 3 datasets and puts them together
#The second part finds all the entries that are NaN either for prediction data or answers and removes the rows containing
#them
#The output is a ready to process dataset with answers for one of the stations

def create_training_example(TempVal,WindMag,WindDir,TempAns, look_back = 1):
    x1 = np.zeros((TempVal.shape[0], look_back, 3*len(location_set)))
    x1[:,:,0:len(location_set)] = (TempVal)
    x1[:,:,len(location_set):2*len(location_set)] = (WindMag*np.sin((WindDir/180)*np.pi))
    x1[:,:,2*len(location_set):3*len(location_set)] = (WindMag*np.cos((WindDir/180)*np.pi))
    index = []

    for i in range((x1.shape[0])):
        if np.isnan(x1[i]).any() == True or np.isnan(TempAns[i]).any() == True:
            index.append(i)
    index_full = np.linspace(0,x1.shape[0],x1.shape[0]+1)
    index_full = np.delete(index_full, index)
    print(index_full)

    x2 = np.zeros((len(index_full), look_back, 3*len(location_set)))
    x2ans = np.zeros((len(index_full),len(location_set)))
    for j in range(len(index_full)-1):
        x2[j,:,:] = x1[int(index_full[j]),:,:]
        x2ans[j,:] = TempAns[int(index_full[j]),:]
    return x2, x2ans

#Just a list of weather stations that are gonna be used

location_set = ['RIAI99PA','RIAL99MS','RIBA99PA','RIDM99MS','RIDO99MS','RIGU99MS','RIJE99PA','RIME99MS','RIPR99PA','RIRE99MS','RIRU99PA','RISA99PA','RISE99MS','RISI99PA','RIST99PA','RIVE99PA','RIZI99PA','RIZO99MS']

#This function creates a dataset for all the weather stations, it combines create_list and create_dataset functions
#and creates a ready to use dataset with answers to input into a recurrent network. Recurrent networks have a particular
#shape that they need to take in order to be correctly read [samples, timesteps, features]. This function also creates
#it into one. As an output it gives [3D trianing array, 2D answer array]

def create_all_dataset(x, look_back = 1, forecast = 1):
    list1 = create_list(x)
    length = len(list1)
    values = np.zeros((length-forecast-look_back,look_back,len(location_set)))
    answers = np.zeros((length-forecast-look_back, len(location_set)))

    for i in range(len(location_set)):
        values[:,:,i], answers[:,i] = create_dataset(create_list(x,station = location_set[i]), look_back = look_back, forecast = forecast)
    return values, answers

#pandas read_csv function is used to read the csv files

data = pd.read_csv("Avg_air_temperature_2004.04_2018.csv")
data1 = pd.read_csv("Wind_direction_2004.04_2018.csv")
data2 = pd.read_csv("Wind_speed_2004.04_2018.csv")

#NB! In here it's necessary to enter the value how many hours ahead the forecasting is performed

look_back = 60
forecast = 12
features = 3*len(location_set)

#Here one can choose the station for which the plot is going to be made

station = 'RIJE99PA'
id = location_set.index(station)

#Here we create datasets and answers for Temperature, Wind speed and direction

valT, ansT = create_all_dataset(data, look_back = look_back, forecast = forecast)
valWM, ansWM = create_all_dataset(data2, look_back = look_back, forecast = forecast)
valWD, ansWD = create_all_dataset(data1, look_back = look_back, forecast = forecast)

#Here all 3 datasets are merged together into a ready to use dataset

Input_val, Ver_val = create_training_example(valT, valWM, valWD, ansT, look_back = look_back)

print("Input shape:",Input_val.shape)

#Train_test_split function splits the training set into a training set and verification set and shuffles the examples
#To improve the testing accuracy

X_train, X_test, Y_train, Y_test = train_test_split(Input_val, Ver_val, train_size=0.75, shuffle=True)

#Here we hand-pick some unshuffled values for plotting

X_test1 = Input_val[10000:10400,:,:]
Y_test1 = Ver_val[10000:10400]

print("Training set made")

#After this we have finished with the data pre-processing and we define our neural network model

#GRU model that is used for predictions

model = Sequential()
model.add(GRU(200, input_shape=(X_train.shape[1], X_train.shape[2]), return_sequences= True))
model.add(Dropout(0.2))
model.add(GRU(200))
model.add(Dropout(0.2))
model.add(Dense(len(location_set)))
model.compile(loss='mean_squared_error', optimizer='adam')

#The following line fits the training data to the model and compiles the results with the given batch size and the number
#of epochs

history = model.fit(X_train, Y_train, epochs=50, batch_size=500, validation_data=(X_test, Y_test),
                        callbacks=[EarlyStopping(monitor='val_loss', patience=10)], verbose=1, shuffle=True)

model.summary()

#This line saves the model. NB! Make sure to name the model accordingly!!!

model.save('GRU200200NetworkAll12hW.h5')

#Obtaining the loss values

train_loss = history.history['loss']
val_loss = history.history['val_loss']

count_epochs = range(1,len(val_loss)+1)

#Making the predictions  on the ready-made model for both training and test data

train_predict = (model.predict(X_train))
test_predict = (model.predict(X_test))
test1_predict = (model.predict(X_test1))

print('Train Mean Absolute Error:', mean_absolute_error(Y_train[:,id], train_predict[:,id]))
print('Train Root Mean Squared Error:',np.sqrt(mean_squared_error(Y_train[:,id], train_predict[:,id])))
print('Test Mean Absolute Error:', mean_absolute_error(Y_test[:,id], test_predict[:,id]))
print('Test Root Mean Squared Error:',np.sqrt(mean_squared_error(Y_test[:,id], test_predict[:,id])))

#Plotting loss results

plt.plot(count_epochs,train_loss)
plt.plot(count_epochs,val_loss)
plt.legend(['Train loss', 'Validation loss'])
plt.show()

#Plotting validation results

aa=[x for x in range(400)]
plt.plot(aa, Y_test1[:,id], marker='.', label="actual")
plt.plot(aa, test1_predict[:,id], 'r', label="prediction")
plt.tight_layout()
plt.subplots_adjust(left=0.07)
plt.ylabel('Temperature', size=15)
plt.xlabel('Time step', size=15)
plt.legend(fontsize=15)
plt.show()


# import libraries
import sys; sys.path.append('/usr/local/python'); sys.path.append('/usr/local/python')
import numpy as np
import cv2
import os
import yaml
from openpose import pyopenpose as op
import time # for measuring elapsted time in YOLO computation
import argparse # check if image path is provided

#setup and import darknet-yolo
sys.path.append('/usr/local/python') # path for CMUopenpose library
sys.path.append(os.environ['DARKNET_PATH'])
sys.path.append(os.environ['DARKNET_PATH']+'/python')
import darknet as dn

import tensorflow as tf
from feedforward_model import *

def prediction(skele,model_name='basic_model'):
	labels = ['high','med','low']
	session = tf.Session()
	with session as sess:
		new_saver = tf.train.import_meta_graph('model/{}.meta'.format(model_name))
		new_saver.restore(sess, 'model/{}'.format(model_name))
		graph = tf.get_default_graph()
		input_ph = graph.get_tensor_by_name("input_ph:0")
		pred = graph.get_tensor_by_name("pred:0")

		predictions = sess.run([pred],feed_dict={input_ph: skele})
		return labels[predictions[0][0]]

def get_bounding_boxes(detections, box_color):
	boxes = []
	for detection in detections:
		bounds = detection[2]
		box = {
			"box_color" : box_color,
			"class" : detection[0],
			"probability" : detection[1],
			"xmin" : int(bounds[0] - bounds[2]/2),
			"ymin" : int(bounds[1] - bounds[3]/2),
			"xmax" : int(bounds[0] + bounds[2]/2),
			"ymax" : int(bounds[1] + bounds[3]/2)
			}
		boxes.append(box)
	return boxes

def box_overlap(human_box, gun_box):
	# Overlapping rectangles overlap both horizontally & vertically
	x_bool = range_overlap(human_box["xmin"], human_box["xmax"], gun_box["xmin"], gun_box["xmax"])
	y_bool = range_overlap(human_box["ymin"], human_box["ymax"], gun_box["ymin"], gun_box["ymax"])
	return x_bool and y_bool

def range_overlap(a_min, a_max, b_min, b_max):
	# Neither range is completely greater than the other
	return (a_min <= b_max) and (b_min <= a_max)

def imcrop(img, bbox):
   x1, y1, x2, y2 = bbox
   if x1 < 0 or y1 < 0 or x2 > img.shape[1] or y2 > img.shape[0]:
        img, x1, x2, y1, y2 = pad_img_to_fit_bbox(img, x1, x2, y1, y2)
   return img[y1:y2, x1:x2, :]

if __name__ == "__main__":
	parser = argparse.ArgumentParser()
	parser.add_argument("--image_path", help="image_path the string you use here")
	args = parser.parse_args()

	# Setup yolo config
	with open("yolo_config.yml", 'r') as ymlfile:
		if sys.version_info[0] > 2:
			cfg = yaml.load(ymlfile, Loader=yaml.FullLoader)
		else:
			cfg = yaml.load(ymlfile)

	yolo_data = cfg['yolo_data']
	yolo_config = cfg['yolo_config']
	yolo_weights = cfg['yolo_weights']
	pistol_net = dn.load_net(yolo_config,yolo_weights,0)
	pistol_meta = dn.load_meta(yolo_data)

	coco_data = cfg['coco_data']
	coco_config = cfg['coco_config']
	coco_weights = cfg['coco_weights']
	coco_net = dn.load_net(coco_config,coco_weights,0)
	coco_meta = dn.load_meta(coco_data)

	# setup openpose config
	# with open("openpose_config.yml", 'r') as opymlfile:
	# 	if sys.version_info[0] > 2:
	# 		op_cfg = yaml.load(opymlfile, Loader=yaml.FullLoader)
	# 	else:
	# 		op_cfg = yaml.load(opymlfile)
	# Custom Params (refer to include/openpose/flags.hpp for more parameters)
	op_params = dict()
	op_params["model_folder"] = cfg['model_folder']
	op_params["model_pose"] = cfg['model_pose']

	# Starting OpenPose
	opWrapper = op.WrapperPython()
	opWrapper.configure(op_params)
	opWrapper.start()
	datum = op.Datum()

	# Load image
	if args.image_path:
		cv_image = cv2.imread(args.image_path,cv2.IMREAD_COLOR) #load image in cv2
	else:
		cv_image = cv2.imread(cfg['single_image'] ,cv2.IMREAD_COLOR) #load image in cv2
	height = cv_image.shape[0]
	width = cv_image.shape[1]
	# print("original height: {}").format(height)
	# print("original width: {}").format(width)
	box_image = cv_image.copy()
	#resize image to match network dimensions
	skeleton_image = cv2.resize(box_image,(dn.network_width(pistol_net),dn.network_height(pistol_net)),interpolation=cv2.INTER_LINEAR)

	# Detect objects
	time_start = time.clock()
	frame_resized = cv2.resize(box_image,(dn.network_width(pistol_net),dn.network_height(pistol_net)),interpolation=cv2.INTER_LINEAR)
	threat_resized = cv2.resize(box_image,(dn.network_width(pistol_net),dn.network_height(pistol_net)),interpolation=cv2.INTER_LINEAR)
	darknet_image = dn.make_image(dn.network_width(pistol_net),dn.network_height(pistol_net),3)
	dn.copy_image_from_bytes(darknet_image,frame_resized.tobytes())
	pistol_detections = dn.detect_image(pistol_net, pistol_meta, darknet_image, thresh=0.25)
	coco_detections = dn.detect_image(coco_net, coco_meta, darknet_image, thresh=0.25)

	# print("resized height: {}").format(dn.network_height(pistol_net))
	# print("resized width: {}").format(dn.network_width(pistol_net))

	height_scale = float(cv_image.shape[0]) / dn.network_height(pistol_net)
	width_scale = float(cv_image.shape[1]) / dn.network_height(pistol_net)

	# print("height_scale: {}").format(height_scale)
	# print("width_scale: {}").format(width_scale)

	# get bounding boxes
	pistol_boxes = get_bounding_boxes(pistol_detections, (255,0,0))
	coco_boxes = get_bounding_boxes(coco_detections, (0,0,255))
	BBoxes = []
	for box in pistol_boxes:
		BBoxes.append(box)
	# for box in coco_boxes:
	# 	BBoxes.append(box)

	fontFace = cv2.FONT_HERSHEY_DUPLEX
	text_color = (0,0,255)
	fontScale = 1
	text_thickness = 1

	box_thickness = 2

	for box in pistol_boxes:
		# print("pistol_boxes  {}").format(box['class'])
		if box['class'] == "pistol":
			box_color = (255,0,0)
		if box['class'] == "person":
			box_color = (0,0,255)
		upper_left = (box['xmin'],box['ymin'])
		lower_right = (box['xmax'],box['ymax'])
		cv2.rectangle(frame_resized,upper_left,lower_right,box_color,box_thickness)
		cv2.putText(frame_resized,box['class'],upper_left, fontFace, fontScale, box_color,text_thickness,cv2.LINE_AA)

	# box_color = (0,0,255)
	# for box in coco_boxes:
	# 	# print("coco_boxes  {}").format(box['class'])
	# 	if box['class'] == "person":
	# 		upper_left = (box['xmin'],box['ymin'])
	# 		lower_right = (box['xmax'],box['ymax'])
			# cv2.rectangle(frame_resized,upper_left,lower_right,box_color,box_thickness)
			# cv2.putText(frame_resized,box['class'],upper_left, fontFace, fontScale, box_color,text_thickness,cv2.LINE_AA)

	if cfg['show_images']:
		cv2.imshow("frame_resized",frame_resized)
		cv2.waitKey(0)
		cv2.destroyWindow("frame_resized")

	# check overlapping boxes
	person_boxes = []
	pistol_boxes = []
	potential_threats = []
	for box in BBoxes:
		# print("box[class] = {}").format(box)
		if box['class'] == "person":
			person_boxes.append({'xmin':box['xmin'],'ymin':box['ymin'],'xmax':box['xmax'],'ymax':box['ymax']})

		if box['class'] == "pistol":
			pistol_boxes.append({'xmin':box['xmin'],'ymin':box['ymin'],'xmax':box['xmax'],'ymax':box['ymax']})

	for person_box in person_boxes:
		for pistol_box in pistol_boxes:
			if box_overlap(person_box, pistol_box):
				potential_threats.append(person_box)

	person_number = 0
	for person_box in potential_threats:
		person_number +=1
		# generate sub images for skeleton
		# print("crop image (ymin:ymax), (xmin:xmax) : ({}:{}), ({}:{})").format(person_box["ymin"],person_box["ymax"], person_box["xmin"], person_box["xmax"])
		if (person_box["ymin"] < 0):
			person_box["ymin"] = 0
		if (person_box["xmin"] < 0):
			person_box["xmin"] = 0

		if (person_box["ymax"] > skeleton_image.shape[0]):
			person_box["ymax"] = skeleton_image.shape[0]
		if (person_box["xmax"] > skeleton_image.shape[1]):
			person_box["xmax"] = skeleton_image.shape[0]

		# print("crop image (ymin:ymax), (xmin:xmax) : ({}:{}), ({}:{})").format(person_box["ymin"],person_box["ymax"], person_box["xmin"], person_box["xmax"])
		crop_img = skeleton_image[person_box["ymin"]:person_box["ymax"],  person_box["xmin"]:person_box["xmax"]]

		# rescale sub image back to original dimensions
		crop_img_height = crop_img.shape[0]
		crop_img_width = crop_img.shape[1]
		# print("crop_img_width * width_scale : {}*{}={}").format(crop_img_width, width_scale, crop_img_width*width_scale)
		# print("crop_img_height * height_scale : {}*{}={}").format(crop_img_height, height_scale, crop_img_height*height_scale)

		if cfg['show_images']:
			cv2.imshow("crop_img", crop_img)
			cv2.waitKey(0)
			cv2.destroyWindow("crop_img")

		# process sub image for skeleton
		newX,newY = crop_img_width*width_scale, crop_img_height*height_scale
		if (newX > 0) and (newY > 0):
			# process sub image for skeleton
			skel_image = cv2.resize(crop_img,(int(newX),int(newY)))
		else:
			continue

		datum.cvInputData = skel_image
		opWrapper.emplaceAndPop([datum])
		if cfg['show_images']:
			cv2.imshow("skeleton", datum.cvOutputData)
			cv2.waitKey(0)
			cv2.destroyWindow("skeleton")

		x = np.empty((1,9,2))
		rele_dexes = [1,2,3,4,5,6,7,9,12]
		right_elbow = 3
		right_wrist = 4

		skeletons = datum.poseKeypoints
		# print("skeletons.type: {}").format(type(skeletons))
		# print("skeletons.size: {}").format(skeletons.size)
		# print("skeletons: {}").format(skeletons)
		if skeletons.size > 1:
			for skele in skeletons:
				# print("\n\nskele.shape = {}").format(skele.shape)
				if skele[right_elbow].all():
					skele[:,0:2] -= skele[right_elbow,0:2] # set right elbow as origin
					# print("skele.right_elbow = {}").format(skele[right_elbow])
					if skele[right_wrist].all():
						# print("skele.right_wrist = {}").format(skele[right_wrist])
						forearm_len = np.sqrt(skele[right_wrist][0]**2+skele[right_wrist][1]**2) # calculate pixel length of forearm
						# print("skele.forearm_len = {}").format(forearm_len)
						skele[:,0:2] /= forearm_len # scale all joints by forearm length
						skele_x = skele[rele_dexes,0:2]
						# print("skele.shape: {}").format(skele_x.shape)
						# print("skele: {}").format(skele_x)

						x = skele_x.reshape([1,skele_x.shape[0]*skele_x.shape[1]])
						# print("x.shape: {}").format(x.shape)
						# print("x: {}").format(x)

						classification = prediction(x)
						# print("~~~ classification = prediction(x)")
						# Add text to person box
						# cv2.rectangle(frame_resized,upper_left,lower_right,box_color,2)

						textSize = cv2.getTextSize(classification, fontFace, fontScale, text_thickness);
						text_width = textSize[0][0]
						text_height = textSize[0][1]
						text_baseline = textSize[1]
						# print("text_width = {}").format(text_width)
						# print("text_height = {}").format(text_height)
						# print("text_baseline = {}").format(text_baseline)

						person_upper_left = (person_box['xmin'],person_box['ymin'])
						person_lower_right = (person_box['xmax'],person_box['ymax'])
						cv2.rectangle(threat_resized,person_upper_left,person_lower_right,text_color,box_thickness)

						black_box_lower_left = (person_box['xmin']+box_thickness,person_box['ymin']+text_height+text_baseline+box_thickness)
						black_box_upper_right = (person_box['xmin']+box_thickness+text_width,person_box['ymin']+box_thickness)
						cv2.rectangle(threat_resized,black_box_lower_left,black_box_upper_right,(0,0,0),-1)

						text_lower_left = (person_box['xmin']+box_thickness,person_box['ymin']+text_height+box_thickness)
						text_upper_right = (person_box['xmin']+box_thickness+text_width,person_box['ymin']+box_thickness)
						cv2.putText(threat_resized,classification,text_lower_left, fontFace, fontScale, text_color,text_thickness,cv2.LINE_AA)

						skele_image_filename = cfg['save_folder'] + cfg['single_filename'] + "_skele_{}.png".format(person_number)
						skele_image = datum.cvOutputData
						skele_height = skele_image.shape[0]
						skele_width = skele_image.shape[1]

						newX,newY = skele_width*width_scale, skele_height*height_scale
						skele_image_resized = cv2.resize(skele_image,(int(newX),int(newY)))
						# cv2.imwrite(skele_image_filename,skele_image_resized)

	elapsed_time = time.clock() - time_start
	f= open("elapsed_time.txt","a+")
	f.write(("{}\n").format(elapsed_time))

	# return bounding boxes to original image dimensions
	box_image = cv2.resize(threat_resized,(cv_image.shape[1], cv_image.shape[0]),interpolation=cv2.INTER_LINEAR)
	box_image_filename = cfg['save_folder'] + cfg['single_filename'] + "_label.png"
	# print("box_image_filename: {}").format(box_image_filename)
	# cv2.imwrite(box_image_filename,box_image)

	# cv2.imshow("box_image", box_image)
	# cv2.waitKey(0)
	# cv2.destroyWindow("box_image")

	# box_image_filename = cfg['save_folder'] + cfg['single_filename'] + "class.png"
	# print("box_image_filename: {}").format(box_image_filename)
	# cv2.imwrite(cfg['save_folder'] ,img) #wait until all the objects are marked and then write out.

	# imageToProcess = cv2.imread(current_image)
	# datum.cvInputData = imageToProcess
	# opWrapper.emplaceAndPop([datum])


from collections import deque
from sys import stdin

INF = 1000000000000000000000000

#경로 path의 유량을 찾는 함수
#경로 path에서 흐를 수 있는 유량은 각 경로를 구성하는 간선이 가지는 잔여 용량들 중 최솟값
def make_flow(s,t,path):

    c = INF #잔여 용량의 최솟값

    #끝점 t부터 시작해서 역으로 돌아가면서
    #경로 path의 흐를 수 있는 최소 용량을 찾는다
    #즉, 경로 path를 구성하는 간선들의 잔여 용량중 최솟값을 찾는다.
    v = t

    while v != s: #경로를 역으로 돌면서 시작점 s에 도달했다면 반복문 탈출

        #v의 이전 정점 path[v]에서 v로 가는 경로의 잔여 용량이 c보다 작다면,
        #c를 해당 값으로 갱신
        if (capacity[path[v]][v] - flow[path[v]][v]) < c:

            c = capacity[path[v]][v] - flow[path[v]][v]

        v = path[v] #v의 이전 정점이 path[v]

    #해당 경로에서 흐르는 유량을 찾았으면, 경로를 역추적
    #해당 경로를 구성하는 간선들에 유량을 더해주자
    v = t

    while v != s:

        #이전 정점 path[v]에서 v로 가는 유량에 c를 더해주고..
        flow[path[v]][v] += c

        #유량 대칭성
        #path[v]에서 v로 들어가는 유량과 v에서 path[v]로 나가는 유량이 같다
        flow[v][path[v]] -= c
        v = path[v]

    return c

#s에서 t로 가는 '유량이 증가하는 경로'를 찾는 bfs
def bfs(s,t):

    queue = deque([s])

    path = [0]*(n+1)

    while queue:

        u = queue.popleft()

        for v in graph[u]:

            #u에서 v로 용량이 남아(residual) 갈 수 있고 아직 방문하지 않았다면...
            if capacity[u][v] - flow[u][v] > 0 and path[v] == 0:

                queue.append(v)
                path[v] = u #v의 이전 정점이 path[v]

                if v == t: #도착점 sink에 도달했다면....

                    return make_flow(s,t,path) #유량 그래프에 현재 경로에서 흐르는 유량을 기록

    return 0 #더 이상 경로를 찾지 못한 경우

def edmonds_karp(s,t):

    path = 0

    #s에서 t로 가는 유량이 존재하는지 찾을 수 없을때까지 BFS로 찾는다
    while 1:

        c = bfs(s,t)

        if c > 0: #s에서 t로 가는 유량이 존재한다 = 경로가 존재한다

            path += 1

        else: #더 이상 흐를 수 있는 유량이 없다 = 경로가 존재하지 않는다

            break

    return path

n,p = map(int,stdin.readline().split())

graph = [[] for _ in range(n+1)] #유량 네트워크

#a > b로 가는 간선이 존재하지 않으면, 용량은 0으로 정의하기 때문에 초기값을 0으로
capacity = [[0]*(n+1) for _ in range(n+1)] #각 간선의 용량

flow = [[0]*(n+1) for _ in range(n+1)] #각 간선에서 흐르는 유량

for _ in range(p):

    a,b = map(int,stdin.readline().split())

    graph[a].append(b)
    graph[b].append(a) #반대 방향으로 흐르는 유량을 생각해야하므로,
    capacity[a][b] = 1 #a > b로 가는 간선이 존재하는 경우 용량은 1로 제한

source = 1 #유량의 시작점
sink = 2 #유량의 도착점

print(edmonds_karp(source,sink))


# -*- coding: utf-8 -*-
# import sqlite3

# conn = sqlite3.connect("./test.db", check_same_thread=False)

# print(dir(conn))
# cur = conn.cursor() # db ni boshqarish uchun kursor

# sql = """CREATE TABLE students( \
#     id INTEGER,
#     name TEXT,
#     age INTEGER,
#     country TEXT
#     );"""
# sql = """INSERT INTO students
#         VALUES(5,'Jason Statham',45,'UK');
# """
# cur.execute() # sql buyruqlarini yozish va amalga oshirish uchun metod

# sql = """SELECT name , age FROM students;"""
# sql = """SELECT * FROM students WHERE age > 50;"""
# try:
#     data = cur.execute(sql)
#     print(data.fetchall()) # sorovdagi barcha yozuvlarni olish
# except sqlite3.DatabaseError as er:
#     print(er)
# else:
#     conn.commit() # o'zgarishlarni tasdiqlash
#     conn.close() # db boglanishni yopish
#     print("Success !")

# View images
import requests
# query = requests.get(f"https://api.unsplash.com/photos/?client_id={ACCESS_KEY}")
# print(query.json())
ACCESS_KEY = "sqwrqufHCGjCQMhgzRGTboWWBg9ANr2Litex75_frGI"
query = "apple"
url = f"https://api.unsplash.com/search/photos/?page=1&{query}=apple&client_id={ACCESS_KEY}"

print(requests.get(url).json())
# Download image from url
# import wget
# url = 'file url'
# out_filepath = "./photos"
# filename = wget.download(url, out=out_filepath)


import re
import os
import json
import boto3

def limpieza(ticket):
    ticket = re.sub(r'[^A-Za-z0-9]+',' ',ticket)
    tokens = ticket.split()
    tokens = [token for token in tokens if token.isalpha()]
    lista_nombres=[]
    lista_apellidos=[]
    for token in tokens:
        if token in lista_nombres:
            tokens.remove(token)
            continue
        elif token in lista_apellidos:
            tokens.remove(token)
            continue
    return " ".join(tokens)

def lambda_handler(event, context):

    nombre_ticket=event['Records'][0]['s3']['object']['key'].split("/")[-1]
    s3 = boto3.client('s3')
    # print(nombre)
    ticket=s3.get_object(
         Bucket='pruebas-clasificador',
         Key=f'transient/{nombre_ticket}',
    )['Body'].read().decode('utf-8')
    # print(ticket)

    ticket = limpieza(ticket)

    s3.put_object(
         Bucket='pruebas-clasificador',
         Key=f'raw/{nombre_ticket}',
         Body=ticket,
    )

    return None


#!/usr/bin/python2

import time
import BaseHTTPServer
import os
import random
import string
import requests
from urlparse import parse_qs, urlparse

HOST_NAME = '0.0.0.0'
PORT_NUMBER = 9999
# A variável MP3_DIR será construida tendo como base o diretório HOME do usuário + Music/Campainha
# (e.g: /home/usuario/Music/Campainha)
MP3_DIR = os.path.join(os.getenv('HOME'), 'Music', 'Campainha')
VALID_CHARS = set(string.ascii_letters + string.digits + '_.')
CHAVE_THINGSPEAK = 'YOURKEY'
# Salva o arquivo de log no diretório do usuário (e.g: /home/usuário/campainha.log)
ARQUIVO_LOG = os.path.join(os.getenv('HOME'), 'campainha.log')

def filtra(mp3):
    if not mp3.endswith('.mp3'):
        return False
    for c in mp3:
        if not c in VALID_CHARS:
            return False
    return True

def log(msg, output_file=None):
    if output_file is None:
        output_file = open(ARQUIVO_LOG, 'a')

    output_file.write('%s: %s\n' % (time.asctime(), msg))
    output_file.flush()

class MyHandler(BaseHTTPServer.BaseHTTPRequestHandler):
    def do_GET(s):
        s.send_header("Content-type", "text/plain")

        query = urlparse(s.path).query
        if not query:
            s.send_response(404)
            s.end_headers()
            s.wfile.write('Not found')
            return

        components = dict(qc.split('=') for qc in query.split('&'))
        if not 'bateria' in components:
            s.send_response(404)
            s.end_headers()
            s.wfile.write('Not found')
            return

        s.send_response(200)
        s.end_headers()

        s.wfile.write('Tocou')
        s.wfile.flush()

        log("Atualizando thingspeak")
        r = requests.post('https://api.thingspeak.com/update',
            data={'api_key': CHAVE_THINGSPEAK, 'field1': components['bateria']})
        log("Thingspeak retornou: %d" % r.status_code)

        log("Tocando MP3")
        mp3s = [f for f in os.listdir(MP3_DIR) if filtra(f)]
        mp3 = random.choice(mp3s)
        os.system("mpv " + os.path.join(MP3_DIR, mp3))

if __name__ == '__main__':
    server_class = BaseHTTPServer.HTTPServer
    httpd = server_class((HOST_NAME, PORT_NUMBER), MyHandler)
    log("Server Starts - %s:%s" % (HOST_NAME, PORT_NUMBER))
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        pass
    httpd.server_close()
    log("Server Stops - %s:%s" % (HOST_NAME, PORT_NUMBER))


# @Author: Manuel Rodriguez <valle>
# @Date:   27-Jun-2018
# @Email:  valle.mrv@gmail.com
# @Last modified by:   valle
# @Last modified time: 29-Jun-2018
# @License: Apache license vesion 2.0

import websocket
import json
from datetime import datetime
from django.db.models import Q, Count, Sum, F
from django.db.models.fields import DecimalField
from django.core.serializers.json import DjangoJSONEncoder
from django.forms.models import model_to_dict
from django.contrib.sites.shortcuts import get_current_site
from django.conf import settings
from tokenapi.http import JsonResponse
from gestion.models import (Pedidos, Lineaspedido, Ticket, Infmesa, Teclas,
                            Arqueocaja, Efectivo, Ticketlineas, Receptores,
                            Mesasabiertas, Camareros)

def imprimir_pedido(request, id):
    pedido = Pedidos.objects.get(pk=id)
    camareo = Camareros.objects.get(pk=pedido.camarero_id)
    mesa = pedido.infmesa.mesasabiertas_set.get().mesa
    lineas = pedido.lineaspedido_set.values("idart",
                                           "nombre",
                                           "precio",
                                           "pedido_id").annotate(can=Count('idart'))
    receptores = {}
    for l in lineas:
        receptor = Teclas.objects.get(id=l['idart']).familia.receptor
        if receptor.nombre not in receptores:
            receptores[receptor.nombre] = {
                "op": "pedido",
                "hora": pedido.hora,
                "receptor": receptor.nomimp,
                "camarero": camareo.nombre + " " + camareo.apellidos,
                "mesa": mesa.nombre,
                "lineas": []
            }
        l["precio"] = float(l["precio"])
        receptores[receptor.nombre]['lineas'].append(l)

    send_pedidos_ws(request, receptores)
    return JsonResponse({})

def imprimir_desglose(request, id):
    arqueo = Arqueocaja.objects.get(pk=id)
    cambio = arqueo.cambio
    efectivo = arqueo.efectivo_set.all().values("moneda").annotate(can = Sum("can")).order_by("-moneda")
    retirar = arqueo.efectivo_set.all().aggregate(total=Sum(F("can") * F("moneda"), output_field=DecimalField()))['total']
    retirar = float(retirar) - float(cambio)
    lineas_retirada = []
    lineas_cambio = []
    parcial = 0
    for linea in efectivo:
        can = linea["can"]
        moneda = linea["moneda"]
        texto = "moneda" if moneda < 5 else "billete"
        texto = texto + "s" if can > 1 else ""
        if retirar <= parcial:
            if can > 0:
                lineas_cambio.append({"titulo": "Cambio", 'can':can,'tipo':float(moneda), 'texto_tipo': texto })
        elif retirar > ((can * float(moneda)) + parcial):
            parcial = parcial + float((can * moneda))
            if can > 0:
                lineas_retirada.append({"titulo": "Retirar", 'can':can,'tipo':float(moneda), 'texto_tipo': texto })
        else:
            diferencia = retirar - parcial
            can_parcial = int(diferencia/float(moneda))
            parcial = parcial + (can_parcial * float(moneda))
            if can_parcial > 0:
                lineas_retirada.append({"titulo": "Retirar", 'can':can_parcial,'tipo':float(moneda), 'texto_tipo': texto })
            texto = "moneda" if moneda < 5 else "billete"
            texto = texto + "s" if can_parcial > 1 else texto
            if can - can_parcial > 0:
                lineas_cambio.append({"titulo": "Cambio", 'can':can - can_parcial, 'tipo':float(moneda), 'texto_tipo': texto })

    obj_cambio = {
        "op": "desglose",
        "receptor": Receptores.objects.get(nombre='Ticket').nomimp,
        "fecha": datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"),
        "lineas": lineas_cambio
    }

    obj_desglose = {
        "op": "desglose",
        "receptor": Receptores.objects.get(nombre='Ticket').nomimp,
        "fecha": datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"),
        "lineas": lineas_retirada
    }
    send_ticket_ws(request, obj_desglose)
    send_ticket_ws(request, obj_cambio)
    return JsonResponse({})

def preimprimir_ticket(request, id):
    mesa = Mesasabiertas.objects.get(mesa_id=id)
    infmesa = mesa.infmesa
    infmesa.numcopias = infmesa.numcopias + 1
    infmesa.save()
    camareo = infmesa.camarero
    mesa = mesa.mesa
    lineas = infmesa.lineaspedido_set.filter(estado="P")
    lineas = lineas.values("idart", "precio").annotate(can=Count('idart'),
                                                       totallinea=Sum("precio"))
    lineas_ticket = []
    for l in lineas:
        try:
            art = Teclas.objects.get(id=l["idart"])
            l["nombre"] = art.nombre
        except:
            art = infmesa.lineaspedido_set.filter(idart=l["idart"])
            if len(art) > 0:
                l["nombre"] = art[0].linea.nombre
        lineas_ticket.append(l)
    obj = {
        "op": "preticket",
        "fecha": datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"),
        "receptor": Receptores.objects.get(nombre='Ticket').nomimp,
        "camarero": camareo.nombre + " " + camareo.apellidos,
        "mesa": mesa.nombre,
        "numcopias": infmesa.numcopias,
        "lineas": lineas_ticket,
        'total': infmesa.lineaspedido_set.filter(estado="P").aggregate(Total=Sum("precio"))['Total']
    }

    send_ticket_ws(request, obj)
    return JsonResponse({})

def imprimir_ticket(request, id):
    ticket = Ticket.objects.get(pk=id)
    camarero = Camareros.objects.get(pk=ticket.camarero_id)
    lineas = ticket.ticketlineas_set.all().annotate(idart=F("linea__idart"),
                                                    precio=F("linea__precio"))

    lineas = lineas.values("idart",
                           "precio").annotate(can=Count('idart'),
                                              totallinea=Sum("precio"))
    lineas_ticket = []
    for l in lineas:
        try:
            art = Teclas.objects.get(id=l["idart"])
            l["nombre"] = art.nombre
        except:
            art = ticket.ticketlineas_set.filter(linea__idart=l["idart"])
            if len(art) > 0:
                l["nombre"] = art[0].linea.nombre
        lineas_ticket.append(l)

    obj = {
        "op": "ticket",
        "fecha": ticket.fecha + " " + ticket.hora,
        "receptor": Receptores.objects.get(nombre='Ticket').nomimp,
        "camarero": camarero.nombre + " " + camarero.apellidos,
        "mesa": ticket.mesa,
        "lineas": lineas_ticket,
        "num": ticket.id,
        "efectivo": ticket.entrega,
        'total': ticket.ticketlineas_set.all().aggregate(Total=Sum("linea__precio"))['Total']
    }
    send_ticket_ws(request, obj)
    return JsonResponse({})

def reenviarlinea(request, id, idl, nombre):
    import urllib.parse
    nombre = urllib.parse.unquote(nombre).replace("+", " ")
    pedido = Pedidos.objects.get(pk=id)
    camareo = Camareros.objects.get(pk=pedido.camarero_id)
    print(pedido.infmesa)
    mesa = pedido.infmesa.mesasabiertas_set.get().mesa
    lineas = pedido.lineaspedido_set.filter(idart=idl, nombre=nombre).values("idart",
                                           "nombre",
                                           "precio",
                                           "pedido_id").annotate(can=Count('idart'))
    receptores = {}
    for l in lineas:
        receptor = Teclas.objects.get(id=l['idart']).familia.receptor
        if receptor.nombre not in receptores:
            receptores[receptor.nombre] = {
                "op": "urgente",
                "hora": pedido.hora,
                "receptor": receptor.nomimp,
                "camarero": camareo.nombre + " " + camareo.apellidos,
                "mesa": mesa.nombre,
                "lineas": []
            }
        l["precio"] = float(l["precio"])
        receptores[receptor.nombre]['lineas'].append(l)

    send_pedidos_ws(request, receptores)
    return JsonResponse({})

def abrircajon(request):
    obj = {
        "op": "open",
        "fecha": datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"),
        "receptor": Receptores.objects.get(nombre='Ticket').nomimp,
    }
    send_ticket_ws(request, obj)
    return JsonResponse({})

def send_pedidos_ws(request, datos):
    for k, v in datos.items():
        url = ''.join(['ws://', get_current_site(request).domain, '/ws/impresion/', v["receptor"], "/"])
        ws = websocket.create_connection(url)
        ws.send(json.dumps({"message": v}))
        ws.close()

def send_ticket_ws(request, v):
    url = ''.join(['ws://', get_current_site(request).domain, '/ws/impresion/', v["receptor"], "/"])
    ws = websocket.create_connection(url)
    ws.send(json.dumps({"message": v}, cls=DjangoJSONEncoder))
    ws.close()


import gspread
from oauth2client.service_account import ServiceAccountCredentials
from pprint import pprint
import functions

# scope = ["https://spreadsheets.google.com/feeds",'https://https://www.googleapis.com/auth/spreadsheets',"https://www.googleapis.com/auth/drive.file","https://www.googleapis.com/auth/drive"]

scope = [
'https://www.googleapis.com/auth/spreadsheets',
'https://www.googleapis.com/auth/drive'
]

creds = ServiceAccountCredentials.from_json_keyfile_name("creds.json", scope)

client = gspread.authorize(creds)

spread_sheet_name = "excelify-test"
spread_sheet = client.open(spread_sheet_name)
sheet = spread_sheet.sheet1
# checkpoint
print(spread_sheet.worksheets())
data = sheet.get_all_records()
pprint(data)

#checkpoint - print stuff

row = sheet.row_values(3)
col = sheet.col_values(2)
cell = sheet.cell(1,2).value
pprint(row)
pprint(col)
pprint(cell)

sheet2 = client.open_by_url("https://docs.google.com/spreadsheets/d/1K8n-Bcn9uqcuNXfVG4PA2zgDvNIUG5bUyh54CescYVc/edit?ts=5f944b9b#gid=0").sheet1

print(sheet2.get_all_records())

# #checkpoint - insert row
# insertRow = ["hello", 5, "red", "blue"]
# sheet.insert_row(insertRow, 4)

# #checkpoint - delete row
# sheet.delete_row(3)

# #checkpoint - update
# sheet.update_cell(8,2, "changed")
sheet.update("b1", "222")

numrows = sheet.row_count
pprint(numrows)


#!/usr/bin/env python3

"""
Smart wrapper around pmstat
Usage:
  smart-pmstat JOBID [ADDITIONAL ARGUMENTS]

Where JOBID is the SLURM id of a running job.
The script will first recover the corresponding nodes using `squeue`, and will
then start a continuous monitoring of the corresponding nodes with `pmstat`.
Additional arguments are passed on directly to `pmstat`.

Since `pmstat` does not exit until it is interrupted, the same is true for
this script.
"""

import subprocess
import re
import argparse

# parse command line arguments
argparser = argparse.ArgumentParser()
argparser.add_argument("jobid")
argparser.add_argument("extra", nargs=argparse.REMAINDER)
args = argparser.parse_args()

# get the nodelist from `squeue`
cmd = f"squeue -j {args.jobid} -O NodeList:100"
df = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE)

stdout = df.stdout.decode("utf-8")
nodelist = stdout.split("\n")[1].strip()

nodes = []
if "[" in nodelist:
    # split into the individual nodes
    # nodes are generally listed as e.g. m[8202,8342,8349-8350]
    for node in nodelist[2:-1].split(","):
        if "-" in node:
            begend = node.split("-")
            for i in range(int(begend[0]), int(begend[1]) + 1):
                nodes.append(i)
        else:
            nodes.append(int(node))
else:
    # there is only one node
    nodes = [nodelist[1:]]

# run `pmstat` on the nodes
pmstat_args = " ".join([f"-h m{node}" for node in nodes])
pmstat_extra = " ".join(args.extra)
cmd = f"pmstat {pmstat_args} {pmstat_extra}"
df = subprocess.run(cmd, shell=True)


from collections import deque
# 방향 배열
direction = [(-1,-1), (-1,0), (-1,1), (0,-1), (0,1), (1,-1), (1,0), (1,1)]

def BFS(start):
    global max_value
    # 덱생성
    q = deque([start])
    while q:
        check = q.popleft()
        # 방향배열 순회
        for k in range(8):
            ny = check[0] + direction[k][0]
            nx = check[1] + direction[k][1]
            # 범위내 존재할 경우
            if ( 0 <= ny < N ) and ( 0 <= nx < M ):
                # matrix 값이 0인 경우
                if matrix[ny][nx] == 0:
                    # 이전값 +1 로 구현
                    matrix[ny][nx] = matrix[check[0]][check[1]] + 1
                    # q에 값 추가
                    q.append((ny,nx))
                # matrix 값이 아닌 경우
                else:
                    # matrix에 저장된 값이랑 이전값 +1 한 값을 비교하여 작은 경우
                    if matrix[ny][nx] > matrix[check[0]][check[1]] + 1:
                        # q에 추가 및 값 갱신
                        q.append((ny,nx))
                        matrix[ny][nx] = matrix[check[0]][check[1]] + 1

N, M = map(int,input().split())

matrix = [ list(map(int,input().split())) for _ in range(N) ]
one_pos = []
# 1인 지점 찾는 반복문
for i in range(N):
    for j in range(M):
        if matrix[i][j] == 1:
            one_pos.append((i,j))
# 최고 값 갱신 변수
max_value = 0
# one_pos 에서 1인 지점 값 뽑아내며 반복
for l in one_pos:
    BFS(l)
# matrix[i][j] 값이 max_value 보다 크면 값 갱신
for i in range(N):
    for j in range(M):
        if matrix[i][j] > max_value:
            max_value = matrix[i][j]

print(max_value-1)


# Variational AutoEncoder

from tensorflow.keras.layers import Input, Dense, Lambda
from tensorflow.keras.models import Model
from tensorflow.keras import backend as K
from tensorflow.keras import metrics
from tensorflow.keras.datasets import mnist
import numpy as np

# Hyperparameter 설정
batch_size = 100
original_dim = 784
latent_dim = 2
intermediate_dim = 256
epochs = 50
epsilon_std = 1.0

# sampling 함수
def sampling(args: tuple):
    z_mean, z_log_var = args
    epsilon = K.random_normal(shape=(K.shape(z_mean)[0], latent_dim), mean=0., stddev=epsilon_std)

    return z_mean + K.exp(z_log_var / 2) * epsilon

# Encoder 만들기
x = Input(shape=(original_dim,), name='input')
h = Dense(intermediate_dim, activation='relu', name='encoding')(x)
z_mean = Dense(latent_dim, name='mean')(h)
z_log_var = Dense(latent_dim, name='log-variance')(h)

z = Lambda(sampling, output_shape=(latent_dim,))([z_mean, z_log_var])

encoder = Model(x, [z_mean, z_log_var, z], name='encoder')

# Decoder 만들기
input_decoder = Input(shape=(latent_dim,), name='decoder_input')
decoder_h = Dense(intermediate_dim, activation='relu', name='decoder_h')(input_decoder)
x_decoded = Dense(original_dim, activation='sigmoid', name='flat_decoded')(decoder_h)

decoder = Model(input_decoder, x_decoded, name='decoder')

# Model 결합
output_combined = decoder(encoder(x)[2])
vae = Model(x, output_combined)
vae.summary()

# 손실 함수 정의
kl_loss = -0.5 * K.sum(1 + z_log_var - K.exp(z_log_var) - K.square(z_mean), axis=-1)

vae.add_loss(K.mean(kl_loss) / 784.)

vae.compile(optimizer='rmsprop', loss='binary_crossentropy')

# 훈련 / 테스트 세트 분할
(x_train, y_train), (x_test, y_test) = mnist.load_data()

x_train = x_train.astype('float32') / 255.
x_test = x_test.astype('float32') / 255.

x_train = x_train.reshape((len(x_train), np.prod(x_train.shape[1:])))
x_test = x_test.reshape((len(x_test), np.prod(x_test.shape[1:])))

# 훈련
vae.fit(x_train, x_train, shuffle=True, epochs=epochs, batch_size=batch_size)

# 테스트
output = vae.predict(x_test)

import matplotlib.pyplot as plt
import random
fig, ((ax1, ax2, ax3, ax4, ax5), (ax6, ax7, ax8, ax9, ax10)) = plt.subplots(2, 5, figsize=(20,7))

# 이미지 다섯 개를 무작위로 고른다.
random_images = random.sample(range(output.shape[0]), 5)

# 원본(입력) 이미지를 맨 위에 그린다.

for i, ax in enumerate([ax1, ax2, ax3, ax4, ax5]):
    ax.imshow(x_test[random_images[i]].reshape(28,28), cmap='gray')

    if i == 0:
        ax.set_ylabel('INPUT', size=20)

    ax.grid(False)
    ax.set_xticks([])
    ax.set_yticks([])

for i, ax in enumerate([ax6, ax7, ax8, ax9, ax10]):
    ax.imshow(output[random_images[i]].reshape(28,28), cmap='gray')

    if i == 0:
        ax.set_ylabel('OUTPUT', size=20)

    ax.grid(False)
    ax.set_xticks([])
    ax.set_yticks([])

plt.show()


import numpy as np
import open3d as o3d
import os
import torch
import torch.nn as nn
import torch.nn.functional as F

import sys
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
ROOR_DIR = os.path.dirname(BASE_DIR)
sys.path.append(ROOR_DIR)

from utils import batch_transform, angle
from models import gather_points, sample_and_group, weighted_icp

class LocalFeatureFused(nn.Module):
    def __init__(self, in_dim, out_dims):
        super(LocalFeatureFused, self).__init__()
        self.blocks = nn.Sequential()
        for i, out_dim in enumerate(out_dims):
            self.blocks.add_module(f'conv2d_{i}',
                                   nn.Conv2d(in_dim, out_dim, 1, bias=False))
            self.blocks.add_module(f'gn_{i}',
                                   nn.GroupNorm(out_dim // 32, out_dim))
            self.blocks.add_module(f'relu_{i}', nn.ReLU(inplace=True))
            in_dim = out_dim

    def forward(self, x):
        '''

        :param x: (B, C1, K, M)
        :return: (B, C2, M)
        '''
        x = self.blocks(x)
        x = torch.max(x, dim=2)[0]
        return x

class LocalFeatue(nn.Module):
    def __init__(self, radius, K, in_dim, out_dims):
        super(LocalFeatue, self).__init__()
        self.radius = radius
        self.K = K
        self.local_feature_fused = LocalFeatureFused(in_dim=in_dim,
                                                     out_dims=out_dims)

    def forward(self, feature, xyz, permute=False, use_ppf=False):
        '''
        :param feature: (B, N, C1) or (B, C1, N) for permute
        :param xyz: (B, N, 3)
        :return: (B, C2, N)
        '''
        if permute:
            feature = feature.permute(0, 2, 1).contiguous()
        new_xyz, new_points, grouped_inds, grouped_xyz = \
            sample_and_group(xyz=xyz,
                             points=feature,
                             M=-1,
                             radius=self.radius,
                             K=self.K)
        if use_ppf:
            nr_d = angle(feature[:, :, None, :], grouped_xyz)
            ni_d = angle(new_points[..., 3:], grouped_xyz)
            nr_ni = angle(feature[:, :, None, :], new_points[..., 3:])
            d_norm = torch.norm(grouped_xyz, dim=-1)
            ppf_feat = torch.stack([nr_d, ni_d, nr_ni, d_norm], dim=-1) # (B, N, K, 4)
            new_points = torch.cat([new_points[..., :3], ppf_feat], dim=-1)
        xyz = torch.unsqueeze(xyz, dim=2).repeat(1, 1, self.K, 1)
        new_points = torch.cat([xyz, new_points], dim=-1)
        feature_local = new_points.permute(0, 3, 2, 1).contiguous() # (B, C1 + 3, K, M)
        feature_local = self.local_feature_fused(feature_local)
        return feature_local

class OverlapAttentionBlock(nn.Module):
    def __init__(self, channels):
        super(OverlapAttentionBlock, self).__init__()
        self.q_conv = nn.Conv1d(channels, channels // 4, 1, bias=False)
        self.k_conv = nn.Conv1d(channels, channels // 4, 1, bias=False)
        self.q_conv.weight = self.k_conv.weight
        self.v_conv = nn.Conv1d(channels, channels, 1)
        self.trans_conv = nn.Conv1d(channels, channels, 1)
        self.after_norm = nn.GroupNorm(channels // 32, channels)
        self.act = nn.ReLU()

    def forward(self, x, ol_score):
        '''
        :param x: (B, C, N)
        :param ol: (B, N)
        :return: (B, C, N)
        '''
        B, C, N = x.size()
        x_q = self.q_conv(x).permute(0, 2, 1).contiguous() # B, N, C
        x_k = self.k_conv(x) # B, C, N
        x_v = self.v_conv(x)
        attention = torch.bmm(x_q, x_k) # B, N, N
        if ol_score is not None:
            ol_score = torch.unsqueeze(ol_score, dim=-1).repeat(1, 1, N) # (B, N, N)
            attention = ol_score * attention
        attention = torch.softmax(attention, dim=-1)
        attention = attention / (1e-9 + attention.sum(dim=1, keepdims=True))
        x_r = torch.bmm(x_v, attention) # B, C, N
        x_r = self.act(self.after_norm(self.trans_conv(x - x_r)))
        x = x + x_r
        return x

class OverlapAttention(nn.Module):
    def __init__(self, dim):
        super(OverlapAttention, self).__init__()
        self.overlap_attention1 = OverlapAttentionBlock(dim)
        self.overlap_attention2 = OverlapAttentionBlock(dim)
        self.overlap_attention3 = OverlapAttentionBlock(dim)
        self.overlap_attention4 = OverlapAttentionBlock(dim)
        self.conv_fuse = nn.Sequential(
            nn.Conv1d(dim*4, dim*4, kernel_size=1, bias=False),
            nn.GroupNorm(16, dim*4),
            nn.LeakyReLU(negative_slope=0.2))

    def forward(self, x, ol):
        x1 = self.overlap_attention1(x, ol)
        x2 = self.overlap_attention2(x1, ol)
        x3 = self.overlap_attention3(x2, ol)
        x4 = self.overlap_attention4(x3, ol)
        x = torch.cat([x1, x2, x3, x4], dim=1)
        x = self.conv_fuse(x)
        return x

class TFMRModule(nn.Module):
    def __init__(self, args):
        super(TFMRModule, self).__init__()
        self.N1s = [args.train_N1, args.test_N1]
        self.M1s = [args.train_M1, args.test_M1]
        self.top_probs = [args.train_top_prob, args.test_top_prob]
        self.similarity_topks = [args.train_similarity_topk, args.test_similarity_topk]
        self.use_ppf = args.use_ppf
        in_dim = 10 if self.use_ppf else 6
        self.local_features = LocalFeatue(radius=args.radius,
                                          K=args.num_neighbors,
                                          in_dim=in_dim,
                                          out_dims=[128, 256, args.feat_dim])
        self.overlap_attention = OverlapAttention(args.feat_dim)

        # to reuse features for tgt
        self.tgt_info = {'f_y_atten': None,
                         'tgt': None}

    def forward(self, src, tgt, x_ol_score, y_ol_score, train, iter,
                normal_src=None, normal_tgt=None):
        '''

        :param src: (B, N, 3)
        :param tgt: (B, M, 3)
        :param x_ol_score: (B, N)
        :param y_ol_score: (B, M)
        :param train: bool, for hyperparameters selection
        :param iter: int
        :param normal_src: (B, N, 3)
        :param normal_tgt: (B, M, 3)
        :return: src: (B, N2, 3); tgt_corr: (B, N2, 3); icp_weights: (B, N2);
        similarity_max_inds: (B, N2), for overlap evaluation.
        '''
        '''

        :param src:
        :param tgt:
        :param x_ol_score:
        return: correspondences points: (N2, 3), (N2, 3)
        '''
        B, _, _ = src.size()
        if train:
            N1, M1, top_prob = self.N1s[0], self.M1s[0], self.top_probs[0]
            similarity_topk = self.similarity_topks[0]
        else:
            N1, M1, top_prob = self.N1s[1], self.M1s[1], self.top_probs[1]
            similarity_topk = self.similarity_topks[1]

        # point feature extraction for tgt
        if self.use_ppf and normal_src is not None:
            f_x_p = self.local_features(feature=normal_src,
                                        xyz=src,
                                        use_ppf=True)
        else:
            f_x_p = self.local_features(feature=None,
                                        xyz=src)

        f_x_attn = self.overlap_attention(f_x_p, ol=None). \
            permute(0, 2, 1).contiguous()
        f_x_atten = f_x_attn / (torch.norm(f_x_attn, dim=-1,
                                        keepdim=True) + 1e-8)  # (B, N, C)
        x_ol_score, x_ol_inds = torch.sort(x_ol_score, dim=-1, descending=True)
        x_ol_inds = x_ol_inds[:, :N1]
        f_x_atten = gather_points(f_x_atten, x_ol_inds)  # (B, N1, C)
        src = gather_points(src, x_ol_inds)

        if iter == 0:
            # point feature extraction for tgt
            if self.use_ppf and normal_tgt is not None:
                f_y_p = self.local_features(feature=normal_tgt,
                                            xyz=tgt,
                                            use_ppf=True)
            else:
                f_y_p = self.local_features(feature=None,
                                            xyz=tgt)

            f_y_atten = self.overlap_attention(f_y_p, ol=None). \
                permute(0, 2, 1).contiguous()
            f_y_atten = f_y_atten / (torch.norm(f_y_atten, dim=-1,
                                                keepdim=True) + 1e-8)  # (B, M, C)
            y_ol_score, y_ol_inds = torch.sort(y_ol_score, dim=-1, descending=True)
            y_ol_inds = y_ol_inds[:, :M1]
            f_y_atten = gather_points(f_y_atten, y_ol_inds)  # (B, M1, C)
            tgt = gather_points(tgt, y_ol_inds)

            self.tgt_info = {'f_y_atten': f_y_atten,
                             'tgt': tgt}
        else:
            f_y_atten, tgt = self.tgt_info['f_y_atten'], self.tgt_info['tgt']

        similarity = torch.bmm(f_x_atten, f_y_atten.permute(0, 2, 1).contiguous()) # (B, N1, M1)

        # feature matching removal
        N2 = int(top_prob * N1)  # train
        similarity_max = torch.max(similarity, dim=-1)[0]  # (B, N1)
        similarity_max_inds = \
            torch.sort(similarity_max, dim=-1, descending=True)[1][:, :N2]
        src = gather_points(src, similarity_max_inds)

        # generate correspondences
        similarity = gather_points(similarity, similarity_max_inds) # (B, N2, M1)
        x_ol_score = torch.squeeze(
            gather_points(torch.unsqueeze(x_ol_score, dim=-1),
                          similarity_max_inds), dim=-1)
        # find topk points in feature space
        device = similarity.device
        similarity_topk_inds = \
            torch.topk(similarity, k=similarity_topk, dim=-1)[1]  # (B, N2, topk)
        mask = torch.zeros_like(similarity).to(device).detach()
        inds1 = torch.arange(B, dtype=torch.long).to(device). \
            reshape((B, 1, 1)).repeat((1, N2, similarity_topk))
        inds2 = torch.arange(N2, dtype=torch.long).to(device). \
            reshape((1, N2, 1)).repeat((B, 1, similarity_topk))
        mask[inds1, inds2, similarity_topk_inds] = 1
        similarity = similarity * mask

        weights = similarity / \
                  (torch.sum(similarity, dim=-1, keepdim=True) + 1e-8)
        tgt_corr = torch.bmm(weights, tgt)
        icp_weights = x_ol_score[:, :N2]

        return src, tgt_corr, icp_weights, similarity_max_inds