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

import collections
import logging
import numpy as np
import sys
import os
import warnings

import paddle
from paddle.fluid import dygraph, core, framework, unique_name
from paddle.fluid.executor import Executor
from paddle.fluid.param_attr import ParamAttr
from paddle.fluid.initializer import Constant
from paddle.fluid.dygraph.io import INFER_MODEL_SUFFIX, INFER_PARAMS_SUFFIX
from paddle.fluid.io import load_inference_model, save_inference_model
from paddle.fluid.log_helper import get_logger
from . import quant_nn
from .. import quantization_pass
from . import utils

__all__ = ['ImperativeQuantAware']

_logger = get_logger(
    __name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s')

class ImperativeQuantAware(object):
    """
    Applying quantization aware training (QAT) to dgraph model.
    """

    def __init__(self,
                 quantizable_layer_type=['Conv2D', 'Linear'],
                 weight_quantize_type='abs_max',
                 activation_quantize_type='moving_average_abs_max',
                 weight_bits=8,
                 activation_bits=8,
                 moving_rate=0.9,
                 weight_preprocess_layer=None,
                 act_preprocess_layer=None,
                 weight_quantize_layer=None,
                 act_quantize_layer=None):
        """
        The constructor for ImperativeQuantAware.

        Args:
            quantizable_layer_type(list[str | layer]): List the type of
                layers that will be quantized. Default is ['Conv2D', 'Linear'].
            weight_quantize_type(str): quantization type for weights,
                which supports 'abs_max' and 'channel_wise_abs_max'.
            activation_quantize_type(str): quantization type for activations,
                which supports 'abs_max' and 'moving_average_abs_max' now.
                If using 'abs_max' mode, the quantization scale will be
                calculated dynamically each step in both training and testing
                period. If using 'moving_average_abs_max', the static
                quantization scale will be calculated during training and
                used in inference.
            weight_bits(int): quantization bit number for weights, whereas
                the bias is not quantized.
            activation_bits(int): quantization bit number for activations.
            moving_rate(float): the parameter for 'moving_average_abs_max'
                quantization.
            weight_preprocess_layer(paddle.nn.Layer, optional): A paddle
                Layer that defines how to preprocess weight before quantization.
                Using this can quickly test if user's preprocess method works
                or not. The input is non-quantized weight and function returns
                processed weight to be quantized.
                If None, the weight will be quantized directly.
                Default is None.
            act_preprocess_layer(paddle.nn.Layer, optional): A paddle Layer
                that defines how to preprocess activation before quantization.
                Using this can quickly test if user's preprocess method works
                or not. The input is non-quantized activation and function returns
                processed activation to be quantized.
                If None, the activation will be quantized directly.
                Default is None.
            weight_quantize_layer(paddle.nn.Layer, optional): A paddle Layer that
                defines how to quantize weight.
                Using this can quickly test if user's quantization method works or not.
                In this layer, user should both define quantization method and
                dequantization method, that is, the function's input is non-quantized
                weight and returns dequantized weight.
                If None, will use uantization op defined by 'weight_quantize_type'.
                Default is None.
            act_quantize_layer(paddle.nn.Layer, optional): A paddle Layer that defines
                how to quantize activation.
                Using this can quickly test if user's quantization method works or not.
                In this layer, user should both define quantization method and
                dequantization method, that is, the function's input is non-quantized
                activation and returns dequantized activation.
                If None, will use quantization op defined by 'activation_quantize_type'.
                Default is None.

        Note:
            If user sets attribute 'skip_quant' to a Layer that support dynamic
            quantization and sets it to true, the layer would not be quantized
            during training. If this attribute is not sets or the attribute is
            false, the Layer would be qunatized in training.

        Examples 1:
        .. code-block:: python

            import paddle
            from paddle.fluid.contrib.slim.quantization \
                import ImperativeQuantAware
            from paddle.vision.models \
                import resnet

            model = resnet.resnet50(pretrained=True)

            imperative_qat = ImperativeQuantAware(
                weight_quantize_type='abs_max',
                activation_quantize_type='moving_average_abs_max')

            # Add the fake quant logical.
            # The original model will be rewrite.
            # The outscale of outputs in supportted layers would be calculated.
            imperative_qat.quantize(model)

            # Fine-tune the quantized model
            # ...

            # Save quant model for the inference.
            imperative_qat.save_quantized_model(
                layer=model,
                model_path="./resnet50_qat",
                input_spec=[
                    paddle.static.InputSpec(
                    shape=[None, 3, 224, 224], dtype='float32')])

        Examples 2:
        .. code-block:: python

            import paddle
            from paddle.fluid.contrib.slim.quantization \
                import ImperativeQuantAware

            class ImperativeModel(paddle.nn.Layer):
                def __init__(self):
                    super(ImperativeModel, self).__init__()
                    # self.linear_0 would skip the quantization.
                    self.linear_0 = paddle.nn.Linear(784, 400)
                    self.linear_0.skip_quant = True

                    # self.linear_1 would not skip the quantization.
                    self.linear_1 = paddle.nn.Linear(400, 10)
                    self.linear_1.skip_quant = False

                def forward(self, inputs):
                    x = self.linear_0(inputs)
                    x = self.linear_1(inputs)
                    return x

            model = ImperativeModel()
            imperative_qat = ImperativeQuantAware(
                weight_quantize_type='abs_max',
                activation_quantize_type='moving_average_abs_max')

            # Add the fake quant logical.
            # The original model will be rewrite.
            #
            # There is only one Layer(self.linear1) would be added the
            # fake quant logical.
            imperative_qat.quantize(model)

            # Fine-tune the quantized model
            # ...

            # Save quant model for the inference.
            imperative_qat.save_quantized_model(
                layer=model,
                model_path="./imperative_model_qat")
        """
        super(ImperativeQuantAware, self).__init__()

        kwargs = {
            "quantizable_layer_type": quantizable_layer_type,
            "weight_quantize_type": weight_quantize_type,
            "activation_quantize_type": activation_quantize_type,
            "weight_bits": weight_bits,
            "activation_bits": activation_bits,
            "moving_rate": moving_rate,
            "weight_preprocess_layer": weight_preprocess_layer,
            "act_preprocess_layer": act_preprocess_layer,
            "weight_quantize_layer": weight_quantize_layer,
            "act_quantize_layer": act_quantize_layer
        }

        self._quantize_inputs = ImperativeQuantizeInputs(**kwargs)

        self._calc_output_scale = ImperativeCalcOutputScale()

    def quantize(self, model):
        """
        According to weights' and activations' quantization types,
        the model will be added some fake quant ops, such as
        fake_quantize_dequantize_moving_average_abs_max,
        fake_quantize_dequantize_abs_max and so on. At the same time,
        the out_scale value of outputs would be calculated.

        Args:
            model(fluid.dygraph.Layer): the model to be quantized.
        Returns:
            None
        """
        assert isinstance(model, dygraph.Layer), \
            "The model must be the instance of dygraph.Layer."
        self._quantize_inputs.apply(model)
        self._calc_output_scale.apply(model)

    def save_quantized_model(self, layer, path, input_spec=None, **config):
        self._calc_output_scale.save_quantized_model(layer, path, input_spec,
                                                     **config)

class ImperativeQuantizeInputs(object):
    """
    Based on the input params, add the quant_dequant computational
    logic both for activation inputs and weight inputs.
    """

    def __init__(self,
                 quantizable_layer_type=['Conv2D', 'Linear'],
                 weight_quantize_type='abs_max',
                 activation_quantize_type='moving_average_abs_max',
                 weight_bits=8,
                 activation_bits=8,
                 moving_rate=0.9,
                 weight_preprocess_layer=None,
                 act_preprocess_layer=None,
                 weight_quantize_layer=None,
                 act_quantize_layer=None):
        """
        The constructor for ImperativeQuantizeInputs.

        Please refer to the args of ImperativeQuantAware.
        """
        super(ImperativeQuantizeInputs, self).__init__()

        self._quantizable_layer_type = tuple(
            utils.supported_quant_layers_map[layer]
            if layer in utils.supported_quant_layers_map else layer
            for layer in quantizable_layer_type)
        for layer in self._quantizable_layer_type:
            assert not isinstance(layer, str), \
                "%s is unspported to be quantized." % layer

        quantize_type = {
            'abs_max', 'moving_average_abs_max', 'channel_wise_abs_max'
        }
        assert weight_quantize_type in quantize_type, \
            "Unsupported weight_quantize_type: %s. It can only " \
            "be abs_max or moving_average_abs_max or " \
            "channel_wise_abs_max." % weight_quantize_type
        assert activation_quantize_type != 'channel_wise_abs_max' \
            and activation_quantize_type in quantize_type, \
            "Unsupported activation_quantize_type: %s. It can " \
            "only be abs_max or moving_average_abs_max now." \
            % activation_quantize_type

        bits_check = lambda bits: isinstance(bits, int) \
            and bits >= 0 and bits <= 16
        assert bits_check(weight_bits), \
            "weight_bits should be 1, 2,... or 16."
        assert bits_check(activation_bits), \
            "activation_bits should be 1, 2,... or 16."

        layer_check = lambda method: method is None or \
            issubclass(method, dygraph.layers.Layer)
        assert layer_check(weight_preprocess_layer), \
            "weight_preprocess should be nn.Layer."
        assert layer_check(act_preprocess_layer), \
            "act_preprocess should be nn.Layer."
        assert layer_check(weight_quantize_layer), \
            "weight_quantize should be nn.Layer."
        assert layer_check(act_quantize_layer), \
            "act_quantize should be nn.Layer."

        self._kwargs = {
            "weight_quantize_type": weight_quantize_type,
            "activation_quantize_type": activation_quantize_type,
            "weight_bits": weight_bits,
            "activation_bits": activation_bits,
            "moving_rate": moving_rate,
            "weight_pre_layer": weight_preprocess_layer,
            "act_pre_layer": act_preprocess_layer,
            "weight_quant_layer": weight_quantize_layer,
            "act_quant_layer": act_quantize_layer
        }

    def apply(self, model):
        assert isinstance(model, dygraph.Layer), \
            "The model must be the instance of dygraph.Layer."

        for name, layer in model.named_sublayers():
            if not isinstance(layer, self._quantizable_layer_type) \
                or (hasattr(layer, "skip_quant") \
                    and layer.skip_quant == True):
                continue

            # TODO(jc): optimize this module
            last_idx = 0
            idx = 0
            obj = model
            while idx < len(name):
                if (name[idx] == '.'):
                    if hasattr(obj, name[last_idx:idx]):
                        obj = getattr(obj, name[last_idx:idx])
                        last_idx = idx + 1
                idx += 1
            target = name[last_idx:idx]

            quant_layer = self._get_quantized_layer(layer)
            setattr(quant_layer, "layer_name", layer.full_name())
            setattr(obj, target, quant_layer)

    def _get_quantized_layer(self, layer):
        quant_layer_name = None
        for key, value in utils.supported_quant_layers_map.items():
            if isinstance(layer, value):
                quant_layer_name = 'Quantized' + key
                break
        assert quant_layer_name is not None, \
            "The layer %s is unsupported to be quantized." \
            % layer.full_name()

        layer_with_weight = ['QuantizedConv2D', 'QuantizedLinear']
        if quant_layer_name not in layer_with_weight:
            quant_layer_name = 'QuantizedNoweightLayer'

        return quant_nn.__dict__[quant_layer_name](layer, **self._kwargs)

class ImperativeCalcOutputScale(object):
    def __init__(self, moving_rate=0.9):
        """
        Add the logic of calculating and setting output scales of some layers.

        Args:
            moving_rate(float): The decay coefficient of moving average.
                                The default value is 0.9.
        """
        super(ImperativeCalcOutputScale, self).__init__()
        self._moving_rate = moving_rate
        self._register_hook_handle_list = []
        self._out_scale_dict = collections.OrderedDict()

    def apply(self, model):
        """
        Insert the `moving_average_abs_max_scale` op to calculate output
        scale of specific layers in model.

        Args:
            model(fluid.dygraph.Layer): The target model which would be
                calculate the output quantization scale.

        Returns:
            None
        """
        assert isinstance(model, dygraph.Layer), \
            "The model must be the instance of dygraph.Layer."
        for _, layer in model.named_sublayers():
            if self._is_target_layer(layer):
                self._init_scale_params(layer)
                hook_handle = layer.register_forward_post_hook(
                    self._calc_output_scale_hook)
                self._register_hook_handle_list.append(hook_handle)

    def save_quantized_model(self, layer, path, input_spec=None, **config):
        """
        Save the quantized model for the inference.

        Args:
            layer (Layer): The Layer to be saved.
            path (str): The path prefix to save model. The format is
                ``dirname/file_prefix`` or ``file_prefix``.
            input_spec (list[InputSpec|Tensor], optional): Describes the input
                of the saved model's forward method, which can be described by
                InputSpec or example Tensor. If None, all input variables of
                the original Layer's forward method would be the inputs of
                the saved model. Default None.
            **configs (dict, optional): Other save configuration options for
                compatibility. We do not recommend using these configurations,
                they may be removed in the future. If not necessary, DO NOT use
                them. Default None.
                The following options are currently supported:
                (1) output_spec (list[Tensor]): Selects the output targets of
                the saved model. By default, all return variables of original
                Layer's forward method are kept as the output of the saved model.
                If the provided ``output_spec`` list is not all output variables,
                the saved model will be pruned according to the given
                ``output_spec`` list.

        Returns:
            None
        """

        assert isinstance(layer, dygraph.Layer), \
            "The model must be the instance of dygraph.Layer."

        # remove handles and collect output scales
        with dygraph.guard():
            layer.eval()
            for handle in self._register_hook_handle_list:
                handle.remove()
            for _, sub_layer in layer.named_sublayers():
                if self._is_target_layer(sub_layer):
                    if hasattr(sub_layer, "layer_name"):
                        layer_name = sub_layer.layer_name
                    else:
                        layer_name = sub_layer.full_name()
                    if hasattr(sub_layer, "_quant_out_scale"):
                        self._out_scale_dict[layer_name] = float(
                            sub_layer._quant_out_scale)

        # save the quantized model that doesn't have output scales
        paddle.jit.save(layer=layer, path=path, input_spec=input_spec, **config)

        # load static model
        is_dynamic_mode = False
        if paddle.in_dynamic_mode():
            is_dynamic_mode = True
            paddle.enable_static()

        place = core.CUDAPlace(0) if core.is_compiled_with_cuda() \
            else core.CPUPlace()
        exe = Executor(place)

        dirname = os.path.dirname(path)
        basename = os.path.basename(path)
        model_filename = basename + INFER_MODEL_SUFFIX
        params_filename = basename + INFER_PARAMS_SUFFIX
        [inference_program, feed_target_names, fetch_targets] = (
            load_inference_model(
                dirname=dirname,
                executor=exe,
                model_filename=model_filename,
                params_filename=params_filename))

        # set output scales to the static model
        check_behind_op = False
        op_count = 0
        ops_list = [key for key, _ in self._out_scale_dict.items()]
        if len(ops_list) == 0:
            warnings.warn(
                "Warning: No Layer of the model while to be saved contains "
                "the out_threshold attribute, so the generated inference "
                "model would not contain the out_threshold.")
        else:
            # Because the Layer in dygraph may correspond to multiple ops
            # in static program after being saved. To ensure correctness,
            # the outscale collected for output of dygraph Layer can only
            # be set to the last op in the corresponding ops in static program.
            #
            # We can judge the execution order of the ops which corresponding
            # to dygraph Layer by check_behind_op
            forward_op = None
            for block in inference_program.blocks:
                for op in block.ops:
                    if op.type in utils.op_real_in_out_name:
                        if op_count > len(ops_list):
                            warnings.warn(
                                "The number of Layer which has "
                                "out_threshold attribute should be bigger than "
                                "the op in inference model")
                            break
                        if check_behind_op:
                            check_behind_op = False
                            if op.type == "elementwise_add":
                                if self._is_op_matched(ops_list[op_count], op,
                                                       block):
                                    op._set_attr("out_threshold",
                                                 self._out_scale_dict[ops_list[
                                                     op_count]])
                                    op_count += 1
                                    forward_op = None
                                continue
                            else:
                                if forward_op is None:
                                    raise ValueError(
                                        "forward_op should not be None")
                                if self._is_op_matched(ops_list[op_count],
                                                       forward_op, block):
                                    forward_op._set_attr(
                                        "out_threshold", self._out_scale_dict[
                                            ops_list[op_count]])
                                    op_count += 1
                                    forward_op = None

                        if op.type in ["conv2d", "depthwise_conv2d", "matmul"]:
                            check_behind_op = True
                            forward_op = op
                            continue
                        if op_count >= len(ops_list):
                            warnings.warn(
                                "The number of Layer which has out_threshold attribute should be bigger than the op in inference model"
                            )
                            break
                        if self._is_op_matched(ops_list[op_count], op, block):
                            op._set_attr(
                                "out_threshold",
                                self._out_scale_dict[ops_list[op_count]])
                            op_count += 1

            self._set_skip_quant_attr(inference_program)

        # save the final quantized model that has output scales
        save_inference_model(
            dirname=dirname,
            feeded_var_names=feed_target_names,
            target_vars=fetch_targets,
            executor=exe,
            main_program=inference_program.clone(),
            model_filename=model_filename,
            params_filename=params_filename)

        if is_dynamic_mode:
            paddle.disable_static()

    def _is_target_layer(self, layer):
        return isinstance(layer, utils.out_scale_layers_list) \
            or 'quantized_' in layer.full_name()

    def _init_scale_params(self, layer, name=None):
        """
        Init the scale params for calculating output scales and save them in the
        target layer.
        After the users define the dygraph model, the hooks for calculating output
        scales will not execute immediately. If the users load parameters form
        checkpoint and save the quantized inference model immediately, the inference
        model would not be saved successfully. Beacuse the dygraph_to_static requires
        that the parameters created in __init__, but the uniqueness of hook make it
        impossible to create parameters in __init__. To avoid this mistake, we define
        the scale parameters in the beginning instead of hook.
        """

        def _create_param(in_layer, first_name, last_name, dtype):
            prefix = '{}.{}'.format(first_name, last_name) \
                if first_name else 'outscale.{}'.format(last_name)
            attr = ParamAttr(
                name=unique_name.generate(prefix),
                initializer=Constant(1),
                trainable=False)
            param = in_layer.create_parameter(shape=[1], attr=attr, dtype=dtype)
            return param

        dtype = layer._dtype if layer._dtype is not None else "float32"
        if dtype not in ["float32", "float64"]:
            return

        layer._quant_out_scale = _create_param(layer, name, "scale", dtype)
        layer._quant_out_scale.stop_gradient = True

        layer._quant_out_state = _create_param(layer, name, "state", dtype)
        layer._quant_out_state.stop_gradient = True

        layer._quant_out_accum = _create_param(layer, name, "accum", dtype)
        layer._quant_out_accum.stop_gradient = True

    # Judge whether the op in program matches the Layer in dynamic model
    def _is_op_matched(self, layer_name, op, block):
        output_var_names = quantization_pass._get_op_output_var_names(op)
        for output_var_name in output_var_names:
            output_var_tensor = block.var(output_var_name)
            if output_var_tensor.dtype not in [
                    core.VarDesc.VarType.FP64, core.VarDesc.VarType.FP32
            ]:
                return False

        # Because the naming styles of static and dynamic graph are different,
        # in order to avoid mistakes, we unify the name here.
        op_type = output_var_names[0].split(".")[0]
        op_type = op_type.rsplit("_", 1)[0]
        if op_type == 'depthwise_conv2d':
            op_type = 'conv2d'
        if 'prelu' in op_type:
            op_type = op_type.replace('prelu', 'p_re_lu')
        if 'relu' in op_type:
            op_type = op_type.replace('relu', 're_lu')
        return op_type in layer_name

    def _set_skip_quant_attr(self, program):
        block = program.global_block()
        for op in block.ops:
            if self._is_skip_quant_op(block, op):
                op._set_attr("skip_quant", True)

    def _is_skip_quant_op(self, block, in_op):
        """
        The input op should be skipped quantization.
        1. the type of input op should be conv2d, depthwise_conv2d or matmul
        2. the previous ops of the input op are not fake_quantize_dequantize ops
        """

        def _find_previous_op(block, var_name):
            for op in block.ops:
                if var_name in op.output_arg_names:
                    return op

        target_op_types = ["conv2d", "depthwise_conv2d", "matmul"]
        if in_op.type not in target_op_types:
            return False

        previous_ops = [_find_previous_op(block, arg_name) \
            for arg_name in in_op.input_arg_names]
        return any(op is not None and op.type not in utils.fake_quantize_dequantize_types \
            for op in previous_ops )

    def _calc_output_scale_hook(self, layer, input, output):
        """
        Create the MovingAverageAbsMaxScale layer for the target layer if needed.
        Execute MovingAverageAbsMaxScale layer to calculate the output scale.
        """
        assert isinstance(output, (core.VarBase, framework.Variable)), \
            "Multiple outputs are not currently supported in ImperativeOutScale."

        fp_types = [core.VarDesc.VarType.FP32, core.VarDesc.VarType.FP64]
        if output.dtype in fp_types:
            if not hasattr(layer, "_out_scale"):
                self._out_scale = quant_nn.MovingAverageAbsMaxScale(
                    layer, output.name, self._moving_rate, output.dtype)
            # TODO (jc): consider the ops that have several outputs
            self._out_scale(output)


import os
from prepare_data.dtm import build_dtm, gen_agl
from prepare_data.pointcloud_util import room2blocks
import torch
from data import FugroDataset, collate_pcs
from model import DGCNN
import numpy as np
from torch.utils.data import DataLoader
from util import *
from prepare_data.process_data import load_pointcloud, save_las_pointcloud
import sklearn.metrics as metrics
from torch.nn.functional import softmax
from tqdm import tqdm

UNCLASSIFIED = 31

def test(k, io,
            data_dir = "",
            num_points = 7000,
            num_classes = 5,
            num_features = 4,
            test_batch_size = 8,
            dropout = 0.5,
            emb_dims = 1024,
            use_all_points = False,
            cuda = False,
            min_class_confidence = 0.8,
            model_label = "dgcnn_model",
            num_threads = 8,
            num_interop_threads = 2,
            model_root = "checkpoints/dgcnn",
            pred_dir = "predict"):
    """Perform inference using a DGCNN model

    Args:
        k (int): Number of neighbours to calculate in feature spaces
        io (IOStream): Stream where log data is sent to
        data_dir (str, optional): Directory containing the dataset in NPY format. Defaults to "".
        num_points (int, optional): Number of points to sample from each block. Defaults to 7000.
        num_classes (int, optional): Number of classes to train on. Defaults to 5.
        test_batch_size (int, optional): Number of test samples in each batch. Defaults to 8.
        dropout (float, optional): Dropout probability for dropout layer in model. Defaults to 0.5.
        emb_dims (int, optional): Dimensions to embed the global feature space into. Defaults to 1024.
        use_all_points (bool, optional): Whether to use all points in a block or to subsample. Defaults to False.
        cuda (bool, optional): Whether to use CUDA device for training. Defaults to False.
        min_class_confidence (float, optional): Minimum confidence value for the model to label a point as belonging to a class
        model_label (str, optional): Label to assign to the model for IO purposes. Defaults to "dgcnn_model".
        num_threads (int, optional): Number of primary threads for PyTorch to use in training. Defaults to 8.
        num_interop_threads (int, optional): Number of secondary threads for PyTorch to use for training. Defaults to 2.
        model_root (str, optional): Directory containing saved model files. Defaults to "checkpoints/dgcnn".
        pred_dir (str, optional): Directory to output the predictions into
    """
    DUMP_DIR = pred_dir
    torch.set_num_threads(num_threads)
    torch.set_num_interop_threads(num_interop_threads)
    all_true_cls = []
    all_pred_cls = []
    all_true_seg = []
    all_pred_seg = []
    for test_area in [1]:
        test_area = str(test_area)
        if (test_area == 'all') or (test_area == test_area):
            dataset = FugroDataset(split='test', data_root=data_dir, num_point=num_points,
                                    use_all_points=use_all_points)
            test_loader = DataLoader(dataset,
                                        batch_size=test_batch_size,
                                        shuffle=False,
                                        drop_last=False,
                                        collate_fn = collate_pcs)

            room_idx = np.array(dataset.room_idxs)

            fout_data_label = []
            true_data_labels = []
            for room_id in np.unique(room_idx):
                out_data_label_filename = 'Area_%s_room_%d_pred_gt.txt' % (test_area, room_id)
                out_data_label_filename = os.path.join(DUMP_DIR, out_data_label_filename)
                out_true_label_filename = 'Area_%s_room_%d_true_labels.txt' % (test_area, room_id)
                out_true_label_filename = os.path.join(DUMP_DIR, out_true_label_filename)
                if os.path.isfile(out_data_label_filename):
                    os.remove(out_data_label_filename)
                if os.path.isfile(out_true_label_filename):
                    os.remove(out_true_label_filename)
                fout_data_label.append(open(out_data_label_filename, 'a'))
                true_data_labels.append(open(out_true_label_filename, 'a'))

            device = torch.device("cuda" if cuda else "cpu")

            io.cprint('Start overall evaluation...')

            model = DGCNN(num_classes, num_features, k, dropout = dropout, emb_dims = emb_dims, cuda = cuda)

            if cuda:
                checkpoint = torch.load(os.path.join(model_root, '%s.t7' % model_label))
            else:
                checkpoint = torch.load(os.path.join(model_root, '%s.t7' % model_label), map_location=torch.device('cpu'))

            print("model: ", os.path.join(model_root, '%s.t7' % model_label))
            model.load_state_dict(checkpoint['model_state_dict'])
            model = model.eval()

            count_parameters(model)

            io.cprint('%s.t7 restored.' % model_label)

            test_acc = 0.0
            count = 0.0
            test_true_cls = []
            test_pred_cls = []
            test_true_seg = []
            test_pred_seg = []

            io.cprint('Start testing ...')
            num_batch = 0

            with tqdm(test_loader, desc = "Testing") as t:
                for data, seg, centers in test_loader:
                    data, seg = data.to(device), seg.to(device)
                    data = data.permute(0, 2, 1).float()
                    batch_size = data.size()[0]

                    seg_pred = model(data)
                    seg_pred = seg_pred.permute(0, 2, 1).contiguous()
                    seg_pred = softmax(seg_pred, dim = 2)
                    vals, pred = seg_pred.max(dim = 2)
                    # vals = seg_pred.amax(dim = 2)
                    pred[torch.where(vals < min_class_confidence)] = UNCLASSIFIED
                    seg_np = seg.cpu().numpy()
                    pred_np = pred.detach().cpu().numpy()
                    test_true_cls.append(seg_np.reshape(-1))
                    test_pred_cls.append(pred_np.reshape(-1))
                    test_true_seg.append(seg_np)
                    test_pred_seg.append(pred_np)

                    # write prediction results
                    for batch_id in range(batch_size):
                        pts = data[batch_id, :, :]

                        pts = pts.permute(1, 0).float()
                        l = seg[batch_id, :]
                        pred_ = pred[batch_id, :]
                        logits = seg_pred[batch_id, :, :]
                        # compute room_id
                        room_id = room_idx[num_batch + batch_id]
                        for i in range(pts.shape[0]):
                            fout_data_label[room_id].write('%f %f %f %d\n' % (
                                (pts[i, 0] + centers[batch_id, 0]), (pts[i, 1] + centers[batch_id, 1]), pts[i, 2],
                                pred_[i]))

                            true_data_labels[room_id].write("%d\n" % l[i])
                    num_batch += batch_size
                    torch.cuda.empty_cache()

                    balanced_acc = metrics.balanced_accuracy_score(seg_np.reshape(-1), pred_np.reshape(-1))
                    acc = metrics.accuracy_score(seg_np.reshape(-1), pred_np.reshape(-1))
                    t.set_postfix(A = acc, BA = balanced_acc)
                    t.update()

            for room_id in np.unique(room_idx):
                fout_data_label[room_id].close()

            test_true_cls = np.concatenate(test_true_cls)
            test_pred_cls = np.concatenate(test_pred_cls)
            test_acc = metrics.accuracy_score(test_true_cls, test_pred_cls)
            avg_per_class_acc = metrics.balanced_accuracy_score(test_true_cls, test_pred_cls)

            outstr = 'Test :: test area: %s, test acc: %.6f, test avg acc: %.6f' % (test_area,
                                                                                test_acc,
                                                                                avg_per_class_acc,)
            io.cprint(outstr)

            # calculate confusion matrix
            conf_mat = metrics.confusion_matrix(test_true_cls, test_pred_cls)
            io.cprint('Confusion matrix:')
            io.cprint(str(conf_mat))

            all_true_cls.append(test_true_cls)
            all_pred_cls.append(test_pred_cls)
            all_true_seg.append(test_true_seg)
            all_pred_seg.append(test_pred_seg)

    if test_area == 'all':
        all_true_cls = np.concatenate(all_true_cls)
        all_pred_cls = np.concatenate(all_pred_cls)
        all_acc = metrics.accuracy_score(all_true_cls, all_pred_cls)
        avg_per_class_acc = metrics.balanced_accuracy_score(all_true_cls, all_pred_cls)
        all_true_seg = np.concatenate(all_true_seg, axis=0)
        all_pred_seg = np.concatenate(all_pred_seg, axis=0)
        all_ious = calculate_sem_IoU(all_pred_seg, all_true_seg, num_classes)
        outstr = 'Overall Test :: test acc: %.6f, test avg acc: %.6f, test iou: %.6f' % (all_acc,
                                                                                avg_per_class_acc,
                                                                                np.mean(all_ious))
        io.cprint(outstr)

def test_args(args, io):
    """Test a DGCNN model using command line args

    Args:
        args (ArgumentParser): Set of command line arguments and their inputs
        io (IOStream): Stream where log data is sent to
    """
    test(
        args.k,
        io,
        data_dir = args.data_dir,
        num_points = args.num_points,
        block_size = args.block_size,
        num_classes = args.num_classes,
        num_features = args.num_features,
        test_batch_size = args.test_batch_size,
        dropout = args.dropout,
        emb_dims = args.emb_dims,
        use_all_points = args.use_all_points,
        cuda = args.cuda,
        min_class_confidence = args.min_class_confidence,
        model_label = args.model_label,
        num_threads = args.num_threads,
        num_interop_threads = args.num_interop_threads,
        model_root = args.model_root,
        pred_dir = args.test_visu_dir
    )

def predict(k, io, pointcloud_file,
            pred_pointcloud_file,
            num_points = 7000,
            block_size = 30.0,
            num_classes = 5,
            num_features = 4,
            dropout = 0.5,
            emb_dims = 1024,
            cuda = False,
            min_class_confidence = 0.8,
            model_label = "dgcnn_model",
            model_root = "checkpoints/dgcnn"):
    """Perform inference on a local pointcloud file

    Args:
        k (int): Number of neighbours to search for
        io (IOStream): Stream where log data is sent to
        pointcloud_file (str): Name of pointcloud file to classify
        pred_pointcloud_file (str): Name of prediction pointcloud file to save
        num_points (int, optional): Number of points to subsample in each block. Defaults to 7000.
        block_size (float, optional): Size of blocks to divide pointcloud into. Defaults to 30.0.
        num_classes (int, optional): Number of classes for the model to consider. Defaults to 5.
        num_features (int, optional): Number of point features fed into the model. Defaults to 4.
        dropout (float, optional): Dropout probability of the model. Defaults to 0.5.
        emb_dims (int, optional): Dimension of space that the model embeds local and global features to during inference. Defaults to 1024.
        cuda (bool, optional): Whether to send data to the GPU. Defaults to False.
        min_class_confidence (float, optional): Minimum value of probability for model to classify point as a class. Defaults to 0.8.
        model_label (str, optional): Label of model being used for inference. Defaults to "dgcnn_model".
        model_root (str, optional): Directory containing model being used for inference. Defaults to "checkpoints/dgcnn".

    Returns:
        data (ndarray): Array of pointcloud data
        preds (ndarray): Vector of classification labels output by the model
    """

    model = DGCNN(num_classes, num_features, k, dropout = dropout, emb_dims = emb_dims, cuda = cuda)

    if cuda:
        checkpoint = torch.load(os.path.join(model_root, '%s.t7' % model_label))
    else:
        checkpoint = torch.load(os.path.join(model_root, '%s.t7' % model_label), map_location=torch.device('cpu'))

    print("model: ", os.path.join(model_root, '%s.t7' % model_label))
    model.load_state_dict(checkpoint['model_state_dict'])
    model = model.eval()

    count_parameters(model)

    io.cprint('%s.t7 restored.' % model_label)

    data, labels = load_pointcloud(pointcloud_file)

    dtm = build_dtm(data)
    agl = gen_agl(dtm, data)

    data = np.hstack((data, np.reshape(agl, (len(agl), 1))))

    # block_data, _ = pointcloud_util.room2blocks(data, labels, num_points, block_size = block_size, stride = block_size, random_sample =False, use_all_points=True)
    block_data, _ = room2blocks(data, labels, num_points,
                                block_size = block_size,
                                stride = block_size,
                                random_sample =False,
                                use_all_points=False)

    preds = np.array([])
    data = []
    n = 1
    with tqdm(block_data, desc = "Classifying") as t:
        for X in block_data:
            np.savetxt("data{}.txt".format(n), X)
            if len(data) == 0:
                data = X
            else:
                data = np.vstack((data, X))

            x_lb = np.amin(X[:, 0])
            y_lb = np.amin(X[:, 1])

            X -= np.array([x_lb, y_lb, 0, 0])

            n += 1
            X = X[:, :, np.newaxis]
            X = torch.tensor(X)
            X = X.permute(2, 1, 0).float()
            logit_pred = model(X)
            logit_pred = logit_pred.permute(0, 2, 1).contiguous()
            logit_pred = softmax(logit_pred, dim = 2)
            probs, pred = logit_pred.max(dim = 2)
            pred[torch.where(probs < min_class_confidence)] = UNCLASSIFIED
            pred = pred.detach().cpu().numpy()
            pred = pred.reshape(pred.shape[1], pred.shape[0])
            X = X.permute(2, 1, 0)
            X = X.detach().cpu().numpy()

            X = X[:, :, 0]

            X += np.array([x_lb, y_lb])

            if len(preds) == 0:
                preds = pred
            else:
                preds = np.vstack((preds, pred))

            save_las_pointcloud(X, pred, "pc_pred_{}.las".format(n))

            t.update()

    save_las_pointcloud(data, preds, pred_pointcloud_file)

    return data, preds


# -*- coding: utf-8 -*-
import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), '../private'))
sys.path.append(os.path.join(os.path.dirname(__file__), '../common'))
sys.path.append(os.path.join(os.path.dirname(__file__), '../XM'))

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
#plt.style.use('seaborn-whitegrid')
from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec

