# -*- coding: utf-8 -*-
"""Local Outlier Factor (LOF). Implemented on scikit-learn library.
"""
# 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
from ..utils.utility import get_batch_index


def get_cosine_similarity(input1, input2):
    # todo: add support of GPU
    '''

    Parameters
    ----------
    input1
    input2

    Returns
    -------

    '''
    return torch.sum(input1 * input2, dim=1) / (
            torch.linalg.norm(input1, dim=1) ** 2 *
            torch.linalg.norm(input2, dim=1) ** 2)


class ABOD(BaseDetector):
    """ABOD class for Angle-base Outlier Detection.
    For an observation, the variance of its weighted cosine scores to all
    neighbors could be viewed as the outlying score.
    See :cite:`kriegel2008angle` for details.

    Two version of ABOD are supported:

    - Fast ABOD: use k nearest neighbors to approximate.
    - Original ABOD: consider all training points with high time complexity at
      O(n^3).


    Parameters
    ----------
    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(ABOD, 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.batch_size is None:
            self.decision_scores_ = self._fit_full(X, return_time)
        else:
            self.decision_scores_ = self._fit_batch(X, return_time)

        self._process_decision_scores()

        return self

    def _fit_full(self, X, return_time):
        n_samples, n_features = X.shape[0], X.shape[1]

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

        # first to identity the k nearst neighbors of each sample
        knn_dist, knn_inds = knn_batch(X, X, self.n_neighbors + 1,
                                       batch_size=self.batch_size,
                                       device=self.device)

        knn_dist, knn_inds = knn_dist[:, 1:], knn_inds[:, 1:]

        # build index list
        # neighbor 1 and neighbor 2
        n_combs = int(self.n_neighbors * (self.n_neighbors - 1) / 2)
        neighbor_indexs = torch.zeros([n_samples * n_combs, 2])

        # create self index, each repeat n_combs times
        self_indexs = torch.arange(0, n_samples).reshape(n_samples, 1).repeat(
            1, n_combs).view(-1)

        # all samples' possible NN combinations are stored
        # it is fine since the loop is only for building the index
        for idx in range(n_samples):
            # print(torch.combinations(knn_inds[0, :], 2))
            neighbor_indexs[idx * n_combs:(idx + 1) * n_combs,
            :] = torch.combinations(knn_inds[idx, :], 2)

        if self.device != 'cpu' and return_time:
            end.record()
            torch.cuda.synchronize()
            # return GPU time in seconds
            self.gpu_time = start.elapsed_time(end) / 1000

        # select the data
        self_feature = torch.index_select(X, 0, self_indexs.long())

        # both need to subtruct the original feature
        nn_1 = torch.index_select(X, 0,
                                  neighbor_indexs[:, 0].long()) - self_feature
        nn_2 = torch.index_select(X, 0,
                                  neighbor_indexs[:, 1].long()) - self_feature

        # calculate cosine similarity
        cos_sims = get_cosine_similarity(nn_1, nn_2).view(-1, n_combs)

        abod_scores = torch.nan_to_num(torch.var(cos_sims, dim=1)).numpy() * -1
        return abod_scores

    def _fit_batch(self, X, return_time):
        n_samples, n_features = X.shape[0], X.shape[1]

        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        # first to identity the k nearst neighbors of each sample
        knn_dist, knn_inds = knn_batch(X, X, self.n_neighbors + 1,
                                       batch_size=self.batch_size,
                                       device=self.device)

        knn_dist, knn_inds = knn_dist[:, 1:], knn_inds[:, 1:]

        # build index list
        # neighbor 1 and neighbor 2
        n_combs = int(self.n_neighbors * (self.n_neighbors - 1) / 2)
        neighbor_indexs = torch.zeros([n_samples * n_combs, 2])

        # create self index, each repeat n_combs times
        self_indexs = torch.arange(0, n_samples).reshape(n_samples, 1).repeat(
            1, n_combs).view(-1)

        # all samples' possible NN combinations are stored
        # it is fine since the loop is only for building the index
        for idx in range(n_samples):
            # print(torch.combinations(knn_inds[0, :], 2))
            neighbor_indexs[idx * n_combs:(idx + 1) * n_combs,
            :] = torch.combinations(knn_inds[idx, :], 2)

        cos_sims = torch.zeros([n_samples * n_combs])
        # recall we need to calculate
        batch_index = get_batch_index(n_samples * n_combs, self.batch_size)

        for i, index in enumerate(batch_index):
            # select the data
            self_feature = torch.index_select(X, 0, self_indexs[
                                                    index[0]:index[1]].long())

            # both need to subtruct the original feature
            nn_1 = torch.index_select(X, 0, neighbor_indexs[:, 0][
                                            index[0]:index[
                                                1]].long()) - self_feature
            nn_2 = torch.index_select(X, 0, neighbor_indexs[:, 1][
                                            index[0]:index[
                                                1]].long()) - self_feature

            # calculate cosine similarity
            cos_sims[index[0]:index[1]] = get_cosine_similarity(nn_1, nn_2)

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

        cos_sims = cos_sims.view(-1, n_combs)

        abod_scores = torch.nan_to_num(torch.var(cos_sims, dim=1)).numpy() * -1
        return abod_scores

    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()
