# MIT License
#
# Copyright (C) The Adversarial Robustness Toolbox (ART) Authors 2021
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
# documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
# Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
"""
This module implements the task specific estimator for PyTorch object detectors.
"""
import logging
from typing import List, Dict, Optional, Tuple, Union, TYPE_CHECKING

import numpy as np

from art.estimators.object_detection.object_detector import ObjectDetectorMixin
from art.estimators.pytorch import PyTorchEstimator

if TYPE_CHECKING:
    # pylint: disable=C0412
    import torch

    from art.utils import CLIP_VALUES_TYPE, PREPROCESSING_TYPE
    from art.defences.preprocessor.preprocessor import Preprocessor
    from art.defences.postprocessor.postprocessor import Postprocessor

logger = logging.getLogger(__name__)

class PyTorchObjectDetector(ObjectDetectorMixin, PyTorchEstimator):
    """
    This module implements the task specific estimator for PyTorch object detection models following the input and
    output formats of torchvision.
    """

    estimator_params = PyTorchEstimator.estimator_params + ["input_shape", "optimizer", "attack_losses"]

    def __init__(
        self,
        model: "torch.nn.Module",
        input_shape: Tuple[int, ...] = (-1, -1, -1),
        optimizer: Optional["torch.optim.Optimizer"] = None,
        clip_values: Optional["CLIP_VALUES_TYPE"] = None,
        channels_first: Optional[bool] = False,
        preprocessing_defences: Union["Preprocessor", List["Preprocessor"], None] = None,
        postprocessing_defences: Union["Postprocessor", List["Postprocessor"], None] = None,
        preprocessing: "PREPROCESSING_TYPE" = None,
        attack_losses: Tuple[str, ...] = (
            "loss_classifier",
            "loss_box_reg",
            "loss_objectness",
            "loss_rpn_box_reg",
        ),
        device_type: str = "gpu",
    ):
        """
        Initialization.

        :param model: Object detection model. The output of the model is `List[Dict[Tensor]]`, one for each input
                      image. The fields of the Dict are as follows:

                      - boxes (FloatTensor[N, 4]): the predicted boxes in [x1, y1, x2, y2] format, with values
                        between 0 and H and 0 and W
                      - labels (Int64Tensor[N]): the predicted labels for each image
                      - scores (Tensor[N]): the scores or each prediction
        :param input_shape: The shape of one input sample.
        :param optimizer: The optimizer for training the classifier.
        :param clip_values: Tuple of the form `(min, max)` of floats or `np.ndarray` representing the minimum and
               maximum values allowed for features. If floats are provided, these will be used as the range of all
               features. If arrays are provided, each value will be considered the bound for a feature, thus
               the shape of clip values needs to match the total number of features.
        :param channels_first: Set channels first or last.
        :param preprocessing_defences: Preprocessing defence(s) to be applied by the classifier.
        :param postprocessing_defences: Postprocessing defence(s) to be applied by the classifier.
        :param preprocessing: Tuple of the form `(subtrahend, divisor)` of floats or `np.ndarray` of values to be
               used for data preprocessing. The first value will be subtracted from the input. The input will then
               be divided by the second one.
        :param attack_losses: Tuple of any combination of strings of loss components: 'loss_classifier', 'loss_box_reg',
                              'loss_objectness', and 'loss_rpn_box_reg'.
        :param device_type: Type of device to be used for model and tensors, if `cpu` run on CPU, if `gpu` run on GPU
                            if available otherwise run on CPU.
        """
        import torch
        import torchvision

        torch_version = list(map(int, torch.__version__.lower().split("+", maxsplit=1)[0].split(".")))
        torchvision_version = list(map(int, torchvision.__version__.lower().split("+", maxsplit=1)[0].split(".")))
        assert not (torch_version[0] == 1 and (torch_version[1] == 8 or torch_version[1] == 9)), (
            "PyTorchObjectDetector does not support torch==1.8 and torch==1.9 because of "
            "https://github.com/pytorch/vision/issues/4153. Support will return for torch==1.10."
        )
        assert not (torchvision_version[0] == 0 and (torchvision_version[1] == 9 or torchvision_version[1] == 10)), (
            "PyTorchObjectDetector does not support torchvision==0.9 and torchvision==0.10 because of "
            "https://github.com/pytorch/vision/issues/4153. Support will return for torchvision==0.11."
        )

        super().__init__(
            model=model,
            clip_values=clip_values,
            channels_first=channels_first,
            preprocessing_defences=preprocessing_defences,
            postprocessing_defences=postprocessing_defences,
            preprocessing=preprocessing,
            device_type=device_type,
        )

        self._input_shape = input_shape
        self._optimizer = optimizer
        self._attack_losses = attack_losses

        if self.clip_values is not None:
            if self.clip_values[0] != 0:
                raise ValueError("This classifier requires un-normalized input images with clip_vales=(0, max_value).")
            if self.clip_values[1] <= 0:  # pragma: no cover
                raise ValueError("This classifier requires un-normalized input images with clip_vales=(0, max_value).")

        if preprocessing is not None:
            raise ValueError("This estimator does not support `preprocessing`.")
        if self.postprocessing_defences is not None:
            raise ValueError("This estimator does not support `postprocessing_defences`.")

        self._model: torch.nn.Module
        self._model.to(self._device)
        self._model.eval()

    @property
    def native_label_is_pytorch_format(self) -> bool:
        """
        Are the native labels in PyTorch format [x1, y1, x2, y2]?
        """
        return True

    @property
    def model(self) -> "torch.nn.Module":
        return self._model

    @property
    def input_shape(self) -> Tuple[int, ...]:
        """
        Return the shape of one input sample.

        :return: Shape of one input sample.
        """
        return self._input_shape

    @property
    def optimizer(self) -> Optional["torch.optim.Optimizer"]:
        """
        Return the optimizer.

        :return: The optimizer.
        """
        return self._optimizer

    @property
    def attack_losses(self) -> Tuple[str, ...]:
        """
        Return the combination of strings of the loss components.

        :return: The combination of strings of the loss components.
        """
        return self._attack_losses

    @property
    def device(self) -> "torch.device":
        """
        Get current used device.

        :return: Current used device.
        """
        return self._device

    def _get_losses(
        self, x: np.ndarray, y: List[Dict[str, Union[np.ndarray, "torch.Tensor"]]]
    ) -> Tuple[Dict[str, "torch.Tensor"], List["torch.Tensor"], List["torch.Tensor"]]:
        """
        Get the loss tensor output of the model including all preprocessing.

        :param x: Samples of shape (nb_samples, height, width, nb_channels).
        :param y: Target values of format `List[Dict[Tensor]]`, one for each input image. The fields of the Dict are as
                  follows:

                  - boxes (FloatTensor[N, 4]): the boxes in [x1, y1, x2, y2] format, with 0 <= x1 < x2 <= W and
                                               0 <= y1 < y2 <= H.
                  - labels (Int64Tensor[N]): the labels for each image
        :return: Loss gradients of the same shape as `x`.
        """
        import torch
        import torchvision

        self._model.train()

        # Apply preprocessing
        if self.all_framework_preprocessing:

            if y is not None and isinstance(y, list) and isinstance(y[0]["boxes"], np.ndarray):
                y_tensor = []
                for i, y_i in enumerate(y):
                    y_t = {}
                    y_t["boxes"] = torch.from_numpy(y_i["boxes"]).type(torch.float).to(self.device)
                    y_t["labels"] = torch.from_numpy(y_i["labels"]).type(torch.int64).to(self.device)
                    if "masks" in y_i:
                        y_t["masks"] = torch.from_numpy(y_i["masks"]).type(torch.int64).to(self.device)
                    y_tensor.append(y_t)
            elif y is not None and isinstance(y, dict):
                y_tensor = []
                for i in range(y["boxes"].shape[0]):
                    y_t = {}
                    y_t["boxes"] = y["boxes"][i]
                    y_t["labels"] = y["labels"][i]
                    y_tensor.append(y_t)
            else:
                y_tensor = y  # type: ignore

            transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
            image_tensor_list_grad = []
            y_preprocessed = []
            inputs_t = []

            for i in range(x.shape[0]):
                if isinstance(x, np.ndarray):
                    if self.clip_values is not None:
                        x_grad = transform(x[i] / self.clip_values[1]).to(self.device)
                    else:
                        x_grad = transform(x[i]).to(self.device)
                    x_grad.requires_grad = True
                else:
                    x_grad = x[i].to(self.device)
                    if x_grad.shape[2] < x_grad.shape[0] and x_grad.shape[2] < x_grad.shape[1]:
                        x_grad = torch.permute(x_grad, (2, 0, 1))

                image_tensor_list_grad.append(x_grad)
                x_grad_1 = torch.unsqueeze(x_grad, dim=0)
                x_preprocessed_i, y_preprocessed_i = self._apply_preprocessing(
                    x_grad_1, y=[y_tensor[i]], fit=False, no_grad=False
                )
                for i_preprocessed in range(x_preprocessed_i.shape[0]):
                    inputs_t.append(x_preprocessed_i[i_preprocessed])
                    y_preprocessed.append(y_preprocessed_i[i_preprocessed])

        elif isinstance(x, np.ndarray):
            x_preprocessed, y_preprocessed = self._apply_preprocessing(x, y=y, fit=False, no_grad=True)

            if y_preprocessed is not None and isinstance(y_preprocessed[0]["boxes"], np.ndarray):
                y_preprocessed_tensor = []
                for i, y_i in enumerate(y_preprocessed):
                    y_preprocessed_t = {}
                    y_preprocessed_t["boxes"] = torch.from_numpy(y_i["boxes"]).type(torch.float).to(self.device)
                    y_preprocessed_t["labels"] = torch.from_numpy(y_i["labels"]).type(torch.int64).to(self.device)
                    if "masks" in y_i:
                        y_preprocessed_t["masks"] = torch.from_numpy(y_i["masks"]).type(torch.uint8).to(self.device)
                    y_preprocessed_tensor.append(y_preprocessed_t)
                y_preprocessed = y_preprocessed_tensor

            transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
            image_tensor_list_grad = []

            for i in range(x_preprocessed.shape[0]):
                if self.clip_values is not None:
                    x_grad = transform(x_preprocessed[i] / self.clip_values[1]).to(self.device)
                else:
                    x_grad = transform(x_preprocessed[i]).to(self.device)
                x_grad.requires_grad = True
                image_tensor_list_grad.append(x_grad)

            inputs_t = image_tensor_list_grad

        else:
            raise NotImplementedError("Combination of inputs and preprocessing not supported.")

        if isinstance(y_preprocessed, np.ndarray):
            labels_t = torch.from_numpy(y_preprocessed).to(self.device)  # type: ignore
        else:
            labels_t = y_preprocessed  # type: ignore

        output = self._model(inputs_t, labels_t)

        return output, inputs_t, image_tensor_list_grad

    def loss_gradient(  # pylint: disable=W0613
        self, x: np.ndarray, y: List[Dict[str, Union[np.ndarray, "torch.Tensor"]]], **kwargs
    ) -> np.ndarray:
        """
        Compute the gradient of the loss function w.r.t. `x`.

        :param x: Samples of shape (nb_samples, height, width, nb_channels).
        :param y: Target values of format `List[Dict[Tensor]]`, one for each input image. The
                  fields of the Dict are as follows:

                  - boxes (FloatTensor[N, 4]): the predicted boxes in [x1, y1, x2, y2] format, with values
                    between 0 and H and 0 and W
                  - labels (Int64Tensor[N]): the predicted labels for each image
                  - scores (Tensor[N]): the scores or each prediction.
        :return: Loss gradients of the same shape as `x`.
        """
        import torch

        grad_list = []

        # Adding this loop because torch==[1.7, 1.8] and related versions of torchvision do not allow loss gradients at
        #  the input for batches larger than 1 anymore for PyTorch FasterRCNN because of a view created by torch or
        #  torchvision. This loop should be revisited with later releases of torch and removed once it becomes
        #  unnecessary.
        for i in range(x.shape[0]):

            x_i = x[[i]]
            y_i = [y[i]]

            output, inputs_t, image_tensor_list_grad = self._get_losses(x=x_i, y=y_i)

            # Compute the gradient and return
            loss = None
            for loss_name in self.attack_losses:
                if loss is None:
                    loss = output[loss_name]
                else:
                    loss = loss + output[loss_name]

            # Clean gradients
            self._model.zero_grad()

            # Compute gradients
            loss.backward(retain_graph=True)  # type: ignore

            if isinstance(x, np.ndarray):
                for img in image_tensor_list_grad:
                    if img.grad is not None:
                        gradients = img.grad.cpu().numpy().copy()
                    else:
                        raise ValueError("Gradient term in PyTorch model is `None`.")
                    grad_list.append(gradients)
            else:
                for img in inputs_t:
                    if img.grad is not None:
                        gradients = img.grad.copy()
                    else:
                        raise ValueError("Gradient term in PyTorch model is `None`.")
                    grad_list.append(gradients)

        if isinstance(x, np.ndarray):
            grads = np.stack(grad_list, axis=0)
            grads = np.transpose(grads, (0, 2, 3, 1))
        else:
            grads = torch.stack(grad_list, dim=0)
            grads = grads.premute(0, 2, 3, 1)

        if self.clip_values is not None:
            grads = grads / self.clip_values[1]

        if not self.all_framework_preprocessing:
            grads = self._apply_preprocessing_gradient(x, grads)

        assert grads.shape == x.shape

        return grads

    def predict(self, x: np.ndarray, batch_size: int = 128, **kwargs) -> List[Dict[str, np.ndarray]]:
        """
        Perform prediction for a batch of inputs.

        :param x: Samples of shape NCHW or NHWC.
        :param batch_size: Batch size.
        :return: Predictions of format `List[Dict[str, np.ndarray]]`, one for each input image. The fields of the Dict
                 are as follows:

                 - boxes [N, 4]: the boxes in [x1, y1, x2, y2] format, with 0 <= x1 < x2 <= W and 0 <= y1 < y2 <= H.
                 - labels [N]: the labels for each image
                 - scores [N]: the scores or each prediction.
        """
        import torch
        import torchvision

        # Set model to evaluation mode
        self._model.eval()

        # Apply preprocessing
        x_preprocessed, _ = self._apply_preprocessing(x, y=None, fit=False)

        # Convert samples into tensors
        transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])

        if self.clip_values is not None:
            norm_factor = self.clip_values[1]
        else:
            norm_factor = 1.0

        if self.channels_first:
            x_preprocessed = [torch.from_numpy(x_i / norm_factor).to(self.device) for x_i in x_preprocessed]
        else:
            x_preprocessed = [transform(x_i / norm_factor).to(self.device) for x_i in x_preprocessed]

        predictions: List[Dict[str, np.ndarray]] = []

        # Run prediction
        num_batch = int(np.ceil(len(x_preprocessed) / float(batch_size)))
        for m in range(num_batch):
            # Batch using indices
            i_batch = x_preprocessed[m * batch_size : (m + 1) * batch_size]

            with torch.no_grad():
                predictions_x1y1x2y2 = self._model(i_batch)

            for prediction_x1y1x2y2 in predictions_x1y1x2y2:
                prediction = {}

                prediction["boxes"] = prediction_x1y1x2y2["boxes"].detach().cpu().numpy()
                prediction["labels"] = prediction_x1y1x2y2["labels"].detach().cpu().numpy()
                prediction["scores"] = prediction_x1y1x2y2["scores"].detach().cpu().numpy()
                if "masks" in prediction_x1y1x2y2:
                    prediction["masks"] = prediction_x1y1x2y2["masks"].detach().cpu().numpy().squeeze()

                predictions.append(prediction)

        return predictions

    def fit(  # pylint: disable=W0221
        self,
        x: np.ndarray,
        y: List[Dict[str, Union[np.ndarray, "torch.Tensor"]]],
        batch_size: int = 128,
        nb_epochs: int = 10,
        drop_last: bool = False,
        scheduler: Optional["torch.optim.lr_scheduler._LRScheduler"] = None,
        **kwargs,
    ) -> None:
        """
        Fit the classifier on the training set `(x, y)`.

        :param x: Samples of shape NCHW or NHWC.
        :param y: Target values of format `List[Dict[Tensor]]`, one for each input image. The
                  fields of the Dict are as follows:

                  - boxes (FloatTensor[N, 4]): the predicted boxes in [x1, y1, x2, y2] format, with values
                    between 0 and H and 0 and W
                  - labels (Int64Tensor[N]): the predicted labels for each image
                  - scores (Tensor[N]): the scores or each prediction.
        :param batch_size: Size of batches.
        :param nb_epochs: Number of epochs to use for training.
        :param drop_last: Set to ``True`` to drop the last incomplete batch, if the dataset size is not divisible by
                          the batch size. If ``False`` and the size of dataset is not divisible by the batch size, then
                          the last batch will be smaller. (default: ``False``)
        :param scheduler: Learning rate scheduler to run at the start of every epoch.
        :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
                       and providing it takes no effect.
        """
        import torch
        import torchvision

        # Set model to train mode
        self._model.train()

        if self._optimizer is None:  # pragma: no cover
            raise ValueError("An optimizer is needed to train the model, but none for provided.")

        # Apply preprocessing
        x_preprocessed, y_preprocessed = self._apply_preprocessing(x, y=y, fit=False, no_grad=True)

        # Convert samples into tensors
        transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])

        if self.clip_values is not None:
            norm_factor = self.clip_values[1]
        else:
            norm_factor = 1.0

        if self.channels_first:
            x_preprocessed = torch.from_numpy(x_preprocessed / norm_factor).to(self.device)
        else:
            x_preprocessed = torch.stack([transform(x_i / norm_factor) for x_i in x_preprocessed]).to(self.device)

        # Convert labels into tensors, if needed
        if isinstance(y_preprocessed[0]["boxes"], np.ndarray):
            y_preprocessed_tensor = []
            for y_i in y_preprocessed:
                y_preprocessed_t = {
                    "boxes": torch.from_numpy(y_i["boxes"]).type(torch.float).to(self.device),
                    "labels": torch.from_numpy(y_i["labels"]).type(torch.int64).to(self.device),
                }
                if "masks" in y_i:
                    y_preprocessed_t["masks"] = torch.from_numpy(y_i["masks"]).type(torch.uint8).to(self.device)
                y_preprocessed_tensor.append(y_preprocessed_t)
            y_preprocessed = y_preprocessed_tensor
        y_preprocessed = np.asarray(y_preprocessed)

        num_batch = len(x_preprocessed) / float(batch_size)
        if drop_last:
            num_batch = int(np.floor(num_batch))
        else:
            num_batch = int(np.ceil(num_batch))
        ind = np.arange(len(x_preprocessed))

        # Start training
        for _ in range(nb_epochs):
            # Shuffle the examples
            np.random.shuffle(ind)

            # Train for one epoch
            for m in range(num_batch):
                i_batch = x_preprocessed[ind[m * batch_size : (m + 1) * batch_size]]
                o_batch = y_preprocessed[ind[m * batch_size : (m + 1) * batch_size]]

                # Zero the parameter gradients
                self._optimizer.zero_grad()

                # Form the loss function
                loss_components = self._model(i_batch, o_batch)
                if isinstance(loss_components, dict):
                    loss = sum(loss_components.values())
                else:
                    loss = loss_components

                # Do training
                loss.backward()  # type: ignore
                self._optimizer.step()

            if scheduler is not None:
                scheduler.step()

    def get_activations(
        self, x: np.ndarray, layer: Union[int, str], batch_size: int, framework: bool = False
    ) -> np.ndarray:
        raise NotImplementedError

    def compute_losses(
        self, x: np.ndarray, y: List[Dict[str, Union[np.ndarray, "torch.Tensor"]]]
    ) -> Dict[str, np.ndarray]:
        """
        Compute all loss components.

        :param x: Samples of shape (nb_samples, nb_features) or (nb_samples, nb_pixels_1, nb_pixels_2,
                  nb_channels) or (nb_samples, nb_channels, nb_pixels_1, nb_pixels_2).
        :param y: Target values of format `List[Dict[Tensor]]`, one for each input image. The
                  fields of the Dict are as follows:

                  - boxes (FloatTensor[N, 4]): the predicted boxes in [x1, y1, x2, y2] format, with values
                    between 0 and H and 0 and W
                  - labels (Int64Tensor[N]): the predicted labels for each image
                  - scores (Tensor[N]): the scores or each prediction.
        :return: Dictionary of loss components.
        """
        output_tensor, _, _ = self._get_losses(x=x, y=y)
        output = {}
        for key, value in output_tensor.items():
            output[key] = value.detach().cpu().numpy()
        return output

    def compute_loss(  # type: ignore
        self, x: np.ndarray, y: List[Dict[str, Union[np.ndarray, "torch.Tensor"]]], **kwargs
    ) -> Union[np.ndarray, "torch.Tensor"]:
        """
        Compute the loss of the neural network for samples `x`.

        :param x: Samples of shape (nb_samples, nb_features) or (nb_samples, nb_pixels_1, nb_pixels_2,
                  nb_channels) or (nb_samples, nb_channels, nb_pixels_1, nb_pixels_2).
        :param y: Target values of format `List[Dict[Tensor]]`, one for each input image. The
                  fields of the Dict are as follows:

                  - boxes (FloatTensor[N, 4]): the predicted boxes in [x1, y1, x2, y2] format, with values
                    between 0 and H and 0 and W
                  - labels (Int64Tensor[N]): the predicted labels for each image
                  - scores (Tensor[N]): the scores or each prediction.
        :return: Loss.
        """
        import torch

        output, _, _ = self._get_losses(x=x, y=y)

        # Compute the gradient and return
        loss = None
        for loss_name in self.attack_losses:
            if loss is None:
                loss = output[loss_name]
            else:
                loss = loss + output[loss_name]

        assert loss is not None

        if isinstance(x, torch.Tensor):
            return loss

        return loss.detach().cpu().numpy()