from CandleChart import CandleChart, makeFig
from Graph import Graph
import Filters
from CalendarTime import Now, DTime, DeltaDay, DeltaHour, Today
from TimeSeries import TimeSeries, OHLCV, OHLC, DATA_TYPE_XM
from datetime import datetime
from analyzers import adf

from MT5Bind import MT5Bind, jstTime

def test1():
    server = MT5Bind('US30Cash')
    t0 = jstTime(2019, 3, 5, 20, 0)
    t1 = jstTime(2019, 3, 6, 8, 0)
    data = server.scrapeRange('M5', t0, t1)

    ts = TimeSeries(data, DATA_TYPE_XM, names=OHLCV, index=[1, 2, 3, 4, 5])
    df = ts.toDataFrame()
    df.to_csv('./dji_M5.csv', index=False)

    fig, ax = makeFig(1, 1, (15, 5))
    graph = CandleChart(ax, ts.time, 'MDHM')
    graph.plotOHLC([ts.dic['open'], ts.dic['high'], ts.dic['low'], ts.dic['close']])

    data = server.scrapeRange('M1', t0, t1)
    ts = TimeSeries(data, DATA_TYPE_XM, names=OHLCV, index=[1, 2, 3, 4, 5])
    df = ts.toDataFrame()
    df.to_csv('./dji_M1.csv', index=False)
    server.close()

