# -*- coding: utf-8 -*-
"""k-Nearest Neighbors Detector (kNN)
"""
# Author: Yue Zhao <zhaoy@cmu.edu>
# License: BSD 2 clause

import numpy as np
import torch

from .base import BaseDetector
from .intermediate_layers import knn_batch


class KNN(BaseDetector):
    # noinspection PyPep8
    """kNN class for outlier detection.
    For an observation, its distance to its kth nearest neighbor could be
    viewed as the outlying score. It could be viewed as a way to measure
    the density. See :cite:`ramaswamy2000efficient,angiulli2002fast` for
    details.

    Three kNN detectors are supported:
    largest: use the distance to the kth neighbor as the outlier score
    mean: use the average of all k neighbors as the outlier score
    median: use the median of the distance to k neighbors as the outlier score

   Parameters
   ----------
   contamination : float in (0., 0.5), optional (default=0.1)
       The amount of contamination of the data set,
       i.e. the proportion of outliers in the data set. Used when fitting to
       define the threshold on the decision function.

   n_neighbors : int, optional (default=20)
       Number of neighbors to use by default for `kneighbors` queries.
       If n_neighbors is larger than the number of samples provided,
       all samples will be used.

   batch_size : integer, optional (default = None)
       Number of samples to process per batch.

   device : str, optional (default = 'cpu')
       Valid device id, e.g., 'cuda:0' or 'cpu'

   Attributes
   ----------
   decision_scores_ : numpy array of shape (n_samples,)
       The outlier scores of the training data.
       The higher, the more abnormal. Outliers tend to have higher
       scores. This value is available once the detector is
       fitted.

   threshold_ : float
       The threshold is based on ``contamination``. It is the
       ``n_samples * contamination`` most abnormal samples in
       ``decision_scores_``. The threshold is calculated for generating
       binary outlier labels.

   labels_ : int, either 0 or 1
       The binary labels of the training data. 0 stands for inliers
       and 1 for outliers/anomalies. It is generated by applying
       ``threshold_`` on ``decision_scores_``.
   """

    def __init__(self, contamination=0.1, n_neighbors=5, batch_size=None,
                 device='cuda:0'):
        super(KNN, self).__init__(contamination=contamination)
        self.n_neighbors = n_neighbors
        self.batch_size = batch_size
        self.device = device

    def fit(self, X, y=None, return_time=False):
        """Fit detector. y is ignored in unsupervised methods.

        Parameters
        ----------
        X : numpy array of shape (n_samples, n_features)
            The input samples.

        y : Ignored
            Not used, present for API consistency by convention.

        return_time : boolean (default=True)
            If True, set self.gpu_time to the measured GPU time.

        Returns
        -------
        self : object
            Fitted estimator.
        """
        # todo: add one for pytorch tensor
        # X = check_array(X)
        self._set_n_classes(y)

        if self.device != 'cpu' and return_time:
            start = torch.cuda.Event(enable_timing=True)
            end = torch.cuda.Event(enable_timing=True)
            start.record()

        knn_dist, _ = knn_batch(X, X, self.n_neighbors + 1,
                                batch_size=self.batch_size,
                                device=self.device)

        if self.device != 'cpu' and return_time:
            end.record()
            torch.cuda.synchronize()

        self.decision_scores_ = knn_dist[:, -1].cpu().numpy()
        self._process_decision_scores()

        # return GPU time in seconds
        if return_time:
            self.gpu_time = start.elapsed_time(end) / 1000

        return self

    def decision_function(self, X):
        """Predict raw anomaly score of X using the fitted detector.
         For consistency, outliers are assigned with larger anomaly scores.
        Parameters
        ----------
        X : numpy array of shape (n_samples, n_features)
            The training input samples. Sparse matrices are accepted only
            if they are supported by the base estimator.
        Returns
        -------
        anomaly_scores : numpy array of shape (n_samples,)
            The anomaly score of the input samples.
        """
        # use multi-thread execution
        if hasattr(self, 'X_train'):
            original_size = X.shape[0]
            X = np.concatenate((self.X_train, X), axis=0)

        # return decision_scores_.ravel()