# -*- coding: utf-8 -*-
"""
author: Diogo Silva
notes: Boruvka implementation based on Sousa's "A Generic and Highly Efficient Parallel Variant of Boruvka ’s Algorithm"
       connected components from Boruvka
"""

import numpy as np
from MyML.helper.scan import exprefixsumNumbaSingle as ex_prefix_sum_cpu, exprefixsumNumba as ex_prefix_sum_cpu2,\
                             scan_gpu as ex_prefix_sum_gpu
from MyML.graph.mst import findMinEdgeNumba, removeMirroredNumba, initColorsNumba, propagateColorsNumba,\
                           buildFlag, countNewEdgesNumba, assignInsertNumba,\
                           compute_cuda_grid_dim, findMinEdge_CUDA, removeMirroredEdges_CUDA, memSet,\
                           initializeColors_CUDA, propagateColors_CUDA, buildFlag_CUDA, countNewEdges_CUDA, assignInsert_CUDA
from numba import jit, cuda, void, boolean, int8, int32, float32

def connected_comps_mst_seq(mst, dest, fe):
    pass

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
#
#                      NUMBA CPU
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

def connected_comps_seq(dest_in, weight_in, firstEdge_in, outDegree_in):
    dest = dest_in
    weight = weight_in
    firstEdge = firstEdge_in
    outDegree = outDegree_in

    n_vertices = firstEdge.size
    n_edges = dest.size

    n_components = n_vertices

    # still need edge_id for conflict resolution in find_minedge
    edge_id = np.arange(n_edges, dtype = dest.dtype)

    #labels = np.empty(n_vertices, dtype = dest.dtype)
    first_iter = True

    # initialize with name top_edge so we can recycle an array between iterations
    top_edge = np.empty(n_components, dtype = dest.dtype)

    final_converged = False
    while(not final_converged):
        vertex_minedge = top_edge
        findMinEdgeNumba(vertex_minedge, weight, firstEdge, outDegree, dest)
        removeMirroredNumba(vertex_minedge, dest)

        # intialize colors of current graph
        colors = np.empty(n_components, dtype = dest.dtype)
        initColorsNumba(vertex_minedge, dest, colors)

        # propagate colors until convergence
        converged = False
        while(not converged):
            converged = propagateColorsNumba(colors)

        # flag marks the vertices that are the representatives of the new supervertices
        # new_vertex will be initialized with he flags
        new_vertex = vertex_minedge
        buildFlag(colors, new_vertex)

        new_n_vertices = ex_prefix_sum_cpu(new_vertex, init = 0)

        if first_iter:
            # first iteration defines labels as the initial colors and updates
            labels = colors.copy()
            first_iter = False

        # update the labels with the new colors
        update_labels_single_pass(labels, colors, new_vertex)

        if new_n_vertices == 1:
            final_converged = True
            break

        # count number of edges for new supervertices and write in new outDegree
        newOutDegree = np.zeros(new_n_vertices, dtype = dest.dtype)
        countNewEdgesNumba(colors, firstEdge, outDegree, dest, new_vertex, newOutDegree)

        # new first edge array for contracted graph
        newFirstEdge = np.empty(newOutDegree.size, dtype = dest.dtype)
        new_n_edges = ex_prefix_sum_cpu2(newOutDegree, newFirstEdge, init = 0)

        # if no edges remain, then MST has converged
        if new_n_edges == 0:
            final_converged = True
            break

        # create arrays for new edges
        new_dest = np.empty(new_n_edges, dtype = dest.dtype)
        new_edge_id = np.empty(new_n_edges, dtype = dest.dtype)
        new_weight = np.empty(new_n_edges, dtype = weight.dtype)
        top_edge = newFirstEdge.copy()

        # assign and insert new edges
        assignInsertNumba(edge_id, dest, weight, firstEdge, outDegree, colors,
                          new_vertex, new_dest, new_edge_id, new_weight, top_edge)

        # delete old graph
        del new_vertex, edge_id, dest, weight, firstEdge, outDegree, colors

        # write new graph
        n_components = newFirstEdge.size
        edge_id = new_edge_id
        dest = new_dest
        weight = new_weight
        firstEdge = newFirstEdge
        outDegree = newOutDegree

    return labels