def breakout(signal, ts, param, stops):
    BUY = 1
    SELL = -1
    CLOSE = 2
    END = 3

    BAND_IN = 0
    BAND_UPPER = 1
    BAND_LOWER = -1
    date = ts.time[0]
    #print(date)

    ohlcv = ts.dataList(OHLCV)
    [buy_threshold, sell_threshold, ma_window, limit] = param
    level_max = len(stops)
    level = 0

    ma = Filters.movingAverage2(signal, ma_window)
    #dif = Filters.dif(ma, 5)

    flag = []
    status = 0
    matches = []
    indices = []
    profit = 0

    count = -1
    longOrShort = None
    prices = []

    if ma_window > 1:
        source = ma
    else:
        source = signal

    signal_state = BAND_IN
    for t, s, o, h, l, c in zip(ts.time, source, ohlcv[0], ohlcv[1], ohlcv[2], ohlcv[3]):
        count += 1
        #if t.hour <= 20 and t.hour >= 8:
        #    continue

        # チェック取引回数リミット
        if len(matches) > limit:
            status = END
            flag.append(0)
            continue

        # エントリ信号チェック
        triggered = False
        if s >= buy_threshold and signal_state == BAND_IN:
            signal_state = BAND_UPPER
            triggered = True
        elif s <= sell_threshold and signal_state == BAND_IN:
            signal_state = BAND_LOWER
            triggered = True

        if status == 0:
            # エントリ
            [stop_profit, stop_loss] = stops[level]
            if triggered and signal_state == BAND_UPPER and (c - o) > 10.0:
                stop_profit_price = c + stop_profit
                stop_loss_price = c + stop_loss
                flag.append(BUY)
                status = BUY
                longOrShort = BUY
                prices = [c]
                indices = [count]
            elif triggered and signal_state == BAND_LOWER and (c - o) < -10.0:
                stop_profit_price = c - stop_profit
                stop_loss_price = c - stop_loss
                flag.append(SELL)
                status = SELL
                longOrShort = SELL
                prices = [c]
                indices = [count]
            else:
                flag.append(0)
        elif status == BUY:
            profit = c - prices[0]
            if stop_profit_price >= l and stop_profit_price <= h:
                while level < level_max:
                    [stop_profit, stop_loss] = stops[level]
                    stop_profit_price = prices[0] + stop_profit
                    stop_loss_price = prices[0] + stop_loss
                    if h > stop_profit_price:
                        break
                    level += 1
                if stop_profit_price >= l and stop_profit_price <= h:
                    flag.append(CLOSE)
                    prices.append(stop_profit_price)
                    indices.append(count)
                    matches.append([longOrShort, prices, indices])
                    longOrShort = 0
                    status = 0
                    level = 0
                else:
                    flag.append(0)
            elif stop_loss_price >= l and stop_loss_price <= h:
                prices.append(stop_loss_price)
                flag.append(CLOSE)
                indices.append(count)
                matches.append([longOrShort, prices, indices])
                status = 0
                longOrShort = 0
                level = 0
        elif status == SELL:
            if stop_profit_price >= l and stop_profit_price <= h:
                while level < level_max:
                    [stop_profit, stop_loss] = stops[level]
                    stop_profit_price = prices[0] - stop_profit
                    stop_loss_price = prices[0] - stop_loss
                    if l > stop_profit_price:
                        break
                    level += 1
                if stop_profit_price >= l and stop_profit_price <= h:
                    flag.append(CLOSE)
                    prices.append(stop_profit_price)
                    indices.append(count)
                    matches.append([longOrShort, prices, indices])
                    status = 0
                    longOrShort = 0
                    level = 0
                else:
                    flag.append(0)
            elif stop_loss_price >= l and stop_loss_price <= h:
                prices.append(stop_loss_price)
                flag.append(CLOSE)
                indices.append(count)
                matches.append([longOrShort, prices, indices])
                status = 0
                longOrShort = 0
                level = 0
        else:
            flag.append(0)

    if len(prices) == 1:
        prices.append(ohlcv[3][-1])
        indices.append(len(ohlcv[0]) - 1)
        matches.append([longOrShort, prices, indices])

    profits = []
    for [lors, pr, idx] in matches:
        if lors == BUY:
            profit = pr[1] - pr[0]
        else:
            profit = pr[0] - pr[1]
        profits.append(profit)
    return (profits, matches, flag)

def purchasePower(ts):
    ohlcv = ts.dataList(OHLCV)
    power = []
    psum = []
    s = 0.0
    for o, h, l, c, v in zip(ohlcv[0], ohlcv[1], ohlcv[2], ohlcv[3], ohlcv[4]):
        p = v * (c - o)
        power.append(p)
        s += p
        psum.append(s)
    return (power, psum)

def drawGraph(title, ts, date, power, psum, result, param):
    op = ts.data('open')
    cl = ts.data('close')

    [buy_threshold, sell_threshold, ma_window, limit] = param

    width = ts.length / 8
    [minvalue, maxvalue] = ts.minmax(OHLC)
    maxmin = maxvalue - minvalue
    height1 = maxmin / 100
    height2 = (np.max(psum) - np.min(psum)) / 50000
    height = height1 + height2
    ratio = height2 / height1

    fig1 = plt.figure(figsize=(width, height))
    grid = GridSpec(nrows=2, ncols=1, height_ratios=[1, ratio])
    grid1 = GridSpecFromSubplotSpec(nrows=1, ncols=1, subplot_spec=grid[0, 0])
    ax1 = fig1.add_subplot(grid1[:, :])
    grid2 = GridSpecFromSubplotSpec(nrows=1, ncols=1, subplot_spec=grid[1, 0])
    ax2 = fig1.add_subplot(grid2[:, :])

    graph1 = CandleChart(ax1, ts.time, 'HM')
    graph1.plotOHLC([op, ts.data('high'), ts.data('low'), cl])
    title_str = title + '    '  + date.strftime('%Y-%m-%d (%A)') + ' ~   Range:' + "{:.1f}".format(maxmin)
    graph1.setTitle(title_str, '', '')
    #graph1.bars(volume, colors, np.max(volume) * 5, 0.01)
    graph1.grid()

    ma = Filters.movingAverage(psum, ma_window)
    dif = Filters.dif(ma, 5)
    graph2 = CandleChart(ax2, ts.time, 'HM')
    graph2.plot(psum, 0, 1)
    graph2.plot(ma, 1, 1)
    #graph2.plot(dif, 7, 1)
    graph2.drawLegend([{'label': 'PurchasePower', 'color':'red'}, {'label': 'MA' + str(ma_window), 'color':'blue'}], None)
    graph2.grid()

    graph2.hline(param[0:2], ['orange', 'violet'], 2)
    profits, matches, flag = result
    profit_str = ''
    for profit, match in zip(profits, matches):
        longOrShort, prices, indices = match
        graph1.box([graph1.time[indices[0]], graph1.time[indices[1]]], prices, 0, 0.2)
        graph1.point([graph1.time[indices[0]], prices[0]], 'o', 'blue', 0.7, 80)
        graph1.point([graph1.time[indices[1]], prices[1]], 'o', 'orange', 0.7, 80)
        profit_str += str(profit) + '  '
    if len(profits) > 0:
        graph1.text(graph1.time[0], graph1.yRange()[1] - 40, 'Profit: ' + profit_str, 'blue', 16)

    #if should_save:
    #    path = './output/' + str(num).zfill(5) +  title + '_' + date.strftime('%Y-%m-%d(%A)') + '.png'
    #    plt.savefig(path)
    return

def simulation(number, title, ts_list, param, stops, should_draw, should_save):
    [buy_threshold, sell_threshold, ma_window, limit] = param
    print (title, 'Buy threshold: ', buy_threshold, 'Sell threshold: ', sell_threshold,  'MA window: ', ma_window, 'StopProfit: ', stops[0][0], 'StopLoss: ', stops[0][1])

    out = []
    num = 0
    profit_sum = 0.0
    equity = []
    for tsvalue in ts_list:
        [date, ts] = tsvalue
        #print('length:', ts.length)
        if ts.length < 10:
            continue

        num += 1
        (power, psum) = purchasePower(ts)
        result = breakout(psum, ts, param, stops)
        profits, matches, flag = result
        for profit, match in zip(profits, matches):
            #print(title, ' ... Profit', profit, 'Prices', prices, '(', indices, ')')
            longOrShort, prices, indices = match
            out.append([num, date, longOrShort, profit, prices, indices])
            profit_sum += profit
            equity.append(profit_sum)

        if should_draw:
            drawGraph(title, ts, date, power, psum, result, param)

    if should_save:
        path = './' + str(number).zfill(5) + '_' + title + '.csv'
        df = pd.DataFrame(out, columns=['No', 'Date', 'Long_Short', 'Profit', 'Price', 'Index'])
        df.to_csv(path, index=False)

    drawdown = np.min(equity)
    print('*', param, stops, '*    Profit: ', profit_sum, '  DrawDown:', drawdown)
    return (equity, profit_sum, drawdown)