@jit
def update_labels_numba(labels, update_array):
    """
    This kernel is dual purpose.

    This kernel updates the color of each vertex in the graph.
    The current colors of the vertices are indices for the current components
    in the graph. Each current component in the graph will have a new color.
    The new color of vertex v will be the new color colors[curr_color] of its
    color labels[v] (index of a component). This step happens because the new
    color propagation is representative of the contracted graph. In this phase
    the update_array is the colors array.

    Old components are merged into new ones with new IDs. The labels need to
    have the new IDs. In this phase the update_array is the new_vertex array.
    """
    n_components = labels.size

    for v in range(n_components):
        curr_color = labels[v]
        labels[v] = update_array[curr_color]

@jit
def update_labels_single_pass(labels, colors, new_vertex):
    """
    Does all the updates on a single pass
    """
    n_components = labels.size

    for v in range(n_components):
        curr_color = labels[v]
        new_color = colors[curr_color]
        new_color_id = new_vertex[new_color]
        labels[v] = new_color_id

def connected_comps_gpu(dest_in, weight_in, firstEdge_in, outDegree_in,
                        MAX_TPB = 512, stream = None):
    if stream is None:
        myStream = cuda.stream()
    else:
        myStream = stream

    dest = cuda.to_device(dest_in, stream = myStream)
    weight = cuda.to_device(weight_in, stream = myStream)
    firstEdge = cuda.to_device(firstEdge_in, stream = myStream)
    outDegree = cuda.to_device(outDegree_in, stream = myStream)

    n_vertices = firstEdge.size
    n_edges = dest.size

    n_components = n_vertices

    # still need edge_id for conflict resolution in find_minedge
    edge_id = cuda.to_device(np.arange(n_edges, dtype = dest.dtype), stream = myStream)

    #labels = np.empty(n_vertices, dtype = dest.dtype)
    first_iter = True

    # initialize with name top_edge so we can recycle an array between iterations
    top_edge = cuda.device_array(n_components, dtype = dest.dtype, stream = myStream)
    labels = cuda.device_array(n_components, dtype = dest.dtype, stream = myStream)

    converged = cuda.device_array(1, dtype = np.int8, stream = myStream)

    gridDimLabels = compute_cuda_grid_dim(n_components, MAX_TPB)
    gridDim = compute_cuda_grid_dim(n_components, MAX_TPB)

    final_converged = False
    while(not final_converged):
        vertex_minedge = top_edge

        findMinEdge_CUDA[gridDim, MAX_TPB, myStream](weight, firstEdge, outDegree, vertex_minedge, dest)

        removeMirroredEdges_CUDA[gridDim, MAX_TPB, myStream](dest, vertex_minedge)

        colors = cuda.device_array(shape = n_components, dtype = np.int32, stream = myStream)
        initializeColors_CUDA[gridDim, MAX_TPB, myStream](dest, vertex_minedge, colors)

        # propagate colors until convergence
        propagateConverged = False
        while(not propagateConverged):
            propagateColors_CUDA[gridDim, MAX_TPB, myStream](colors, converged)
            converged_num = converged.getitem(0, stream = myStream)
            propagateConverged = True if converged_num == 1 else False

        # first we build the flags in the new_vertex array
        new_vertex = vertex_minedge # reuse the vertex_minedge array as the new new_vertex
        buildFlag_CUDA[gridDim, MAX_TPB, myStream](colors, new_vertex)

        # new_n_vertices is the number of vertices of the new contracted graph
        new_n_vertices = ex_prefix_sum_gpu(new_vertex, MAX_TPB = MAX_TPB, stream = myStream).getitem(0, stream = myStream)
        new_n_vertices = int(new_n_vertices)

        if first_iter:
            # first iteration defines labels as the initial colors and updates
            labels.copy_to_device(colors, stream = myStream)
            first_iter = False

        # other iterations update the labels with the new colors
        update_labels_single_pass_cuda[gridDimLabels, MAX_TPB, myStream](labels, colors, new_vertex)

        if new_n_vertices == 1:
            final_converged = True
            del new_vertex
            break

        newGridDim = compute_cuda_grid_dim(n_components, MAX_TPB)

        # count number of edges for new supervertices and write in new outDegree
        newOutDegree = cuda.device_array(shape = new_n_vertices, dtype = np.int32, stream = myStream)
        memSet[newGridDim, MAX_TPB, myStream](newOutDegree, 0) # zero the newOutDegree array
        countNewEdges_CUDA[gridDim, MAX_TPB, myStream](colors, firstEdge, outDegree, dest, new_vertex, newOutDegree)

        # new first edge array for contracted graph
        newFirstEdge = cuda.device_array_like(newOutDegree, stream = myStream)
        newFirstEdge.copy_to_device(newOutDegree, stream = myStream) # copy newOutDegree to newFirstEdge
        new_n_edges = ex_prefix_sum_gpu(newFirstEdge, MAX_TPB = MAX_TPB, stream = myStream)

        new_n_edges = new_n_edges.getitem(0, stream = myStream)
        new_n_edges = int(new_n_edges)

        # if no edges remain, then MST has converged
        if new_n_edges == 0:
            final_converged = True
            del newOutDegree, newFirstEdge, new_vertex
            break

        # create arrays for new edges
        new_dest = cuda.device_array(new_n_edges, dtype = np.int32, stream = myStream)
        new_edge_id = cuda.device_array(new_n_edges, dtype = np.int32, stream = myStream)
        new_weight = cuda.device_array(new_n_edges, dtype = weight.dtype, stream = myStream)

        top_edge = cuda.device_array_like(newFirstEdge, stream = myStream)
        top_edge.copy_to_device(newFirstEdge, stream = myStream)

        # assign and insert new edges
        assignInsert_CUDA[gridDim, MAX_TPB, myStream](edge_id, dest, weight, firstEdge,
                                            outDegree, colors, new_vertex,
                                            new_dest, new_edge_id, new_weight, top_edge)

        # delete old graph
        del new_vertex, edge_id, dest, weight, firstEdge, outDegree, colors

        # write new graph
        n_components = newFirstEdge.size
        edge_id = new_edge_id
        dest = new_dest
        weight = new_weight
        firstEdge = newFirstEdge
        outDegree = newOutDegree
        gridDim = newGridDim

    returnLabels = labels.copy_to_host()

    del dest, weight, edge_id, firstEdge, outDegree, converged, labels

    return returnLabels

@cuda.jit
def update_labels_cuda(labels, update_array):
    """
    CUDA version of update_labels.
    """

    v = cuda.grid(1)
    n_components = labels.size

    if v >= n_components:
        return

    curr_color = labels[v]
    labels[v] = update_array[curr_color]

@cuda.jit
def update_labels_single_pass_cuda(labels, colors, new_vertex):
    """
    Does all the updates on a single pass
    """
    v = cuda.grid(1)
    n_components = labels.size

    if v >= n_components:
        return

    curr_color = labels[v]
    new_color = colors[curr_color]
    new_color_id = new_vertex[new_color]
    labels[v] = new_color_id


#PySCF input file for CuO calculations
import json
from pyscf import gto,scf,mcscf, fci,lo,ci,cc,lib
from pyscf.scf import ROHF,ROKS,UHF,UKS, addons
import numpy as np
import pandas as pd

df=json.load(open("trail.json"))
charge=0