def dataSource():
    now = Now()
    begin1 = datetime(2019, 3, 1)
    end1 = datetime(2020, 2, 28)
    begin2 = datetime(2020, 8, 20)
    server = MT5Bind('US30Cash')
    ts_list = []
    t = begin2
    while t <= now:
        t0 = jstTime(t.year, t.month, t.day, 19, 0)
        tmp = t0 + DeltaDay(1)
        t1 = jstTime(tmp.year, tmp.month, tmp.day, 5, 30)
        data = server.scrapeRange('M5', t0, t1)
        ts = TimeSeries(data, DATA_TYPE_XM, names=OHLCV, index=[1, 2, 3, 4, 5])
        #print(t, 'length:', ts.length)
        if ts.length > 20:
            ts_list.append([t, ts])
        t += DeltaDay(1)
    server.close()
    return ts_list

def evaluate1(should_draw, should_save):
    header = 'No, BuyTh, SellTh, StopProfit, StopLos, Profit, Drawdown'
    path = './evaluation1.csv'

    if os.path.exists(path):
        file = open(path, 'a')
    else:
        file = open(path, 'w')
        file.write(header + '\n')
    ts_list = dataSource()

    stops_list = [[[180, -180]],
                  [[180, -180], [250, 150], [300, 200], [350, 250], [400, 300], [450, 350], [500, 400]],
                  [[100, -100], [150, 50], [200, 120], [250, 170], [300, 220], [350, 270], [400, 320], [450, 370], [500, 420]],
                  [[50, -50], [100, 20], [150, 80], [200, 140], [250, 180], [300, 230], [350, 280], [400, 330], [450, 380], [500, 430]] ]
    num = 1
    for buy_threshold in range(10000, 30001, 1000):
        for sell_threshold in range(-10000, -30001, -1000):
            k = 1
            for stops in stops_list:
                equity, profit, drawdown = simulation(num, 'DJI-M5', ts_list, [buy_threshold, sell_threshold], stops, should_draw, should_save)
                s = str(num) + ',' + str(buy_threshold) + ',' + str(sell_threshold) + ',' + str(k) + ',' + str(profit) + ',' + str(drawdown)
                file.write(s + '\n')
                k += 1
                num += 1
    file.close()
    print('*** All Done ***')
    return

def evaluate2(should_draw, should_save):
    header = 'No, BuyTh, SellTh, ma_window, stop_profit, stop_loss, limit, Profit, Drawdown'
    path = './evaluation5.csv'
    if os.path.exists(path):
        file = open(path, 'a')
    else:
        file = open(path, 'w')
        file.write(header + '\n')
    ts_list = dataSource()

    stops_list= []
    for a in range(50, 410, 50):
        for b in range(50, 410, 50):
            if a > b:
                stops_list.append( [[a, -b]] )

    num = 1
    limit = 3
    for buy_threshold in range(10000, 30001, 5000):
        for sell_threshold in range(-10000, -30001, -5000):
            for ma in [1, 3, 5, 7]:
                k = 1
                for stops in stops_list:
                    equity, profit, drawdown = simulation(num, 'DJI-M5', ts_list, [buy_threshold, sell_threshold, ma, limit], stops, should_draw, should_save)
                    s = str(num) + ',' + str(buy_threshold) + ',' + str(sell_threshold) + ',' + str(ma) + ',' + str(stops[0][0]) + ',' + str(stops[0][1]) + ',' + str(limit) + ',' + str(profit) + ',' + str(drawdown)
                    file.write(s + '\n')
                    k += 1
                    num += 1
    file.close()
    print('*** All Done ***')
    return

def test():
    ts_list = dataSource()
    param = [10000, -10000, 1, 5]
    #stops = [[100, -100], [150, 90], [200, 140], [250, 190]]
    stops = [[100, -50]]
    equity, profit, drawdown = simulation(1, 'DJI-M5', ts_list, param, stops, True, True)
    fig, ax = makeFig(1, 1, (15, 5))
    graph = Graph(ax)
    prop = {'color':'green', 'style':'solid', 'width':1}
    graph.plot(range(len(equity)), equity, prop)
    graph.setTitle('DJI Day Trade Equity Curve Profit: $' + str(profit) + ' Drawdown: $' + str(drawdown), '', 'US$')
    print(profit, drawdown)
    return

if __name__ == '__main__':
    #evaluate2(False, False)
    test()


# -*- coding: UTF-8 -*-
import numpy as np
import random
import logging

logging.basicConfig(level=logging.WARNING)

class Action(object):
    def __init__(self):
        self.move = None
        self.collect = None
        self.offloading = []
        self.bandwidth = 0
        self.execution = []

class AgentState(object):
    def __init__(self):
        self.position = None
        self.obs = None

class EdgeDevice(object):
    edge_count = 0

    def __init__(self, obs_r, pos, spd, collect_r, max_buffer_size, movable=True, mv_bt=0, trans_bt=0):  # pos(x,y,h)
        self.no = EdgeDevice.edge_count
        EdgeDevice.edge_count += 1
        self.obs_r = obs_r  # observe radius
        self.init_pos = pos
        self.position = pos
        self.move_r = spd
        self.collect_r = collect_r
        self.mv_battery_cost = mv_bt
        self.trans_battery_cost = trans_bt
        self.data_buffer = {}
        self.max_buffer_size = max_buffer_size
        self.idle = True  # collecting idle
        self.movable = movable
        self.state = AgentState()
        self.action = Action()
        self.done_data = []
        self.offloading_idle = True
        self.total_data = {}
        self.computing_rate = 2e4
        self.computing_idle = True
        self.index_dim = 2
        self.collecting_sensors = []
        self.ptr = 0.2
        self.h = 5
        self.noise = 1e-13
        self.trans_rate = 0

    def move(self, new_move, h):
        if self.idle:
            self.position += new_move
            self.mv_battery_cost += np.linalg.norm(new_move)

    def get_total_data(self):
        total_data_state = np.zeros([self.index_dim, self.max_buffer_size])
        if self.total_data:
            for j, k in enumerate(list(self.total_data.keys())):
                # print(self.total_data[k])
                total_data_state[0, j] = self.total_data[k][0]
                total_data_state[1, j] = self.total_data[k][1]
        return total_data_state

    def get_done_data(self):
        done_data_state = np.zeros([self.index_dim, self.max_buffer_size])
        if self.done_data:
            for m, k in enumerate(self.done_data):
                done_data_state[0, m] = k[0]
                done_data_state[1, m] = k[1]
        return done_data_state

    def data_update(self, pak):
        if pak[1] in self.data_buffer.keys():
            self.data_buffer[pak[1]].append(pak)
        else:
            self.data_buffer[pak[1]] = [pak]

    def edge_exe(self, tmp_size, t=1):  # one-sum local execution
        if not self.total_data:
            return [0] * self.max_buffer_size
        # age update
        for k in self.total_data.keys():
            self.total_data[k][1] += t
        if len(self.done_data) >= self.max_buffer_size:
            return tmp_size
        # process
        if self.total_data and sum(self.action.execution):
            data2process = [[k, d] for k, d in self.total_data.items()]
            self.computing_idle = False
            if np.argmax(self.action.execution) >= len(data2process):
                self.action.execution = [0] * self.max_buffer_size
                self.action.execution[np.random.randint(len(data2process))] = 1
            for i, data in enumerate(data2process):
                if len(self.done_data) >= self.max_buffer_size:
                    break
                # print([i, tmp_size])
                tmp_size[i] += min(self.total_data[data2process[i][0]][0], self.computing_rate * self.action.execution[i] * t)
                self.total_data[data2process[i][0]][0] -= self.computing_rate * self.action.execution[i] * t
                if self.total_data[data2process[i][0]][0] <= 0:
                    self.total_data[data2process[i][0]][0] = tmp_size[i]
                    self.done_data.append(self.total_data[data2process[i][0]])
                    self.total_data.pop(data2process[i][0])
                    tmp_size[i] = 0
        return tmp_size

    def process(self, tmp_size, t=1):  # one-hot local execution
        if not self.total_data:
            return 0
        # age update
        for k in self.total_data.keys():
            self.total_data[k][1] += t
        if len(self.done_data) >= self.max_buffer_size:
            return 0
        # process
        if self.total_data and sum(self.action.execution):
            data2process = [[k, d] for k, d in self.total_data.items()]

            if self.action.execution.index(1) >= len(data2process):
                self.action.execution[self.action.execution.index(1)] = 0
                self.action.execution[np.random.randint(len(data2process))] = 1
                # print(self.action.execution)
            self.computing_idle = False
            tmp_size += min(self.total_data[data2process[self.action.execution.index(
                1)][0]][0], self.computing_rate * t)
            self.total_data[data2process[self.action.execution.index(
                1)][0]][0] -= self.computing_rate * t
            if self.total_data[data2process[self.action.execution.index(1)][0]][0] <= 0:
                self.total_data[data2process[self.action.execution.index(
                    1)][0]][0] = tmp_size
                self.done_data.append(self.total_data[data2process[self.action.execution.index(
                    1)][0]])
                self.total_data.pop(data2process[self.action.execution.index(
                    1)][0])
                tmp_size = 0
        return tmp_size

def agent_com(agent_list):
    age_dict = {}
    for u in agent_list:
        for k, v in u.data_buffer.items():
            if k not in age_dict:
                age_dict[k] = v[-1][1]
            elif age_dict[k] > v[-1][1]:
                age_dict[k] = v[-1][1]
    return age_dict

class Sensor(object):
    sensor_cnt = 0

    def __init__(self, pos, data_rate, bandwidth, max_ds, lam=0.5, weight=1):
        self.no = Sensor.sensor_cnt
        Sensor.sensor_cnt += 1
        self.position = pos
        self.weight = weight
        self.data_rate = data_rate  # generate rate
        self.bandwidth = bandwidth
        self.trans_rate = 8e3  # to be completed
        self.data_buffer = []
        self.max_data_size = max_ds
        self.data_state = bool(len(self.data_buffer))
        self.collect_state = False
        self.lam = lam
        self.noise_power = 1e-13 * self.bandwidth
        self.gen_threshold = 0.3

    def data_gen(self, t=1):
        # update age
        if self.data_buffer:
            for i in range(len(self.data_buffer)):
                self.data_buffer[i][1] += t
        new_data = self.data_rate * np.random.poisson(self.lam)
        # new_data = min(new_data, self.max_data_size)
        if new_data >= self.max_data_size or random.random() >= self.gen_threshold:
            return
        if new_data:
            self.data_buffer.append([new_data, 0, self.no])
            self.data_state = True

collecting_channel_param = {'suburban': (4.88, 0.43, 0.1, 21),
                            'urban': (9.61, 0.16, 1, 20),
                            'dense-urban': (12.08, 0.11, 1.6, 23),
                            'high-rise-urban': (27.23, 0.08, 2.3, 34)}

collecting_params = collecting_channel_param['urban']
a = collecting_params[0]
b = collecting_params[1]
yita0 = collecting_params[2]
yita1 = collecting_params[3]
carrier_f = 2.5e9

def collecting_rate(sensor, agent):
    d = np.linalg.norm(np.array(sensor.position) - np.array(agent.position))
    Pl = 1 / (1 + a * np.exp(-b * (np.arctan(agent.h / d) - a)))
    L = Pl * yita0 + yita1 * (1 - Pl)
    gamma = agent.ptr_col / (L * sensor.noise_power**2)
    rate = sensor.bandwidth * np.log2(1 + gamma)
    return rate