S=[1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1]
symm_dict=[
#3d10 sector
{'A1':(5,6),'E1x':(3,3),'E1y':(3,3),'E2x':(1,0),'E2y':(1,0)}, #GS
'''
#3d10 sector
{'A1':(5,5),'E1x':(3,3),'E1y':(3,2),'E2x':(1,1),'E2y':(1,1)}, #GS
{'A1':(5,4),'E1x':(3,3),'E1y':(3,3),'E2x':(1,1),'E2y':(1,1)}, #(z -> pi)
{'A1':(5,6),'E1x':(3,2),'E1y':(3,2),'E2x':(1,1),'E2y':(1,1)}, #(pi -> s, J~-small)

#3d9 sector
{'A1':(6,5),'E1x':(3,3),'E1y':(2,3),'E2x':(1,1),'E2y':(1,0)}, #(dd -> s, Jh~1/4)
{'A1':(6,5),'E1x':(3,3),'E1y':(3,2),'E2x':(1,1),'E2y':(0,1)}, #(dd -> s, Jh~-1/4)

#3d10 sector
{'A1':(6,5),'E1x':(3,2),'E1y':(2,3),'E2x':(1,1),'E2y':(1,1)}, #(pi -> s)
{'A1':(6,4),'E1x':(3,3),'E1y':(2,3),'E2x':(1,1),'E2y':(1,1)}, #(z -> s)

#3d9 sector
{'A1':(5,5),'E1x':(3,3),'E1y':(3,3),'E2x':(1,1),'E2y':(1,0)}, #(dd -> pi)
]
'''
]

'''
#Mirrored states
S=[1,1,1,1,1,1]
symm_dict=[
#3d10 sector
{'A1':(5,5),'E1x':(3,2),'E1y':(3,3),'E2x':(1,1),'E2y':(1,1)}, #GS

#3d9 sector
{'A1':(6,5),'E1x':(2,3),'E1y':(3,3),'E2x':(1,0),'E2y':(1,1)}, #(dd -> s, Jh~1/4)
{'A1':(6,5),'E1x':(3,2),'E1y':(3,3),'E2x':(0,1),'E2y':(1,1)}, #(dd -> s, Jh~-1/4)

#3d10 sector
{'A1':(6,5),'E1x':(2,3),'E1y':(3,2),'E2x':(1,1),'E2y':(1,1)}, #(pi -> s)
{'A1':(6,4),'E1x':(2,3),'E1y':(3,3),'E2x':(1,1),'E2y':(1,1)}, #(z -> s)

#3d9 sector
{'A1':(5,5),'E1x':(3,3),'E1y':(3,3),'E2x':(1,0),'E2y':(1,1)}, #(dd -> pi)
]
'''

datacsv={}

for nm in['run','method','basis','pseudopotential','bond-length','S','E','conv']:
  datacsv[nm]=[]

for run in [0]:
  for r in [1.725]:
    for method in ['UB3LYP']:
      for basis in ['vdz','vtz']:
        for el in ['Cu']:
          molname=el+'O'
          mol=gto.Mole()

          mol.ecp={}
          mol.basis={}
          for e in [el,'O']:
            mol.ecp[e]=gto.basis.parse_ecp(df[e]['ecp'])
            mol.basis[e]=gto.basis.parse(df[e][basis])
          mol.charge=charge
          mol.spin=S[run]
          mol.build(atom="%s 0. 0. 0.; O 0. 0. %g"%(el,r),verbose=4,symmetry=True)

          if("U" in method):
            if("HF" in method):
              m=UHF(mol)
            else:
              m=UKS(mol)
              m.xc=method[1:]
          else:
            if(method=="ROHF"):
              m=ROHF(mol)
            else:
              m=ROKS(mol)
              m.xc=method

          if basis=='vdz':
            #m=m.newton()
            m.chkfile= '3d8_vdz.chk' #el+basis+"_r"+str(r)+"_s"+str(S[run])+"_"+method+"_"+str(run)+".chk"
            m.irrep_nelec = symm_dict[run]
            m.max_cycle=100
            m = addons.remove_linear_dep_(m)
            m.conv_tol=1e-5
            m.diis=scf.ADIIS()
            total_energy=m.kernel()

            #Compute the Mulliken orbital occupancies...
            m.analyze()
            #m.stability(external=True)
            assert(np.sum(m.mo_occ)==25)

          #Once we get past the vdz basis, just read-in the existing chk file...
          else:
            dm= m.from_chk('3d8_vdz.chk') #m.from_chk(el+'vdz'+"_r"+str(r)+"_s"+str(S[run])+"_"+method+"_"+str(run)+".chk")
            m.chkfile= '3d8_vtz.chk' #el+basis+"_r"+str(r)+"_s"+str(S[run])+"_"+method+"_"+str(run)+".chk"
            m.irrep_nelec = symm_dict[run]
            m.max_cycle=100
            m = addons.remove_linear_dep_(m)
            m.conv_tol=1e-5
            m.diis=scf.ADIIS()
            total_energy=m.kernel(dm)

            m.analyze()
            #m.stability(external=True)
            assert(np.sum(m.mo_occ)==25)

          '''
          datacsv['run'].append(run)
          datacsv['bond-length'].append(r)
          datacsv['S'].append(S[run])
          datacsv['method'].append(method)
          datacsv['basis'].append(basis)
          datacsv['pseudopotential'].append('trail')
          datacsv['E'].append(total_energy)
          datacsv['conv'].append(m.converged)
          pd.DataFrame(datacsv).to_csv("cuo_additional.csv",index=False)
          '''


from __future__ import absolute_import

import mock

from datetime import datetime, timedelta

from changes.constants import Result, Status
from changes.config import db
from changes.jobs.sync_job import has_timed_out, sync_job
from changes.models import FailureReason, ItemOption, ItemStat, Job, Step, Task
from changes.testutils import TestCase

class HasTimedOutTest(TestCase):
    def test_simple(self):
        project = self.create_project()
        plan = self.create_plan()
        plan.projects.append(project)

        build = self.create_build(project=project)
        job = self.create_job(build=build, status=Status.queued)
        job_plan = self.create_job_plan(job, plan)

        option = ItemOption(
            item_id=plan.id,
            name='build.timeout',
            value='5',
        )
        db.session.add(option)
        db.session.commit()

        assert not has_timed_out(job, job_plan)

        job.status = Status.in_progress
        job.date_started = datetime.utcnow()
        db.session.add(job)
        db.session.commit()

        assert not has_timed_out(job, job_plan)

        job.date_started = datetime.utcnow() - timedelta(seconds=400)
        db.session.add(job)
        db.session.commit()

        assert has_timed_out(job, job_plan)

        option.value = '0'
        db.session.add(option)
        db.session.commit()

        assert not has_timed_out(job, job_plan)

        option.value = '500'
        db.session.add(option)
        db.session.commit()

        assert not has_timed_out(job, job_plan)

class SyncJobTest(TestCase):
    def setUp(self):
        super(SyncJobTest, self).setUp()
        self.project = self.create_project()
        self.build = self.create_build(project=self.project)
        self.job = self.create_job(build=self.build)
        self.jobphase = self.create_jobphase(self.job)
        self.jobstep = self.create_jobstep(self.jobphase)

        self.task = self.create_task(
            parent_id=self.build.id,
            task_id=self.job.id,
            task_name='sync_job',
        )

        self.plan = self.create_plan()
        self.plan.projects.append(self.project)
        self.step = self.create_step(self.plan, implementation='test', order=0)
        self.jobplan = self.create_job_plan(self.job, self.plan)

    @mock.patch('changes.jobs.sync_job.queue.delay')
    @mock.patch.object(Step, 'get_implementation')
    def test_in_progress(self, get_implementation,
                         queue_delay):
        implementation = mock.Mock()
        get_implementation.return_value = implementation

        build, job, task = self.build, self.job, self.task

        self.create_task(
            task_name='sync_job_step',
            task_id=job.phases[0].steps[0].id,
            parent_id=job.id,
            status=Status.in_progress,
        )

        def mark_in_progress(job):
            job.status = Status.in_progress

        implementation.update.side_effect = mark_in_progress

        sync_job(
            job_id=job.id.hex,
            task_id=job.id.hex,
            parent_task_id=build.id.hex
        )

        get_implementation.assert_called_once_with()

        implementation.update.assert_called_once_with(
            job=self.job,
        )

        queue_delay.assert_any_call('sync_job', kwargs={
            'job_id': job.id.hex,
            'task_id': job.id.hex,
            'parent_task_id': build.id.hex,
        }, countdown=5)

        task = Task.query.get(task.id)

        assert task.status == Status.in_progress

    @mock.patch('changes.jobs.sync_job.fire_signal')
    @mock.patch('changes.jobs.sync_job.queue.delay')
    @mock.patch.object(Step, 'get_implementation')
    def test_finished(self, get_implementation, queue_delay,
                      mock_fire_signal):
        implementation = mock.Mock()
        get_implementation.return_value = implementation

        assert self.jobplan

        build, job, task = self.build, self.job, self.task

        step = job.phases[0].steps[0]

        self.create_task(
            task_name='sync_job_step',
            task_id=step.id,
            parent_id=job.id,
            status=Status.finished,
        )
        self.create_test(job)
        self.create_test(job)

        db.session.add(ItemStat(item_id=step.id, name='tests_missing', value=1))
        db.session.add(ItemStat(item_id=step.id, name='lines_covered', value=10))
        db.session.add(ItemStat(item_id=step.id, name='lines_uncovered', value=25))
        db.session.commit()

        sync_job(
            job_id=job.id.hex,
            task_id=job.id.hex,
            parent_task_id=build.id.hex,
        )

        job = Job.query.get(job.id)

        assert job.status == Status.finished

        queue_delay.assert_any_call('update_project_plan_stats', kwargs={
            'project_id': self.project.id.hex,
            'plan_id': self.plan.id.hex,
        }, countdown=1)

        mock_fire_signal.delay.assert_any_call(
            signal='job.finished',
            kwargs={'job_id': job.id.hex},
        )

        task = Task.query.get(task.id)

        assert task.status == Status.finished

        stat = ItemStat.query.filter(
            ItemStat.item_id == job.id,
            ItemStat.name == 'tests_missing',
        ).first()
        assert stat.value == 1

        stat = ItemStat.query.filter(
            ItemStat.item_id == job.id,
            ItemStat.name == 'lines_covered',
        ).first()
        assert stat.value == 10

        stat = ItemStat.query.filter(
            ItemStat.item_id == job.id,
            ItemStat.name == 'lines_uncovered',
        ).first()
        assert stat.value == 25

    @mock.patch('changes.jobs.sync_job.has_timed_out')
    @mock.patch('changes.jobs.sync_job.queue.delay')
    @mock.patch.object(Step, 'get_implementation')
    def test_timed_out(self, get_implementation,
                       queue_delay, mock_has_timed_out):
        implementation = mock.Mock()
        get_implementation.return_value = implementation

        build, job, task = self.build, self.job, self.task

        other_build = self.create_build(self.project, status=Status.in_progress)
        other_job = self.create_job(other_build)
        other_jobphase = self.create_jobphase(other_job)
        other_jobstep = self.create_jobstep(other_jobphase)

        mock_has_timed_out.return_value = True

        sync_job(
            job_id=job.id.hex,
            task_id=job.id.hex,
            parent_task_id=build.id.hex
        )

        mock_has_timed_out.assert_called_once_with(job, self.jobplan)

        get_implementation.assert_called_once_with()
        implementation.cancel.assert_called_once_with(
            job=self.job,
        )

        assert job.result == Result.failed
        assert job.status == Status.finished

        assert self.jobphase.result == job.result
        assert self.jobphase.status == job.status

        assert self.jobstep.result == job.result
        assert self.jobstep.status == job.status

        assert FailureReason.query.filter(
            FailureReason.step_id == self.jobstep.id,
            FailureReason.reason == 'timeout',
        )

        # ensure we haven't updated an incorrect jobstep as well
        assert other_jobstep.status == Status.in_progress
        assert other_jobphase.status == Status.in_progress
        assert other_job.status == Status.in_progress
        assert other_build.status == Status.in_progress


#!/usr/bin/env python
# coding: utf-8

from django.utils.safestring import mark_safe
from table.columns import Column

class CheckboxColumn(Column):
    def __init__(self, header=None, **kwargs):
        kwargs["field"] = None
        kwargs["safe"] = False
        kwargs["sortable"] = False
        kwargs["searchable"] = False
        super(CheckboxColumn, self).__init__(header=header, **kwargs)

    def render(self, obj):
        return mark_safe('<input type="checkbox">')


#
# BSD 3-Clause License
#
# Copyright (c) 2017-2018, plures
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

from xnd import xnd
from gumath import sin
import sys, time
import unittest
import argparse

try:
    import numpy as np
except ImportError:
    np = None

class TestCall(unittest.TestCase):

    def test_sin(self):
        test_cases = [
          ([[float(i) for i in range(1000000)],
            [float(i+1) for i in range(1000000)]],
           "2 * 1000000 * float64", "float64"),

          (1000000 * [[float(i+1) for i in range(2)]],
           "1000000 * 2 * float64", "float64"),

          ([[float(i) for i in range(1000000)],
            [float(i+1) for i in range(1000000)]],
           "2 * 1000000 * float32", "float32"),

          (1000000 * [[float(i+1) for i in range(2)]],
           "1000000 * 2 * float32", "float32")]

        for lst, t, dtype in test_cases:
            x = xnd(lst, type=t)

            start = time.time()
            y = sin(x)
            end = time.time()
            sys.stderr.write("\ngumath: time=%.3f\n" % (end-start))

            if np is not None:
                a = np.array(lst, dtype=dtype)

                start = time.time()
                b = np.sin(lst)
                end = time.time()
                sys.stderr.write("numpy: time=%.3f\n" % (end-start))

                np.testing.assert_almost_equal(y, b, 7)

ALL_TESTS = [
  TestCall
]

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-f", "--failfast", action="store_true",
                        help="stop the test run on first error")
    args = parser.parse_args()

    suite = unittest.TestSuite()
    loader = unittest.TestLoader()

    for case in ALL_TESTS:
        s = loader.loadTestsFromTestCase(case)
        suite.addTest(s)

    runner = unittest.TextTestRunner(failfast=args.failfast, verbosity=2)
    result = runner.run(suite)
    ret = not result.wasSuccessful()

    sys.exit(ret)


import csv
from pyspark import SparkContext, PickleSerializer, MarshalSerializer
import datetime
import time
import os
import os.path

def To_CSV(data):
    """Write out CSV lines"""
    dir = r"F:/Do an chuyen nganh/TT-TeamRC/DEMO-PROJECT.git/trunk/Save/LAST_ACTION/"
    if not os.path.exists(dir):
        os.mkdir(dir)
    # arrayList=[]
    # for line in data:
    #     array = []
    #     array.append(int(line[0]))
    #     array.append(line[1])
    #     print(array)
    #     arrayList.append(array)
    with open(os.path.join(dir, "Last_Action" + '.csv'), "a", newline='') as csvfile:
        writer = csv.writer(csvfile, delimiter=',',
                            quotechar=',', quoting=csv.QUOTE_MINIMAL)
        writer.writerow(data)

def convert_String_To_TimeStamp(s):
    return time.mktime(time.strptime(s, "%m:%d:%Y %H:%M:%S"))

if __name__ == "__main__":
    sc = SparkContext(appName="last action")
    movies_raw_data = sc.textFile("F:/Do an chuyen nganh/TT-TeamRC/DEMO-PROJECT.git/trunk/Metadata/IMDBMovies.csv")
    small_ratings_raw_data_header = movies_raw_data.take(1)[0]
    movies_data = movies_raw_data.filter(lambda line: line != small_ratings_raw_data_header)
    movies_data = movies_data.map(lambda line: line.split(",")).map(lambda token: (int(token[0]), token[1]))

    events_raw_data = sc.textFile("F:/Do an chuyen nganh/TT-TeamRC/DEMO-PROJECT.git/trunk/UserProfile/userevent.csv")
    events_data = events_raw_data.map(lambda line: line.split(";")).map(lambda x: (x[0], int(x[1]), x[2]))
    listUser = sorted(events_data.map(lambda x: x[0]).distinct().collect())
    print("list User:")
    print(listUser)
    print(events_data.take(10))
    #
    # user_id = str("Account0")
    for user_id in listUser:
        rating_data_for_user_id = events_data.filter(lambda line: user_id in line).map(
            lambda x: (int(x[1]), (x[2]), convert_String_To_TimeStamp(x[2])))
        result = rating_data_for_user_id.takeOrdered(10, key=lambda x: -x[2])
        list = []
        list.append(user_id)
        for line in result:
            for firm in movies_data.collect():
                if line[0] == firm[0]:
                    # array = []
                    # array.append(firm[1] + "|" + line[1])
                    # array.append(line[1])
                    # list.append(array)
                    list.append(firm[1] + "|" + line[1])
                    break

        # result = rating_data_for_user_id.takeOrdered(5)
        # print(result)
        # print("ket qua")
        # print(list)
        To_CSV(list)

    # result = rating_data_for_user_id.sortByKey(True,3,keyfunc=lambda x: str(x[1]))
    # print(result)

    #
    # for line in rating_data_for_user_id.take(5):
    #     print(line)

    # movies_raw_data = sc.textFile("C:/recommendation/ml-1m/movies.dat")
    # movies_data = movies_raw_data.map(lambda line: line.split("::")).map(lambda token: (token[0], token[1]))
    # complete_movies_titles = movies_data.map(lambda x: (int(x[0]), x[1]))
    # result = complete_movies_titles.join(rating_data_for_user_id).takeOrdered(5, key=lambda x: -x[1][1])
    # result = result.map(lambda x: (x[1][0],x[1][1]))
    # writeRecords(result)

    # s="06:11:2017 00:48:30"
    # rs = time.mktime(time.strptime(s,"%m:%d:%Y %H:%M:%S"))
    # print(rs)


"""
Utility to create tile pyramids out of an input raster.

Also provides various options to rescale data if necessary
"""

import click
import os
import rasterio

import mapchete
from mapchete.cli import utils
from mapchete.io import get_best_zoom_level

# ranges from rasterio
# https://github.com/mapbox/rasterio/blob/master/rasterio/dtypes.py#L61
DTYPE_RANGES = {
    'uint8': (0, 255),
    'uint16': (0, 65535),
    'int16': (-32768, 32767),
    'uint32': (0, 4294967295),
    'int32': (-2147483648, 2147483647),
    'float32': (-3.4028235e+38, 3.4028235e+38),
    'float64': (-1.7976931348623157e+308, 1.7976931348623157e+308)
}

@click.command(help="Create a tile pyramid from an input raster dataset.")
@utils.arg_input_raster
@utils.arg_out_dir
@utils.opt_pyramid_type_mercator
@utils.opt_output_format
@utils.opt_resampling_method
@utils.opt_scale_method
@utils.opt_zoom
@utils.opt_bounds
@utils.opt_overwrite
@utils.opt_debug
def pyramid(
    input_raster,
    output_dir,
    pyramid_type=None,
    output_format=None,
    resampling_method=None,
    scale_method=None,
    zoom=None,
    bounds=None,
    overwrite=False,
    debug=False
):
    """Create tile pyramid out of input raster."""
    bounds = bounds if bounds else None
    options = dict(
        pyramid_type=pyramid_type,
        scale_method=scale_method,
        output_format=output_format,
        resampling=resampling_method,
        zoom=zoom,
        bounds=bounds,
        overwrite=overwrite
    )
    raster2pyramid(input_raster, output_dir, options)

def raster2pyramid(input_file, output_dir, options):
    """Create a tile pyramid out of an input raster dataset."""
    pyramid_type = options["pyramid_type"]
    scale_method = options["scale_method"]
    output_format = options["output_format"]
    resampling = options["resampling"]
    zoom = options["zoom"]
    bounds = options["bounds"]
    mode = "overwrite" if options["overwrite"] else "continue"

    # Prepare process parameters
    minzoom, maxzoom = _get_zoom(zoom, input_file, pyramid_type)
    with rasterio.open(input_file, "r") as input_raster:
        output_bands = input_raster.count
        input_dtype = input_raster.dtypes[0]
        output_dtype = input_raster.dtypes[0]
        nodataval = input_raster.nodatavals[0]
        nodataval = nodataval if nodataval else 0
        if output_format == "PNG" and output_bands > 3:
            output_bands = 3
            output_dtype = 'uint8'
        scales_minmax = ()
        if scale_method == "dtype_scale":
            for index in range(1, output_bands+1):
                scales_minmax += (DTYPE_RANGES[input_dtype], )
        elif scale_method == "minmax_scale":
            for index in range(1, output_bands+1):
                band = input_raster.read(index)
                scales_minmax += ((band.min(), band.max()), )
        elif scale_method == "crop":
            for index in range(1, output_bands+1):
                scales_minmax += ((0, 255), )
        if input_dtype == "uint8":
            scale_method = None
            scales_minmax = ()
            for index in range(1, output_bands+1):
                scales_minmax += ((None, None), )

    # Create configuration
    config = dict(
        process="mapchete.processes.pyramid.tilify",
        output={
            "path": output_dir,
            "format": output_format,
            "bands": output_bands,
            "dtype": output_dtype
            },
        pyramid=dict(pixelbuffer=5, grid=pyramid_type),
        scale_method=scale_method,
        scales_minmax=scales_minmax,
        input={"raster": input_file},
        config_dir=os.getcwd(),
        zoom_levels=dict(min=minzoom, max=maxzoom),
        nodataval=nodataval,
        resampling=resampling,
        bounds=bounds,
        baselevel={"zoom": maxzoom, "resampling": resampling},
        mode=mode
    )

    # create process
    with mapchete.open(config, zoom=zoom, bounds=bounds) as mp:
        # prepare output directory
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        # run process
        mp.batch_process(zoom=[minzoom, maxzoom])

def _get_zoom(zoom, input_raster, pyramid_type):
    """Determine minimum and maximum zoomlevel."""
    if not zoom:
        minzoom = 1
        maxzoom = get_best_zoom_level(input_raster, pyramid_type)
    elif len(zoom) == 1:
        minzoom = zoom[0]
        maxzoom = zoom[0]
    elif len(zoom) == 2:
        if zoom[0] < zoom[1]:
            minzoom = zoom[0]
            maxzoom = zoom[1]
        else:
            minzoom = zoom[1]
            maxzoom = zoom[0]
    return minzoom, maxzoom