def data_collecting(sensors, agent, hovering_time):
    for k in agent.total_data.keys():
        agent.total_data[k][1] += 1
    if agent.idle and (len(agent.total_data.keys()) < agent.max_buffer_size):
        # obs_sensor = []
        data_properties = []
        # for k in agent.data_buffer.keys():
        #     for i, d in enumerate(agent.data_buffer[k]):
        #         agent.data_buffer[k][i][1] += 1
        for sensor in sensors:
            if not sensor.data_buffer:
                continue
            if (np.linalg.norm(np.array(sensor.position) - np.array(agent.position)) <= agent.collect_r) and not(sensor.collect_state) and not(sensor.no in agent.total_data.keys()):
                sensor.collect_state = True
                agent.collecting_sensors.append(sensor.no)
                agent.idle = False
                if len(agent.total_data.keys()) >= agent.max_buffer_size:
                    continue
                # obs_sensor.append(sensor)
                # if not (sensor.no in agent.data_buffer.keys()):
                #     agent.data_buffer[sensor.no] = []
                tmp_size = 0
                # trans_rate = collecting_rate(sensor, agent)
                for data in sensor.data_buffer:
                    tmp_size += data[0]
                    # data[1] += tmp_size / self.trans_rate  # age update
                if sensor.no in agent.data_buffer.keys():
                    agent.data_buffer[sensor.no].append(tmp_size)
                else:
                    agent.data_buffer[sensor.no] = [tmp_size]
                data_properties.append(tmp_size / sensor.trans_rate)
                agent.total_data[sensor.no] = [tmp_size, sensor.data_buffer[-1][1], sensor.no]
                # agent.total_data[sensor.no] = [tmp_size, np.average([x[1] for x in sensor.data_buffer]), sensor.no]
                sensor.data_buffer = []

        if data_properties:
            hovering_time = max(data_properties)
            # print([data_properties, hovering_time])
            return hovering_time
        else:
            return 0
    # finish collection
    elif not agent.idle:
        hovering_time -= 1
        if hovering_time <= 0:
            agent.idle = True
            for no in agent.collecting_sensors:
                sensors[no].collect_state = False
            agent.collecting_sensors = []
            hovering_time = 0
        return hovering_time
    else:
        return 0

def offloading(agent, center_pos, t=1):
    if not agent.done_data:
        #print('no done')
        return (False, {})
    for data in agent.done_data:
        data[1] += t

    if sum(agent.action.offloading):
        if agent.action.offloading.index(1) >= len(agent.done_data):
            agent.action.offloading[agent.action.offloading.index(1)] = 0
            agent.action.offloading[np.random.randint(len(agent.done_data))] = 1
        agent.offloading_idle = False
        dist = np.linalg.norm(np.array(agent.position) - np.array(center_pos))
        agent.trans_rate = trans_rate(dist, agent)  # to be completed
    else:
        return False, {}
    # print(agent.done_data)
    # print(agent.action.offloading)
    agent.done_data[agent.action.offloading.index(1)][0] -= agent.trans_rate * t
    if agent.done_data[agent.action.offloading.index(1)][0] <= 0:
        sensor_indx = agent.done_data[agent.action.offloading.index(1)][2]
        sensor_aoi = agent.done_data[agent.action.offloading.index(1)][1]
        sensor_data = agent.data_buffer[sensor_indx][0]
        del agent.data_buffer[sensor_indx][0]
        del agent.done_data[agent.action.offloading.index(1)]
        # return finish flag & total data
        agent.offloading_idle = True
        return True, {sensor_indx: [sensor_data, sensor_aoi]}
    return False, {}

def trans_rate(dist, agent):  # to be completed
    W = 1e6 * agent.action.bandwidth
    Pl = 1 / (1 + a * np.exp(-b * (np.arctan(agent.h / dist) - a)))
    fspl = (4 * np.pi * carrier_f * dist / (3e8))**2
    L = Pl * fspl * 10**(yita0 / 20) + 10**(yita1 / 20) * fspl * (1 - Pl)
    rate = W * np.log2(1 + agent.ptr / (L * agent.noise * W))
    print('agent-{} rate: {},{},{},{},{}'.format(agent.no, dist, agent.action.bandwidth, Pl, L, rate))
    return rate

class MEC_world(object):
    def __init__(self, map_size, agent_num, sensor_num, obs_r, speed, collect_r, max_size=1, sensor_lam=0.5):
        self.agents = []
        self.sensors = []
        self.map_size = map_size
        self.center = (map_size / 2, map_size / 2)
        self.sensor_count = sensor_num
        self.agent_count = agent_num
        self.max_buffer_size = max_size
        sensor_bandwidth = 1000
        max_ds = sensor_lam * 2
        data_gen_rate = 1
        self.offloading_slice = 1
        self.execution_slice = 1
        self.time = 0
        self.DS_map = np.zeros([map_size, map_size])
        self.DS_state = np.ones([map_size, map_size, 2])
        self.hovering_list = [0] * self.agent_count
        self.tmp_size_list = [0] * self.agent_count
        # [self.tmp_size_list.append([0] * self.max_buffer_size) for i in range(self.agent_count)]
        self.offloading_list = []
        self.finished_data = []
        self.obs_r = obs_r
        self.move_r = speed
        self.collect_r = collect_r
        self.sensor_age = {}
        # random.seed(7)
        self.sensor_pos = [random.choices([i for i in range(int(0.1 * self.map_size), int(0.9 * self.map_size))], k=sensor_num),
                           random.choices([i for i in range(int(0.1 * self.map_size), int(0.9 * self.map_size))], k=sensor_num)]
        # self.sensor_pos = [random.choices([i for i in range(int(0.1 * self.map_size), int(0.5 * self.map_size))], k=int(sensor_num / 2)) + random.choices(
        #     [i for i in range(int(0.5 * self.map_size), int(0.9 * self.map_size))], k=int(sensor_num / 2)), random.choices([i for i in range(int(0.1 * self.map_size), int(0.9 * self.map_size))], k=sensor_num)]
        for i in range(sensor_num):
            self.sensors.append(
                Sensor(np.array([self.sensor_pos[0][i], self.sensor_pos[1][i]]), data_gen_rate, sensor_bandwidth, max_ds, lam=sensor_lam))
            self.sensor_age[i] = 0
            self.DS_map[self.sensor_pos[0][i], self.sensor_pos[1][i]] = 1
        self.agent_pos_init = [random.sample([i for i in range(int(0.15 * self.map_size), int(0.85 * self.map_size))], agent_num),
                               random.sample([i for i in range(int(0.15 * self.map_size), int(0.85 * self.map_size))], agent_num)]

        for i in range(agent_num):
            self.agents.append(
                EdgeDevice(self.obs_r, np.array([self.agent_pos_init[0][i], self.agent_pos_init[1][i]]), speed, collect_r, self.max_buffer_size))

    def step(self):
        # update sensor age
        for k in self.sensor_age.keys():
            self.sensor_age[k] += 1
        # data generation & DS_state update
        logging.info("data generation")
        for sensor in self.sensors:
            sensor.data_gen()
            if sensor.data_buffer:
                data_size = sum(i[0] for i in sensor.data_buffer)
                # update data source state, note that the (x,y) is reversed to the matrix index (i,j)
                self.DS_state[sensor.position[1], sensor.position[0]] = [
                    data_size, sensor.data_buffer[0][1]]

        # edge process  offloading collect
        logging.info("edge operation")
        age_dict = {}
        for i, agent in enumerate(self.agents):
            # edge process
            self.tmp_size_list[i] = agent.process(self.tmp_size_list[i])
            # offloading
            finish_flag, data_dict = offloading(agent, self.center)
            # update reward state
            # print([i, finish_flag, data_dict])
            if finish_flag:
                for sensor_id, data in data_dict.items():
                    self.finished_data.append([data[0], data[1], sensor_id])
                    if sensor_id in age_dict.keys():
                        age_dict[sensor_id].append(data[1])
                    else:
                        age_dict[sensor_id] = [data[1]]
                    # self.sensor_age[sensor_id] -=data[1]
            # collecting
            self.hovering_list[i] = data_collecting(self.sensors, agent, self.hovering_list[i])
            # print(self.hovering_list[i])
        for id in age_dict.keys():
            self.sensor_age[id] = sorted(age_dict[id])[0]
        print('hovering:{}'.format(self.hovering_list))


import sys
sys.path.append('../')
import os
import glob
import math
# os.environ["CUDA_VISIBLE_DEVICES"]="0"

from keras import backend as K
from keras.models import load_model
from keras.models import Model
from keras.optimizers import Adam
from scipy.misc import imread
import numpy as np
from matplotlib import pyplot as plt

from models.keras_ssd512_Siamese_analyze_contrastive_loss import ssd_512
from keras_loss_function.keras_ssd_loss import SSDLoss
from keras_layers.keras_layer_AnchorBoxes import AnchorBoxes
from keras_layers.keras_layer_DecodeDetections import DecodeDetections
from keras_layers.keras_layer_DecodeDetectionsFast import DecodeDetectionsFast
from keras_layers.keras_layer_L2Normalization import L2Normalization
from data_generator.object_detection_2d_data_generator_paired_input import DataGenerator
from eval_utils.average_precision_evaluator import Evaluator

# %matplotlib inline

# Set a few configuration parameters.
img_height = 512
img_width = 512
# model_mode indicates the way the pretrained model was created.
# In training model, Model_Build should == 'Load_Model'. decode_detections will be called in the Evaluator.
# However, decode_detections is run on CPU and is very slow.
# In inference model, Model_Build should == 'New_Model_Load_Weights'.
# DecodeDetections will be called when build the model. DecodeDetections is writen in tensorflow and is run GPU.
# It seems that the result under inference model is slightly better than that under training model.
# Maybe DecodeDetections and decode_detections are not exactly the same.
model_mode = 'inference'  # 'training'#
assert model_mode == 'inference'

model_path = '../trained_weights/SSD512_City_to_foggy0_01_resize_400_800/current/pool12_loss_weights_0_000005/epoch-55_loss-3.9684_val_loss-5.2015.h5'

batch_size = 8

DatasetName = 'City_to_foggy0_01_paired_input'
processed_dataset_path = './processed_dataset_h5/' + DatasetName
if not os.path.exists(processed_dataset_path):
    os.makedirs(processed_dataset_path)

if len(glob.glob(os.path.join(processed_dataset_path, '*.h5'))):
    Dataset_Build = 'Load_Dataset'
else:
    Dataset_Build = 'New_Dataset'

# The anchor box scaling factors used in the original SSD512
scales_coco = [0.07, 0.15, 0.3, 0.45, 0.6, 0.75, 0.9, 1.05]
scales = scales_coco
top_k = 200
confidence_thresh = 0.01
nms_iou_threshold = 0.5

if DatasetName == 'City_to_foggy0_01_paired_input':
    resize_image_to = (400, 800)

    # Our model will produce predictions for these classes.
    classes = ['background',
               'person', 'rider', 'car', 'truck',
               'bus', 'train', 'motorcycle', 'bicycle']
    train_classes = classes
    train_include_classes = 'all'
    # Number of positive classes, 8 for domain Cityscapes, 20 for Pascal VOC, 80 for MS COCO, 1 for SIM10K
    n_classes = len(classes) - 1

else:
    raise ValueError('Undefined dataset name.')

# 1: Build the Keras model

K.clear_session()  # Clear previous models from memory.

# import tensorflow as tf
# from keras.backend.tensorflow_backend import set_session
#
# config = tf.ConfigProto()
# config.gpu_options.allow_growth = True  # dynamically grow the memory used on the GPU
# config.log_device_placement = True  # to log device placement (on which device the operation ran)
# # (nothing gets printed in Jupyter, only if you run it standalone)
# sess = tf.Session(config=config)
# set_session(sess)  # set this TensorFlow session as the default session for Keras

# model.output = `predictions`: (batch, n_boxes_total, n_classes + 4 + 8)
# In inference mode, the predicted locations have been converted to absolute coordinates.
# In addition, we have performed confidence thresholding, per-class non-maximum suppression, and top-k filtering.
model = ssd_512(image_size=(img_height, img_width, 3),
                n_classes=n_classes,
                mode=model_mode,
                l2_regularization=0.0005,
                scales=scales,
                aspect_ratios_per_layer= [[1.0, 2.0, 0.5],
                                          [1.0, 2.0, 0.5, 3.0, 1.0/3.0],
                                          [1.0, 2.0, 0.5, 3.0, 1.0/3.0],
                                          [1.0, 2.0, 0.5, 3.0, 1.0/3.0],
                                          [1.0, 2.0, 0.5, 3.0, 1.0/3.0],
                                          [1.0, 2.0, 0.5],
                                          [1.0, 2.0, 0.5]],
                two_boxes_for_ar1=True,
                steps=[8, 16, 32, 64, 128, 256, 512],
                offsets=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
                clip_boxes=False,
                variances=[0.1, 0.1, 0.2, 0.2],
                normalize_coords=True,
                subtract_mean=[123, 117, 104],
                swap_channels=[2, 1, 0],
                confidence_thresh=confidence_thresh,
                iou_threshold=nms_iou_threshold,
                top_k=top_k,
                nms_max_output_size=400)

# 2: Load the trained weights into the model
model.load_weights(model_path, by_name=True)

if Dataset_Build == 'New_Dataset':
    # 1: Instantiate two `DataGenerator` objects: One for training, one for validation.

    # Optional: If you have enough memory, consider loading the images into memory for the reasons explained above.

    train_dataset = DataGenerator(dataset='train', load_images_into_memory=False, hdf5_dataset_path=None)

    Cityscapes_images_dir = '../../datasets/Cityscapes/JPEGImages'
    Cityscapes_target_images_dir = '../../datasets/CITYSCAPES_beta_0_01/JPEGImages'

    # The directories that contain the annotations.
    Cityscapes_annotation_dir = '../../datasets/Cityscapes/Annotations'

    # The paths to the image sets.
    Cityscapes_train_source_image_set_filename = '../../datasets/Cityscapes/ImageSets/Main/train_source.txt'
    Cityscapes_train_target_image_set_filename = '../../datasets/Cityscapes/ImageSets/Main/train_target.txt'

    # images_dirs, image_set_filenames, and annotations_dirs should have the same length
    train_dataset.parse_xml(images_dirs=[Cityscapes_images_dir,
                                         Cityscapes_target_images_dir],
                            image_set_filenames=[Cityscapes_train_source_image_set_filename,
                                                 Cityscapes_train_target_image_set_filename],
                            annotations_dirs=[Cityscapes_annotation_dir,
                                              Cityscapes_annotation_dir],
                            classes=classes,
                            include_classes='all',
                            exclude_truncated=False,
                            exclude_difficult=False,
                            ret=False)

    # Optional: Convert the dataset into an HDF5 dataset. This will require more disk space, but will
    # speed up the training. Doing this is not relevant in case you activated the `load_images_into_memory`
    # option in the constructor, because in that cas the images are in memory already anyway. If you don't
    # want to create HDF5 datasets, comment out the subsequent two function calls.

    # After create these h5 files, if you have resized the input image, you need to reload these files. Otherwise,
    # the images and the labels will not change.

    train_dataset.create_hdf5_dataset(file_path=os.path.join(processed_dataset_path, 'dataset_train.h5'),
                                      resize=resize_image_to,
                                      variable_image_size=True,
                                      verbose=True)

    train_filenames = []
    with open(Cityscapes_train_source_image_set_filename, 'r') as f:
        train_filenames.extend([os.path.join(Cityscapes_images_dir, line.strip()) for line in f])
    with open(Cityscapes_train_target_image_set_filename, 'r') as f:
        train_filenames.extend([os.path.join(Cityscapes_target_images_dir, line.strip()) for line in f])

    train_dataset = DataGenerator(dataset='train',
                                  load_images_into_memory=False,
                                  hdf5_dataset_path=os.path.join(processed_dataset_path, 'dataset_train.h5'),
                                  filenames=train_filenames,
                                  filenames_type='text',
                                  images_dir=Cityscapes_images_dir)

elif Dataset_Build == 'Load_Dataset':
    # 1: Instantiate two `DataGenerator` objects: One for training, one for validation.

    # The directories that contain the images.
    Cityscapes_images_dir = '../../datasets/Cityscapes/JPEGImages'
    Cityscapes_target_images_dir = '../../datasets/CITYSCAPES_beta_0_01/JPEGImages'

    # The paths to the image sets.
    Cityscapes_train_source_image_set_filename = '../../datasets/Cityscapes/ImageSets/Main/train_source.txt'
    Cityscapes_train_target_image_set_filename = '../../datasets/Cityscapes/ImageSets/Main/train_target.txt'

    train_filenames = []
    with open(Cityscapes_train_source_image_set_filename, 'r') as f:
        train_filenames.extend([os.path.join(Cityscapes_images_dir, line.strip()) for line in f])
    with open(Cityscapes_train_target_image_set_filename, 'r') as f:
        train_filenames.extend([os.path.join(Cityscapes_target_images_dir, line.strip()) for line in f])

    train_dataset = DataGenerator(dataset='train',
                                  load_images_into_memory=False,
                                  hdf5_dataset_path=os.path.join(processed_dataset_path, 'dataset_train.h5'),
                                  filenames=train_filenames,
                                  filenames_type='text',
                                  images_dir=Cityscapes_images_dir)

else:
    raise ValueError('Undefined Dataset_Build. Dataset_Build should be New_Dataset or Load_Dataset.')

# Make predictions:
# 1: Set the generator for the predictions.

# For the test generator:
from data_generator.object_detection_2d_geometric_ops import Resize
from data_generator.object_detection_2d_photometric_ops import ConvertTo3Channels
from data_generator.object_detection_2d_misc_utils import apply_inverse_transforms

# First convert the input image to 3 channels and size img_height X img_width
# Also, convert the groundtruth bounding box
# Remember, if you want to visualize the predicted box on the original image,
# you need to apply the corresponding reverse transformation.
convert_to_3_channels = ConvertTo3Channels()
resize = Resize(height=img_height, width=img_width)

test_generator = train_dataset.generate(batch_size=batch_size,
                                        shuffle=False,
                                        transformations=[convert_to_3_channels,
                                                         resize],
                                        label_encoder=None,
                                        returns={'processed_images',
                                                 'filenames',
                                                 'inverse_transform',
                                                 'original_images',
                                                 'original_labels'},
                                        keep_images_without_gt=False)

n_images = train_dataset.get_dataset_size() // 2
n_batches = int(math.ceil(n_images / batch_size))
print("Number of images in the dataset:\t{:>6}".format(n_images))

contrastive_loss = [0] * 9

for i in range(n_batches):
    batch_images, batch_filenames, batch_inverse_transforms, batch_original_images, batch_original_labels = next(test_generator)
    y_pred = model.predict(batch_images)
    for k, batch_item in enumerate(y_pred):
        contrastive_loss[k] += np.mean(np.sum(np.square(batch_item), axis=1))


import sys
import os
import glob
import argparse
import numpy as np
from os.path import expanduser
import h5py
import pickle
from time import time
import subprocess

velocity_bound=10000.

def retrieve_commit_hash(path_to_repo):
    """ Return the commit hash of the git branch currently live in the input path.
    Parameters
    ----------
    path_to_repo : string
    Returns
    -------
    commit_hash : string
    """
    cmd = 'cd {0} && git rev-parse HEAD'.format(path_to_repo)
    return subprocess.check_output(cmd, shell=True).strip()

def velocity_bug_fix(output_snapshot, scalefactor=1.0):
    """
    output_snapshot: dict containing mock data to be modified
    scale: correction factor for velocities
    """
    #overwrite with corrected values
    for d in ['x', 'y', 'z']:
        print(".....Correcting v{}'s".format(d))
        corrected_halo_velocity = output_snapshot['target_halo_v{}'.format(d)]*scalefactor
        host_centric_v = output_snapshot['host_centric_v{}'.format(d)]
        corrected_galaxy_velocity = corrected_halo_velocity + host_centric_v
        output_snapshot['target_halo_v{}'.format(d)] = corrected_halo_velocity
        output_snapshot['v{}'.format(d)] = corrected_galaxy_velocity

    return output_snapshot

def mask_large_velocities(output_snapshot, max_value=velocity_bound):

    total = len(output_snapshot['target_halo_vx'])
    mask = np.ones(total, dtype=bool)
    for d in ['x', 'y', 'z']:
        mask &= np.abs(output_snapshot['target_halo_v{}'.format(d)]) < max_value

    nbad = np.sum(~mask)
    print('.....Masking {} galaxy(ies); total = {}; fraction = {}'.format(nbad, total, nbad/float(total)))

    return mask

def apply_mask(output_snapshot, mask):

    for k in output_snapshot.keys():
        output_snapshot[k] = output_snapshot[k][mask]

    return output_snapshot

def healpix_mock_modify(healpix_filename, commit_hash, functions=None, correction_data=None):

    output_mock = {}
    masks_used = {}
    print('Starting correction of {}'.format(os.path.basename(healpix_filename)))
    with h5py.File(healpix_filename, 'r') as fh:
        print('...copying input to output_mock')
        #copy input mock to output mock
        for (k, v) in fh.items():
            if k.isdigit():
                output_snapshot = {}
                for kk, vv  in v.items():
                    output_snapshot[kk] = vv.value

                output_mock[k] = output_snapshot

        print('...Keys copied to output mock: {}'.format(', '.join(output_mock.keys())))

        for f in functions:
            corrections = correction_data.get(str(f),{})
            for (k, v) in output_mock.items():
                if k.isdigit():
                    if len(v) == 0:
                        print('Skipping empty snap {}'.format(k))
                        continue
                    print('...Processing snap {} with {} and data-correction value(s) {}'.format(k, str(f), corrections[int(k)]))
                    output_mock[k] = f(v, corrections[int(k)])

                    #apply masks
                    mask = mask_large_velocities(output_mock[k], max_value=velocity_bound)
                    output_mock[k] = apply_mask(output_mock[k], mask)
                    print('...Masked length of arrays in snapshot {} = {}'.format(k, len(output_mock[k]['galaxy_id'])))
                    masks_used['large_velocities'] =  mask_large_velocities

                    del mask

        # copy and correct metaData
        k = 'metaData'
        output_mock[k] = {}
        for tk, v in fh[k].items():
            output_mock[k][tk] = v.value

        output_mock[k]['versionMinorMinor'] += 1
        for n, f in enumerate(functions):
            ckey = 'comment_'+str(n)
            output_mock[k][ckey] = ' '.join(['Corrected with', str(f)])
            print('...Adding metaData comment: {}'.format(output_mock[k][ckey]))

        for tk, v in masks_used.items():
            ckey = 'mask_'+tk
            output_mock[k][ckey] = ' '.join(['Corrected with', str(v)])
            print('...Adding metaData comment: {}'.format(output_mock[k][ckey]))

        output_mock[k]['current_commit_hash'] = commit_hash

    return output_mock

def write_output_mock(output_mock, output_healpix_file):
    hdfFile = h5py.File(output_healpix_file, 'w')

    for k, v in output_mock.items():
        gGroup = hdfFile.create_group(k)
        for tk in v.keys():
            gGroup[tk] = v[tk]

    hdfFile.close()
    print('...Wrote {} to disk'.format(output_healpix_file))
    return