# -*- coding: utf-8 -*-
"""
Created on Tue May 28 15:53:48 2019

@author: Josh
"""

import spacy

#Define the custom component
def length_component(doc):
    #Get the doc's length
    doc_length = len(doc)
    print("This document is {} tokens long.".format(doc_length))
    #Return the doc
    return doc

#Load the small English model
nlp = spacy.load("en_core_web_sm")

#Add the component first in the pipeline and print the pipe names
nlp.add_pipe(length_component, first=True)
print(nlp.pipe_names)

#Process a text
doc = nlp("This is a sentence.")

Chapter 3 - Processing Pipelines/6 - Simple components.py


from django.db import models

class Pin(models.Model):
	user = models.ForeignKey(PinterestUser, null=True, blank=True)
	board = models.ForeignKey(Board, null=True, blank=True)
	# Always 18 characters, but set to 31 for room to change
	pin_id = models.CharField(max_length=31, unique=True)
	# taken from the image url
	md5 = models.CharField(max_length=32, blank=True, default='')
	description = models.TextField()
	source = models.CharField(max_length=255, blank=True, default='')
	likes = models.IntegerField(default=0, db_index=True)
	repins = models.IntegerField(default=0, db_index=True)
	comments = models.IntegerField(default=0, db_index=True)
	actions = models.IntegerField(default=0, db_index=True)
	price = models.BigIntegerField(default=0)
	currency = models.CharField(max_length=1)
	etag = models.CharField(max_length=63, blank=True)
	pinned = models.DateTimeField(null=True, blank=True)
	created = models.DateTimeField(auto_now_add=True)
	updated = models.DateTimeField(auto_now=True)


""" Clean the dataset """
import argparse
import hashlib
import logging
from urllib.parse import urlparse

import pandas as pd
import nltk
from nltk.corpus import stopwords
nltk.download('punkt')
nltk.download('stopwords')

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def _read_data(filename):
    logger.info('Reading file %s', filename)

    return pd.read_csv(filename)

def _extract_newspaper_uid(filename):
    logger.info('Extracting newspaper uid')
    newspaper_uid = filename.split('_')[0]

    logger.info('Newspaper uid detected: %s', newspaper_uid)

    return newspaper_uid

def _add_newspaper_uid_column(df, newspaper_uid):
    logger.info('Filling newspaper_uid with %s', newspaper_uid)
    df['newspaper_uid'] = newspaper_uid

    return df

def _extract_host(df):
    logger.info('Extracting host from urls')
    df['host'] = df['url'].apply(lambda url: urlparse(url).netloc)

    return df

def _fill_missing_titles(df):
    logger.info('Filling missing titles.')
    missing_titles_mask = df['title'].isna()

    missing_titles = (df[missing_titles_mask]['url']
                      .str.extract(r'(?P<missing_titles>[^/]+)$')
                      .applymap(lambda title: title.split('-'))
                      .applymap(lambda title_word_list: ' '.join(title_word_list))
                      )

    df.loc[missing_titles_mask, 'title'] = missing_titles.loc[:, 'missing_titles']

    return df

def _generate_uids_for_rows(df):
    logger.info('Generating uids for each row')
    uids = (df
            .apply(lambda row: hashlib.md5(bytes(row['url'].encode())), axis=1)
            .apply(lambda hash_object: hash_object.hexdigest())
            )

    df['uid'] = uids

    return df.set_index('uid')

def _remove_new_lines_from_body(df):
    logger.info('Removing new lines from body.')

    stripped_body = (df
                     .apply(lambda row: row['body'], axis=1)
                     .apply(lambda body: list(body))
                     .apply(lambda letters: list(map(lambda letter: letter.replace('\n', ' '), letters)))
                     .apply(lambda letters: ''.join(letters)))

    df['body'] = stripped_body

    return df

def _tokenize_column(df, column_name):
    stop_words = set(stopwords.words('spanish'))

    column_names = []
    if type(column_name) is not list:
        column_names.append(column_name)
    else:
        column_names = column_name

    logger.info(
        f'Tokenizing column(s) {[column for column in column_names]}')

    for column in column_names:
        df['n_token_'+column] = (df
                                 .dropna()
                                 .apply(lambda row: nltk.word_tokenize(row[column]), axis=1)
                                 .apply(lambda tokens: list(filter(lambda token: token.isalpha(), tokens)))
                                 .apply(lambda tokens: list(map(lambda token: token.lower(), tokens)))
                                 .apply(lambda word_list: list(filter(lambda word: word not in stop_words, word_list)))
                                 .apply(lambda valid_word_list: len(valid_word_list)))

    return df

# noinspection PyShadowingNames
def main(filename):
    """ Funcion principal """
    logger.info('Starting cleaning process')

    df = _read_data(filename)
    print('Dataframe in firts step')
    print(df)

    newspaper_uid = _extract_newspaper_uid(filename)
    df = _add_newspaper_uid_column(df, newspaper_uid)
    df = _extract_host(df)
    df = _fill_missing_titles(df)
    df = _generate_uids_for_rows(df)
    df = _remove_new_lines_from_body(df)
    df = _tokenize_column(df, 'title')
    df = _tokenize_column(df, 'body')

    return df

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('filename', help='The path to the dirty data', type=str)

    arg = parser.parse_args()

    df = main(arg.filename)

    print(df)
    df.to_csv('cleaned_dataframe_el_universal.csv', sep='\t', encoding='utf-8')


from flask import Flask, render_template, request
from book_parser import BookParser
from book_request_handler import BookAPIRequestHandler
application = Flask(__name__, template_folder='templates', static_folder='static')

@application.route('/', methods=['GET'])
def open():
    return render_template('index.html', book_data=None)

@application.route('/', defaults={'term':None, 'page': 1}, methods=['POST'])
@application.route('/<term>/<page>', methods=['GET'])
def search(term, page):
    if not term:
        term = request.form['term']

    request_handler = BookAPIRequestHandler()
    response = request_handler.request_book_search(term, page)

    if ('items' in response):
        page = int(page)+1
        parser = BookParser()
        books = parser.get_book_data_from_json(response)
        return render_template('index.html', search_word=term, book_data=books, page_number=page)

    elif ('error' in response):
        return response['error']

    else:
        return render_template('index.html', book_data=None)


import numpy as np
import pylab as py
import matplotlib

rho = np.zeros((500,500))
gs = (6000.0*2.0/500.0)
v_e_1 = 950.0
v_e_2 = 300.0
v_out_2 = 600.0
vz_off_1 = 250.0
vz_off_2 = -1500.0

n1 = 'twoclump.dat'
n2 = 'OI_features/OI_268d.txt'

scale = 2.
fig = py.figure( figsize=(5.5*scale, 4.5*scale), facecolor='w', edgecolor='k' )
font = {'weight' : 'bold',
        'size'   : 10*scale}
matplotlib.rc('font', **font)  #set the font properties for entire plot
matplotlib.rc('xtick', labelsize=8*scale) #but have the tickmarks a little smaller

[x1,y1] = np.genfromtxt(n2,usecols=(0,1),unpack=True)
[x2,y2] = np.genfromtxt(n1,usecols=(0,1),unpack=True)
max1 = np.max(y1)
max2 = np.max(y2)
py.axis([6100, 6800,-.5,10.5])
py.plot(x1,y1*10/max1, 'k', lw=1, label='2011dh, 268 days')
py.plot(x2,y2*10/max2, 'k--', lw=2, label='Model')
leg = py.legend(loc=(.57,.86),fancybox=True,fontsize='small')
leg.get_frame().set_alpha(0.0)
py.xlabel(r'Wavelength ($\AA$)')
py.ylabel(r'Scaled F$_\lambda$')

for i in range(500):
    for j in range(500):
        z = gs*(1.0*i - 250)
        x = gs*(1.0*j - 250)+vz_off_1
        r = (x**2 + z**2.0)**0.5
        #FIRST SPHERE / TORUS
        sinsq = 0
        if (r != 0): sinsq = (1 - (z/r)**2.0)
        rho[i][j] = np.exp(-r/v_e_1) #*sinsq

        #SECOND SPHERE
        onoff = 1
        z = gs*(1.0*i - 250)
        x = gs*(1.0*j - 250)+vz_off_2
        r = (x**2 + z**2.0)**0.5
        if (r < v_out_2): rho[i][j] +=  4.0*np.exp(-r/v_e_2)*onoff

a = py.axes([.561, .5, .27, .27])
a.axis([138, 363, 150, 350])
py.setp(a, xticks=[], yticks=[])
a.imshow(np.log10(rho))#, cmap=py.cm.afmhot_r, py.cm.gray_r
a.scatter( 250,250, c='k', alpha=.5, marker='+', s=1000)

b = py.axes([.57, .23, .25, .25])
b.tick_params(labelsize='x-small')
b.set_xticks([6300+(40*l) for l in range(4)])
b.set_yticks([2*l for l in range(1,6)])
b.axis([6290, 6420, 1,10.5])
#b.set_xlabel(r'($\AA$)', size='x-small')
b.plot(x1,y1*10/max1,'k-', lw=1)
b.plot(x2,y2*10/max2,'k--', lw=2)

# add the eye
import matplotlib.image as image
im = image.imread('Eye_Symbol.png')[125:475,:,3]
c = py.axes([.38, .595, .3, .08], frameon=False)
c.axis([0,750,375,0])
py.setp(c, xticks=[], yticks=[])
c.imshow(im, cmap=py.cm.binary)
c.scatter( 450,176, marker=r'$\rightarrow$', s=500 )

py.savefig( 'twoclump.png', dpi=300 )


# -*- coding: utf-8 -*-

import numpy as np
import pandas as pd
import xlwings as xw
import os

wb = xw.Book('result.xlsm')
sht = wb.sheets('业绩数据')
start_date = sht.range('C1').value
end_date = sht.range('C2').value

dir_path = os.path.dirname(wb.fullname)
data = pd.read_excel(dir_path + '\\weekly_return.xlsx', header=0)

def get_month_rtn(df):
    sub_data = data[data['fund_name'] == df.fund_name.min()]
    nav1 = df[df['date'] == df['date'].min()]['last_nav'].max()
    if df['date'].min() > sub_data['date'].min():
        date1 = df['date'].min() - pd.Timedelta(days=7)
        while sub_data[sub_data['date'] == date1].size == 0:
            date1 = date1 - pd.Timedelta(days=7)
        nav1 = sub_data[sub_data['date'] == date1]['nav'].max()

    nav2 = df[df['date'] == df['date'].max()]['nav'].max()
    return nav2/nav1 - 1

def get_std_a(df):
    return np.std(df['weekly_return'])*np.sqrt(52)

def get_max_retra(df, start_date):
    nav_list = list(df['nav'])
    if (df['date'].min() - pd.Timedelta(days=7)) >= start_date:
        nav_list.insert(0, df['last_nav'][df.index[0]])

    peak_list = [max(nav_list[:i+1]) for i in range(len(nav_list))]
    return ((pd.Series(nav_list) / pd.Series(peak_list)) - 1).min()