home = expanduser("~")
path_to_cosmodc2 = os.path.join(home, 'cosmology/cosmodc2')
if 'mira-home' in home:
    sys.path.insert(0, '/gpfs/mira-home/ekovacs/.local/lib/python2.7/site-packages')
sys.path.insert(0, path_to_cosmodc2)

parser = argparse.ArgumentParser()
parser.add_argument("-cutout",
    help="healpix cutout number to modify",
    default='*')
parser.add_argument("-healpix_fname_template",
    help="Template filename of healpix cutout to modify",
    default='baseDC2*cutout_{}.hdf5')
parser.add_argument("-input_master_dirname",
    help="Directory name (relative to home) storing input and output healpix file directories",
    default='cosmology/DC2/OR_Production')
parser.add_argument("-output_mock_dirname",
    help="Directory name (relative to input_master_dirname) storing output mock healpix files",
    default='baseDC2_9.8C_v1.1_velocity_bug_fixes')
#    default='baseDC2_min_9.8_centrals_v0.4.7_velocity_bug_fixes')
parser.add_argument("-input_mock_dirname",
    help="Directory name (relative to input_master_dirname) storing input mock healpix files",
    default='baseDC2_9.8C_v1.1')
#    default='baseDC2_min_9.8_centrals_v0.4.5')
#    default='baseDC2_min_9.8_centrals_v0.4.5_test')
parser.add_argument("-modify_functions",
    help="Functions applied to modify input -> output",
    nargs='+', choices=[velocity_bug_fix],
    default=[velocity_bug_fix])

args = parser.parse_args()

#setup directory names
input_master_dirname = os.path.join(home, args.input_master_dirname)
input_mock_dirname = os.path.join(input_master_dirname, args.input_mock_dirname)
output_mock_dirname = os.path.join(input_master_dirname, args.output_mock_dirname)
healpix_filename = args.healpix_fname_template.format(args.cutout)
function_list = args.modify_functions

print('Reading input from {}\n'.format(input_mock_dirname))
print('Writing output to {}\n'.format(output_mock_dirname))
print('Modifying healpix files matching {} with function(s) {}\n'.format(healpix_filename, function_list))

current_commit_hash = retrieve_commit_hash(path_to_cosmodc2)[0:7]
print('Using current commit hash {}'.format(current_commit_hash))

#load additional data needed for corrections
correction_data = {}
function_names = map(str, function_list)
for f in map(str, function_list):
    if 'velocity_bug_fix' in f:
        datfile = os.path.join(path_to_cosmodc2, 'scripts/velocity_correction_factors.pkl')
        #text file option
        #datfile = os.path.join(path_to_cosmodc2, 'scripts/velocity_correction_factors.txt')
        #ts, sf = np.loadtxt(datfile, unpack=True, usecols=[0, 1])
        #correction_data[f] =dict(zip(ts.astype(int), sf))
        with open(datfile, 'rb') as handle:
            correction_data[f] = pickle.load(handle)
        print('Using correction data input from {}\n'.format(datfile))

healpix_files = sorted(glob.glob(os.path.join(input_mock_dirname, healpix_filename)))
start_time = time()

for hpx in healpix_files:
    start_file = time()
    output_mock = healpix_mock_modify(hpx, current_commit_hash, functions=function_list, correction_data=correction_data)
    output_healpix_file = os.path.join(output_mock_dirname, os.path.basename(hpx))
    write_output_mock(output_mock, output_healpix_file)
    end_file = time()
    print('Processed {} in {:.2f} minutes\n'.format(os.path.basename(output_healpix_file), (end_file - start_file)/60.))

time_stamp = time()
msg = "End-to-end runtime = {0:.2f} minutes\n"
print(msg.format((time_stamp-start_time)/60.))


from sklearn.feature_extraction.text import CountVectorizer
from sklearn.externals import joblib
from pathlib import Path
from micro_influencer_utilities import topic_selection
import numpy as np
import datasetUtils as dsu
import embeddings
import sys
import os
import re

dataset_path = "FastText/dataset.vec"
tweet_threshold = 3

vectorizer = CountVectorizer(stop_words="english", analyzer="word")
analyzer = vectorizer.build_analyzer()

if len(sys.argv)== 2:
    topic_selected = sys.argv[1]
    if not topic_selected.startswith('#'):
        topic_selected = "#"+topic_selected
else:
    topic_selected = topic_selection()
pathToUserList = "./Data/"+topic_selected+"/00_potential_micro_influencers_users/user_list.csv"
f = open(pathToUserList, "r")
usernames = f.read().split(",")

print("Loading embeddings dataset...")
wordDictionary = dsu.parseFastText(dataset_path)
print("Data successfully loaded.")

outfile = []
user_progress = 0
for username in usernames:
    user_progress +=1
    print("\nUser:",username)
    print("progress: " + str(user_progress) + "/" + str(len(usernames)))
    tweet_file_path = "Data/"+topic_selected+"/02_users_tweets/"+username
    if not os.path.isfile(tweet_file_path):
        print("Error. Cannot find tweets for username",username)
        continue

    filteredTweets = []
    #preprocess tweets
    for tweet in open(tweet_file_path, "r"):

        if re.match(r'^(RT)', tweet):
            #ingore pure retweets (= with no associated text)
            continue
        #remove links starting with "http"
        tweet = re.sub(r'((http)([^\s]*)(\s|$))|((http)([^\s]*)$)', "", tweet)
        #remove links with no http (probably unnecessary)
        tweet = re.sub(r'(\s([^\s]*)\.([^\s]*)\/([^\s]*)\s)|(^([^\s]*)\.([^\s]*)\/([^\s]*)(\s|$))|(\s([^\s]*)\.([^\s]*)\/([^\s]*)$)', " ", tweet)
        #remove mentions
        tweet = re.sub(r'(\s(@)([^\s]*)\s)|((^@)([^\s]*)(\s|$))|(@([^\s]*)$)', " ", tweet)
        #hashtags are removed by countvectorizer

        filteredTweets.append(tweet)

        if len(filteredTweets) == 0:
            print("Not enough tweets for prediction.")
            continue

    #now we can process the tweet using embeddings.transofrmTextForTraining
    try:
        tweetEmbeddings = embeddings.transformTextForTesting(wordDictionary, tweet_threshold, filteredTweets, "conc")
        print("Embeddings computed.")
    except:
        #most of tweets are ingored for brevity/no embedding correspondence
        print("Not enough tweets for prediction.")
        continue

    scores = {}
    #load the saved ML models
    for trait in ["O","C","E","A","N"]:
        # model = joblib.load("Models/SVM_fasttext_conc_"+trait+".pkl")
        model = joblib.load("Models/SVM/SVM_"+trait+".pkl")
        mean = np.mean(tweetEmbeddings, axis = 0)
        score = model.predict([mean])
        scores[trait] = float(str(score[0])[0:5])
        print("\tScore for",trait,"is:",str(score[0])[0:5])

    jung = ""
    if scores["E"] > 3:
        jung = "E"
    else:
        jung = "I"
    if scores["O"] > 3:
        jung = jung + "N"
    else:
        jung = jung + "S"
    if scores["A"] > 3:
        jung = jung + "F"
    else:
        jung = jung + "T"
    if scores["C"] > 3:
        jung = jung + "J"
    else:
        jung = jung + "P"

    print("\tJungian type is",jung)

    outfile.append(username+","+str(scores["O"])+","+str(scores["C"])+","+str(scores["E"])+","+str(scores["A"])+","+str(scores["N"]))

pathToResults = "./Data/"+topic_selected+"/03_users_parameters/big5/"
filename = pathToResults + "big5.csv"
fp = open(filename,"w")
for elem in outfile:
    fp.write(elem+"\n")
fp.close()


#coding=utf-8

from sklearn.ensemble import RandomForestClassifier as RF
from sklearn import cross_validation
from sklearn.metrics import confusion_matrix
from collections import *
import pandas as pd
import numpy as np
import re, os

def getOpcodeSequence(filename):
    opcode_seq = []
    p = re.compile(r'\s([a-fA-F0-9]{2}\s)+\s+([a-z]+)') # 需要再匹配一个制表符
    with open(filename) as f:
        for line in f:
            if line.startswith(".text"): # 指定解析代码行
                m = re.findall(p,line)
                if m:
                    # print m # 结果匹配opcode及前一个制表符16进制数据
                    # eg:[('FF\t', 'jmp')]
                    opc = m[0][1]
                    #  and opc != "a" and opc != "db" and opc != "dword" and opc != "unicode" and opc != "dd" and opc != "word"
                    if opc != "align":
                        opcode_seq.append(opc)
    return opcode_seq

def getOpcodeNgram(ops, n=3):
    opngramlist = [tuple(ops[i:i+n]) for i in range(len(ops)-n)]
    opngram = Counter(opngramlist)
    return opngram

subtrainfeature = pd.read_csv("3gramfeature.csv")
# 取csv文件的第一行中除去id列，并转化为list
features = subtrainfeature.columns.tolist()[1:]
print features

'''
9
iouyxgUJX7pjI1LBTSDF
IGX67kgpuNz8JD4Ciwd2
gGfFjZ8Xab24eQ6zADN1
ETH8lrLD7Yx05Wu16hcK
8
ESMo7G9RVnYNgB3fHdpP
EMaDXLFSlr9TyW68JNOz
56KhNHUTXZpryYaGqoj4
5102w8Hp4TRkWbIUXSgZ
7
EZ7FwiO2p4DCrH6SJYBT
ewsfhA3Dunp49SFNMoVR
DpNvYX1qedHLiQnRhF98
28U1hRkQ6Yl57493ZdXD
6
k9H1WjlnQPASt0mg25hB
K51kg4QuOW3zMeft7lXU
IZWuOMidGxg1HYNP9fJr
f0jhPJ2913c56yY4mwEi
5
e2r6IncxE1LQOKFGgphj
cIojVJGQOtrL0S1ApeDY
4UTMdcZkxzLvwygO8EuK
1KB3Z7gd5aN4Xmx8W0sf
'''

filepath = '/home/zrx/code/malwaredata/train/4UTMdcZkxzLvwygO8EuK.asm'
tops = getOpcodeSequence(filepath)
t3gram = getOpcodeNgram(tops)
print t3gram
testdata = {}
for feature in features:
    # print feature
    # print t3gram[feature]
    feature = tuple(eval(feature)) # 转""数据为元组
    if feature in t3gram.keys():
        print "get"
        testdata[feature] = t3gram[feature] # 传递统计数值
    else:
        testdata[feature] = 0
# testdata = pd.DataFrame(testdata)
print testdata.values()
dd = np.array(testdata.values(), dtype = int).reshape((1,-1))
print dd

subtrainLabel = pd.read_csv('subtrainLabels.csv')
subtrainfeature = pd.read_csv("3gramfeature.csv")
subtrain = pd.merge(subtrainLabel,subtrainfeature,on='Id')
labels = subtrain.Class
subtrain.drop(["Class","Id"], axis=1, inplace=True)
subtrain = subtrain.as_matrix()

X_train, X_test, y_train, y_test = cross_validation.train_test_split(subtrain,labels,test_size=0.3)
srf = RF(n_estimators=500, n_jobs=-1)
# srf.fit(X_train,y_train)
srf.fit(subtrain, labels)
print srf.predict(dd)
print srf.predict_proba(dd)
print srf.predict_log_proba(dd)
# print srf.score(X_test,y_test)