def main_fun():
    sub_data = data[(data['date'] >= start_date) & (data['date'] < end_date)]
    # filter funds with less than 3 records in a month
    grouped = sub_data.groupby('fund_name')
    name_list = (grouped.apply(len) >= 3)[(grouped.apply(len) >= 3)].index
    sub_data = sub_data[sub_data['fund_name'].isin(name_list)]
    # calculate annual return, annual std, sharpe ratio, max retracement
    grouped = sub_data.groupby('fund_name')

    cate = grouped['type'].apply(min)
    rtn_m = grouped.apply(lambda x: get_month_rtn(x))
    rtn_a = rtn_m.apply(lambda x: (1+x)**(12) - 1)
    std_a = grouped.apply(get_std_a)
    sharpe = (rtn_a - 0.015)/std_a
    max_retra = grouped.apply(lambda x: get_max_retra(x, start_date))
    result = pd.DataFrame({'type':cate, 'annual_return':rtn_a, 'annual_std':std_a, 'sharpe':sharpe, 'max_retra':max_retra, 'month_return':rtn_m},
                      index = rtn_a.index)
    print(result)
    sht.range('A3').expand().clear_contents()
    sht.range('A3').value = result

main_fun()

code for work in CICC/week_to_all_time.py


from typing import Union
import warnings
from pathlib import Path

import click

from audiomidi import audio_utils, params

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

warnings.filterwarnings(action='ignore', category=UserWarning)

@click.command()
@click.option('-m', '--max-files', default=None, type=int)
def main(max_files: Union[int, None]):

    output_dir = params.features_dir
    output_dir.mkdir(exist_ok=True, parents=True)

    # process training files

    train_dir = params.nsynth_train_audio
    test_dir = params.nsynth_test_audio
    valid_dir = params.nsynth_valid_audio

    dirs = [train_dir, test_dir, valid_dir]
    dir_names = ['train', 'test', 'valid']

    for d, d_name in zip(dirs, dir_names):

        names, chroma_stfts, mfcc_stfts, _ = audio_utils.process_files(
            d,
            max_files=max_files,
            calc_chroma_stft=True,
            calc_mfcc_stft=True,
            calc_mfcc=False,
            label=d_name)

        names_file = audio_utils.dump_to_file(names, d_name + '_name',
                                              output_dir)

        if not Path(names_file[0]).exists():
            print('Error writing names')

        chroma_stft_file = audio_utils.dump_to_file(
            chroma_stfts, d_name + '_chroma_stft', output_dir)

        if not Path(chroma_stft_file[0]).exists():
            print(f'Error writing chroma_stft_{d_name}')

        mfcc_stft_file = audio_utils.dump_to_file(
            mfcc_stfts, d_name + '_mfcc_stft', output_dir)

        if not Path(mfcc_stft_file[0]).exists():
            print(f'Error writing mfcc_stft_{d_name}')

if __name__ == "__main__":
    main()


# uncompyle6 version 3.6.7
# Python bytecode 3.6 (3379)
# Decompiled from: Python 3.8.2 (tags/v3.8.2:7b3ab59, Feb 25 2020, 23:03:10) [MSC v.1916 64 bit (AMD64)]
# Embedded file name: build/bdist.macosx-10.7-x86_64/egg/airflow/contrib/operators/spark_sql_operator.py
# Compiled at: 2019-09-11 03:47:34
# Size of source mod 2**32: 4307 bytes
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.contrib.hooks.spark_sql_hook import SparkSqlHook

class SparkSqlOperator(BaseOperator):
    """SparkSqlOperator"""
    template_fields = [
     '_sql']
    template_ext = ['.sql', '.hql']

    @apply_defaults
    def __init__(self, sql, conf=None, conn_id='spark_sql_default', total_executor_cores=None, executor_cores=None, executor_memory=None, keytab=None, principal=None, master='yarn', name='default-name', num_executors=None, yarn_queue='default', *args, **kwargs):
        (super(SparkSqlOperator, self).__init__)(*args, **kwargs)
        self._sql = sql
        self._conf = conf
        self._conn_id = conn_id
        self._total_executor_cores = total_executor_cores
        self._executor_cores = executor_cores
        self._executor_memory = executor_memory
        self._keytab = keytab
        self._principal = principal
        self._master = master
        self._name = name
        self._num_executors = num_executors
        self._yarn_queue = yarn_queue
        self._hook = None

    def execute(self, context):
        """
        Call the SparkSqlHook to run the provided sql query
        """
        self._hook = SparkSqlHook(sql=(self._sql), conf=(self._conf),
          conn_id=(self._conn_id),
          total_executor_cores=(self._total_executor_cores),
          executor_cores=(self._executor_cores),
          executor_memory=(self._executor_memory),
          keytab=(self._keytab),
          principal=(self._principal),
          name=(self._name),
          num_executors=(self._num_executors),
          master=(self._master),
          yarn_queue=(self._yarn_queue))
        self._hook.run_query()

    def on_kill(self):
        self._hook.kill()


import argparse
import os
import time

import gym
import numpy as np
import rc_gym
import torch

from agents.ddpg_selfplay import DDPGActor
from agents.sac import GaussianPolicy

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--cuda", default=False,
                        action="store_true", help="Enable cuda")
    parser.add_argument("-c_atk", "--checkpoint_atk", required=True,
                        help="checkpoint to load attacker")
    parser.add_argument("-c_gk", "--checkpoint_gk", required=True,
                        help="checkpoint to load goalkeeper")
    args = parser.parse_args()
    device = "cuda" if args.cuda else "cpu"

    checkpoint_atk = torch.load(args.checkpoint_atk)
    checkpoint_gk = torch.load(args.checkpoint_gk)

    env = gym.make(checkpoint_atk['ENV_NAME'])

    if checkpoint_atk['AGENT'] == 'ddpg_async':
        pi_atk = DDPGActor(checkpoint_atk['N_OBS'], checkpoint_atk['N_ACTS']).to(device)
        pi_gk = DDPGActor(checkpoint_gk['N_OBS'], checkpoint_gk['N_ACTS']).to(device)
    elif checkpoint['AGENT'] == 'sac_async':
        pi_atk = GaussianPolicy(checkpoint_atk['N_OBS'], checkpoint_atk['N_ACTS'],
                            checkpoint_atk['LOG_SIG_MIN'],
                            checkpoint_atk['LOG_SIG_MAX'], checkpoint['EPSILON']).to(device)
        pi_gk = GaussianPolicy(checkpoint_gk['N_OBS'], checkpoint_gk['N_ACTS'],
                            checkpoint_gk['LOG_SIG_MIN'],
                            checkpoint_gk['LOG_SIG_MAX'], checkpoint_gk['EPSILON']).to(device)
    else:
        raise AssertionError

    pi_atk.load_state_dict(checkpoint_atk['pi_state_dict'])
    pi_gk.load_state_dict(checkpoint_gk['pi_state_dict'])
    pi_atk.eval()
    pi_gk.eval()

    while True:
        done = False
        s = env.reset()
        info = {}
        ep_steps = 0
        ep_rw_atk = 0
        ep_rw_gk = 0
        st_time = time.perf_counter()
        for i in range(checkpoint_atk['MAX_EPISODE_STEPS']):
            # Step the environment
            s_v_atk = torch.Tensor(s['observation_atk']).to(device)
            s_v_gk = torch.Tensor(s['observation_gk']).to(device)
            a_atk = pi_atk.get_action(s_v_atk)
            a_gk = pi_gk.get_action(s_v_gk)
            s_next, r, done, info = env.step({'action_atk': a_atk,'action_gk': a_gk})
            ep_steps += 1
            ep_rw_atk += r['reward_atk']
            ep_rw_gk += r['reward_gk']
            env.render()
            if done:
                break

            # Set state for next step
            s = s_next

        info['fps'] = ep_steps / (time.perf_counter() - st_time)
        info['ep_steps'] = ep_steps
        info['ep_rw_atk'] = ep_rw_atk
        info['ep_rw_gk'] = ep_rw_gk
        print(info)


import os
import shutil
import multiprocessing
from subprocess import CalledProcessError, check_output, STDOUT
from glob import glob
from concurrent.futures import ThreadPoolExecutor, wait
from utils import log_output, create_path, extract_tgz, filter_series, get_scanner_meta
from threading import Semaphore

LOG_MESSAGES = {
    'success_converted':
        'Converted {} to {}\n'
        'Command:\n{}\n'
        'Return Code:\n{}\n\n',
    'output':
        'Output:\n{}\n\n',
    'dcm2niix_error':
        'Error running dcm2niix on DICOM series in {} directory.\n'
        'Command:\n{}\n'
        'Return Code:\n{}\n\n',
    'dimon_error':
        'Error running Dimon on DICOM series in {} directory.\n'
        'Command:\n{}\n'
        'Return Code:\n{}\n\n',
}

MAX_WORKERS = multiprocessing.cpu_count() * 5

class NiftyConversionFailure(Exception):
    def __init__(self, message):
        self.message = message

class DuplicateFile(Exception):
    def __init__(self, message):
        self.message = message

def dcm_to_nifti(dcm_dir, out_fname, out_dir, conversion_tool, logger=None, bids_meta=False, semaphore=None):

    if conversion_tool == 'dcm2niix':

        dcm2niix_workdir = dcm_dir

        if bids_meta:
            cmd = [
                "dcm2niix",
                "-z",
                "y",
                "-b",
                "y",
                "-f",
                out_fname,
                dcm_dir
            ]
        else:
            cmd = [
                "dcm2niix",
                "-z",
                "y",
                "-f",
                out_fname,
                dcm_dir
            ]

        try:

            result = check_output(cmd, stderr=STDOUT, cwd=dcm2niix_workdir, universal_newlines=True)

            # The following line is a hack to get the actual filename returned by the dcm2niix utility. When converting
            # the B0 dcm files, or files that specify which coil they used, or whether they contain phase information,
            # the utility appends some prefixes to the filename it saves, instead of just using
            # the specified output filename. There is no option to turn this off (and the author seemed unwilling to
            # add one). With this hack I retrieve the actual filename it used to save the file from the utility output.
            # This might break on future updates of dcm2niix
            actual_fname = \
                [s for s in ([s for s in str(result).split('\n') if "Convert" in s][0].split(" "))
                 if s[0] == '/'][0].split("/")[-1]

            # Move nifti file and json bids file to anat folder
            shutil.move(os.path.join(dcm_dir, "{}.nii.gz".format(actual_fname)),
                        os.path.join(out_dir, "{}.nii.gz".format(out_fname)))
            shutil.move(os.path.join(dcm_dir, "{}.json".format(actual_fname)),
                        os.path.join(out_dir, "{}.json".format(out_fname)))

            dcm_file = [f for f in os.listdir(dcm_dir) if ".dcm" in f][0]

            log_str = LOG_MESSAGES['success_converted'].format(os.path.join(dcm_dir, dcm_file), out_fname,
                                                               " ".join(cmd), 0)

            if result:
                log_str += LOG_MESSAGES['output'].format(result)

            log_output(log_str, logger=logger, semaphore=semaphore)

            return ("/".join(dcm2niix_workdir.split("/")[-3:]),
                    os.path.join("/".join(out_dir.split("/")[-4:]), out_fname + ".nii.gz"),
                    True)

        except CalledProcessError as e:

            log_str = LOG_MESSAGES['dcm2niix_error'].format(dcm_dir, " ".join(cmd), e.returncode)

            if e.output:
                log_str += LOG_MESSAGES['output'].format(e.output)

            log_output(log_str, level="ERROR", logger=logger, semaphore=semaphore)

            return ("/".join(dcm2niix_workdir.split("/")[-3:]),
                    os.path.join("/".join(out_dir.split("/")[-4:]), out_fname + ".nii.gz"),
                    False)

        finally:

            # Clean up temporary files
            tmp_files = glob(os.path.join(dcm2niix_workdir, "*.nii.gz"))
            tmp_files.extend(glob(os.path.join(dcm2niix_workdir, "*.json")))

            if tmp_files:
                list(map(os.remove, tmp_files))

    elif conversion_tool == 'dimon':

        dimon_workdir = dcm_dir

        # IMPLEMENT GENERATION OF BIDS METADATA FILES WHEN USING DIMON FOR CONVERSION OF DCM FILES

        cmd = [
            "Dimon",
            "-infile_pattern",
            os.path.join(dcm_dir, "*.dcm"),
            "-gert_create_dataset",
            "-gert_quit_on_err",
            "-gert_to3d_prefix",
            "{}.nii.gz".format(out_fname)
        ]

        dimon_env = os.environ.copy()
        dimon_env['AFNI_TO3D_OUTLIERS'] = 'No'

        try:

            result = check_output(cmd, stderr=STDOUT, env=dimon_env, cwd=dimon_workdir, universal_newlines=True)

            # Check the contents of stdout for the -quit_on_err flag because to3d returns a success code
            # even if it terminates because the -quit_on_err flag was thrown
            if "to3d kept from going into interactive mode by option -quit_on_err" in result:

                log_str = LOG_MESSAGES['dimon_error'].format(dcm_dir, " ".join(cmd), 0)

                if result:
                    log_str += LOG_MESSAGES['output'].format(result)

                log_output(log_str, level="ERROR", logger=logger, semaphore=semaphore)

                return ("/".join(dimon_workdir.split("/")[-3:]),
                        os.path.join("/".join(out_dir.split("/")[-4:]), out_fname + ".nii.gz"),
                        False)

            shutil.move(os.path.join(dimon_workdir, "{}.nii.gz".format(out_fname)),
                        os.path.join(out_dir, "{}.nii.gz".format(out_fname)))

            dcm_file = [f for f in os.listdir(dcm_dir) if ".dcm" in f][0]

            log_str = LOG_MESSAGES['success_converted'].format(os.path.join(dcm_dir, dcm_file), out_fname,
                                                               " ".join(cmd), 0)

            if result:
                log_str += LOG_MESSAGES['output'].format(result)

            log_output(log_str, logger=logger, semaphore=semaphore)

            return ("/".join(dimon_workdir.split("/")[-3:]),
                    os.path.join("/".join(out_dir.split("/")[-4:]), out_fname + ".nii.gz"),
                    True)

        except CalledProcessError as e:

            log_str = LOG_MESSAGES['dimon_error'].format(dcm_dir, " ".join(cmd), e.returncode)

            if e.output:
                log_str += LOG_MESSAGES['output'].format(e.output)

            log_output(log_str, level="ERROR", logger=logger, semaphore=semaphore)

            return ("/".join(dimon_workdir.split("/")[-3:]),
                    os.path.join("/".join(out_dir.split("/")[-4:]), out_fname + ".nii.gz"),
                    False)

        finally:

            # Clean up temporary files
            tmp_files = glob(os.path.join(dimon_workdir, "GERT_Reco_dicom*"))
            tmp_files.extend(glob(os.path.join(dimon_workdir, "dimon.files.run.*")))

            if tmp_files:
                list(map(os.remove, tmp_files))

    else:

        raise NiftyConversionFailure("Tool Error: {} is not a supported conversion tool. Please select 'dcm2niix' or "
                                     "'dimon'".format(conversion_tool))

def convert_to_bids(bids_dir, oxygen_dir, mapping_guide=None, conversion_tool='dcm2niix', logger=None,
                    nthreads=MAX_WORKERS, overwrite=False, filters=None, scanner_meta=False):

    if nthreads > 0:
        thread_semaphore = Semaphore(value=1)
    else:
        thread_semaphore = None

    # If BIDS directory exists, verify that it's either empty, or that overwrite is allowed. Otherwise create directory.
    if os.path.isdir(bids_dir):

        bids_files = glob(os.path.join(bids_dir, '*'))

        if bids_files:
            if not overwrite:
                raise DuplicateFile("The BIDS directory is not empty, and overwrite is set to False. Aborting...")
            else:
                rm_files = glob(os.path.join(bids_dir, '*'))
                list(map(shutil.rmtree, rm_files))
    else:
        create_path(bids_dir)

    # Uncompress any compressed Oxygen DICOM files
    raw_files = os.path.join(oxygen_dir, '*')

    # Check if there are compressed oxygen files, and if so, uncompress them
    compressed_files = [d for d in glob(raw_files) if os.path.isfile(d)]

    log_output("Extracting compressed files...", logger=logger)

    if nthreads > 0:   # Run in multiple threads

        futures = []

        with ThreadPoolExecutor(max_workers=nthreads) as executor:

            for f in compressed_files:

                futures.append(executor.submit(extract_tgz, f, oxygen_dir, logger, thread_semaphore))

        wait(futures)

    else:   # Run sequentially
        for f in compressed_files:
            extract_tgz(f, oxygen_dir, logger=logger)

    log_output("Compressed file extractions complete.", logger=logger)

    # Now we can get a list of uncompressed directories
    uncompressed_files = [d for d in glob(raw_files) if os.path.isdir(d)]

    mapping = {}

    # If a BIDS mapping has not be provided to guide the conversion process, attempt to generate mapping from
    # available information.
    if not mapping_guide:

        subject_counter = 1

        mapping = {}

        for unc_file in uncompressed_files:

            subject_id = unc_file.split("/")[-1].split("-")[-1]

            if subject_id not in mapping.keys():

                mapping[subject_id] = {
                    "bids_subject": "{:0>4d}".format(subject_counter),
                    "sessions": {}
                }

                subject_counter += 1

            session_dirs = [d for d in glob(os.path.join(unc_file, '*')) if os.path.isdir(d)]

            session_counter = 1

            for ses_dir in session_dirs:

                session_id = ses_dir.split("/")[-1]

                mapping[subject_id]["sessions"][session_id] = {
                    "bids_session": "{:0>4d}".format(session_counter),
                    "oxygen_file": "{}-{}-DICOM.tgz".format(ses_dir.split("/")[-2], ses_dir.split("/")[-1]),
                    "scans": {}
                }

                session_counter += 1

                scan_dirs = [d for d in glob(os.path.join(ses_dir, '*')) if os.path.isdir(d) and "mr_" in d]

                scan_counter = 1

                for sc_dir in scan_dirs:

                    scan_id = sc_dir.split("/")[-1]

                    # Filter this series directory
                    if filter_series(sc_dir, filters=filters, logger=logger):
                        continue

                    mapping[subject_id]["sessions"][session_id]["scans"][scan_id] = {
                        "series_dir": "/".join(sc_dir.split("/")[-3:]),
                        "bids_fpath": "",
                        "conversion_status": False,
                        "meta": {
                            "type": "func",
                            "modality": "bold",
                            "description": "task-fmri",
                            "run": "{:0>4d}".format(scan_counter)
                        }
                    }

                    if scanner_meta:
                        meta = get_scanner_meta(sc_dir)
                        mapping[subject_id]["sessions"][session_id]["scans"][scan_id]["scanner_meta"] = meta

                    scan_counter += 1

    # Mapping has been generated
    # Iterate through the mapping to create execution list to be split into threads
    exec_list = []

    for subject in mapping.keys():
        for session in mapping[subject]["sessions"].keys():
            for scan in mapping[subject]["sessions"][session]["scans"].keys():

                series_dir = os.path.join(oxygen_dir,
                                          mapping[subject]["sessions"][session]["scans"][scan]["series_dir"])
                bids_subject = "sub-{}".format(mapping[subject]["bids_subject"])
                bids_session = "ses-{}".format(mapping[subject]["sessions"][session]["bids_session"])
                bids_desc = mapping[subject]["sessions"][session]["scans"][scan]["meta"]["description"]
                bids_type = mapping[subject]["sessions"][session]["scans"][scan]["meta"]["type"]
                bids_modality = mapping[subject]["sessions"][session]["scans"][scan]["meta"]["modality"]
                bids_run = "run-{}".format(mapping[subject]["sessions"][session]["scans"][scan]["meta"]["run"])
                bids_fname = "{}_{}_{}_{}".format(bids_subject, bids_session, bids_desc, bids_run)

                if bids_modality:
                    bids_fname += "_{}".format(bids_modality)

                bids_fname = "{}.nii.gz".format(bids_fname)

                bids_fpath = os.path.join(bids_dir, bids_subject, bids_session, bids_type, bids_fname)

                exec_list.append((series_dir, bids_fpath))

    # Iterate through executable list and convert to nifti
    if nthreads > 0:    # Run in multiple threads

        futures = []

        with ThreadPoolExecutor(max_workers=nthreads) as executor:

            for dcm_dir, bids_fpath in exec_list:

                out_bdir = "/".join(bids_fpath.split("/")[:-1])
                if not os.path.isdir(out_bdir):
                    create_path(out_bdir)

                out_fname = bids_fpath.split("/")[-1].split(".")[0]

                futures.append(executor.submit(dcm_to_nifti, dcm_dir, out_fname, out_bdir,
                                               conversion_tool=conversion_tool, bids_meta=True, logger=logger,
                                               semaphore=thread_semaphore))
                ## FOR TESTING
                # break
                #######

            wait(futures)

            for future in futures:
                series_dir, bids_fpath, success = future.result()

                subject = series_dir.split("/")[0].split("-")[1]
                session = series_dir.split("/")[1]
                scan = series_dir.split("/")[2]

                if success:
                    mapping[subject]["sessions"][session]["scans"][scan]["bids_fpath"] = bids_fpath
                    mapping[subject]["sessions"][session]["scans"][scan]["conversion_status"] = True

    else:   # Run sequentially

        for dcm_dir, bids_fpath in exec_list:

            out_bdir = "/".join(bids_fpath.split("/")[:-1])
            if not os.path.isdir(out_bdir):
                create_path(out_bdir)

            out_fname = bids_fpath.split("/")[-1].split(".")[0]

            series_dir, bids_fpath, success = dcm_to_nifti(dcm_dir, out_fname, out_bdir, conversion_tool='dcm2niix',
                                                           bids_meta=True, logger=logger)

            subject = series_dir.split("/")[0].split("-")[1]
            session = series_dir.split("/")[1]
            scan = series_dir.split("/")[2]

            if success:
                mapping[subject]["sessions"][session]["scans"][scan]["bids_fpath"] = bids_fpath
                mapping[subject]["sessions"][session]["scans"][scan]["conversion_status"] = True

    return mapping