#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

"""
Python package for feature in MLlib.
"""
from __future__ import absolute_import

import sys
import warnings
import random
import binascii
if sys.version >= '3':
    basestring = str
    unicode = str

from py4j.protocol import Py4JJavaError

from pyspark import SparkContext
from pyspark.rdd import RDD, ignore_unicode_prefix
from pyspark.mllib.common import callMLlibFunc, JavaModelWrapper
from pyspark.mllib.linalg import (
    Vector, Vectors, DenseVector, SparseVector, _convert_to_vector)
from pyspark.mllib.regression import LabeledPoint
from pyspark.mllib.util import JavaLoader, JavaSaveable

__all__ = ['Normalizer', 'StandardScalerModel', 'StandardScaler',
           'HashingTF', 'IDFModel', 'IDF', 'Word2Vec', 'Word2VecModel',
           'ChiSqSelector', 'ChiSqSelectorModel', 'ElementwiseProduct']


class VectorTransformer(object):
    """
    .. note:: DeveloperApi

    Base class for transformation of a vector or RDD of vector
    """
    def transform(self, vector):
        """
        Applies transformation on a vector.

        :param vector: vector to be transformed.
        """
        raise NotImplementedError


class Normalizer(VectorTransformer):
    """
    .. note:: Experimental

    Normalizes samples individually to unit L\ :sup:`p`\  norm

    For any 1 <= `p` < float('inf'), normalizes samples using
    sum(abs(vector) :sup:`p`) :sup:`(1/p)` as norm.

    For `p` = float('inf'), max(abs(vector)) will be used as norm for
    normalization.

    :param p: Normalization in L^p^ space, p = 2 by default.

    >>> v = Vectors.dense(range(3))
    >>> nor = Normalizer(1)
    >>> nor.transform(v)
    DenseVector([0.0, 0.3333, 0.6667])

    >>> rdd = sc.parallelize([v])
    >>> nor.transform(rdd).collect()
    [DenseVector([0.0, 0.3333, 0.6667])]

    >>> nor2 = Normalizer(float("inf"))
    >>> nor2.transform(v)
    DenseVector([0.0, 0.5, 1.0])
    """
    def __init__(self, p=2.0):
        assert p >= 1.0, "p should be greater than 1.0"
        self.p = float(p)

    def transform(self, vector):
        """
        Applies unit length normalization on a vector.

        :param vector: vector or RDD of vector to be normalized.
        :return: normalized vector. If the norm of the input is zero, it
                 will return the input vector.
        """
        sc = SparkContext._active_spark_context
        assert sc is not None, "SparkContext should be initialized first"
        if isinstance(vector, RDD):
            vector = vector.map(_convert_to_vector)
        else:
            vector = _convert_to_vector(vector)
        return callMLlibFunc("normalizeVector", self.p, vector)


class JavaVectorTransformer(JavaModelWrapper, VectorTransformer):
    """
    Wrapper for the model in JVM
    """

    def transform(self, vector):
        """
        Applies transformation on a vector or an RDD[Vector].

        Note: In Python, transform cannot currently be used within
              an RDD transformation or action.
              Call transform directly on the RDD instead.

        :param vector: Vector or RDD of Vector to be transformed.
        """
        if isinstance(vector, RDD):
            vector = vector.map(_convert_to_vector)
        else:
            vector = _convert_to_vector(vector)
        return self.call("transform", vector)


class StandardScalerModel(JavaVectorTransformer):
    """
    .. note:: Experimental

    Represents a StandardScaler model that can transform vectors.
    """
    def transform(self, vector):
        """
        Applies standardization transformation on a vector.

        Note: In Python, transform cannot currently be used within
              an RDD transformation or action.
              Call transform directly on the RDD instead.

        :param vector: Vector or RDD of Vector to be standardized.
        :return: Standardized vector. If the variance of a column is
                 zero, it will return default `0.0` for the column with
                 zero variance.
        """
        return JavaVectorTransformer.transform(self, vector)

    def setWithMean(self, withMean):
        """
        Setter of the boolean which decides
        whether it uses mean or not
        """
        self.call("setWithMean", withMean)
        return self

    def setWithStd(self, withStd):
        """
        Setter of the boolean which decides
        whether it uses std or not
        """
        self.call("setWithStd", withStd)
        return self


class StandardScaler(object):
    """
    .. note:: Experimental

    Standardizes features by removing the mean and scaling to unit
    variance using column summary statistics on the samples in the
    training set.

    :param withMean: False by default. Centers the data with mean
                     before scaling. It will build a dense output, so this
                     does not work on sparse input and will raise an
                     exception.
    :param withStd: True by default. Scales the data to unit
                    standard deviation.

    >>> vs = [Vectors.dense([-2.0, 2.3, 0]), Vectors.dense([3.8, 0.0, 1.9])]
    >>> dataset = sc.parallelize(vs)
    >>> standardizer = StandardScaler(True, True)
    >>> model = standardizer.fit(dataset)
    >>> result = model.transform(dataset)
    >>> for r in result.collect(): r
    DenseVector([-0.7071, 0.7071, -0.7071])
    DenseVector([0.7071, -0.7071, 0.7071])
    """
    def __init__(self, withMean=False, withStd=True):
        if not (withMean or withStd):
            warnings.warn("Both withMean and withStd are false. The model does nothing.")
        self.withMean = withMean
        self.withStd = withStd

    def fit(self, dataset):
        """
        Computes the mean and variance and stores as a model to be used
        for later scaling.

        :param dataset: The data used to compute the mean and variance
                     to build the transformation model.
        :return: a StandardScalarModel
        """
        dataset = dataset.map(_convert_to_vector)
        jmodel = callMLlibFunc("fitStandardScaler", self.withMean, self.withStd, dataset)
        return StandardScalerModel(jmodel)


class ChiSqSelectorModel(JavaVectorTransformer):
    """
    .. note:: Experimental

    Represents a Chi Squared selector model.
    """
    def transform(self, vector):
        """
        Applies transformation on a vector.

        :param vector: Vector or RDD of Vector to be transformed.
        :return: transformed vector.
        """
        return JavaVectorTransformer.transform(self, vector)


class ChiSqSelector(object):
    """
    .. note:: Experimental

    Creates a ChiSquared feature selector.

    :param numTopFeatures: number of features that selector will select.

    >>> data = [
    ...     LabeledPoint(0.0, SparseVector(3, {0: 8.0, 1: 7.0})),
    ...     LabeledPoint(1.0, SparseVector(3, {1: 9.0, 2: 6.0})),
    ...     LabeledPoint(1.0, [0.0, 9.0, 8.0]),
    ...     LabeledPoint(2.0, [8.0, 9.0, 5.0])
    ... ]
    >>> model = ChiSqSelector(1).fit(sc.parallelize(data))
    >>> model.transform(SparseVector(3, {1: 9.0, 2: 6.0}))
    SparseVector(1, {0: 6.0})
    >>> model.transform(DenseVector([8.0, 9.0, 5.0]))
    DenseVector([5.0])
    """
    def __init__(self, numTopFeatures):
        self.numTopFeatures = int(numTopFeatures)

    def fit(self, data):
        """
        Returns a ChiSquared feature selector.

        :param data: an `RDD[LabeledPoint]` containing the labeled dataset
                     with categorical features. Real-valued features will be
                     treated as categorical for each distinct value.
                     Apply feature discretizer before using this function.
        """
        jmodel = callMLlibFunc("fitChiSqSelector", self.numTopFeatures, data)
        return ChiSqSelectorModel(jmodel)


class PCAModel(JavaVectorTransformer):
    """
    Model fitted by [[PCA]] that can project vectors to a low-dimensional space using PCA.
    """


class PCA(object):
    """
    A feature transformer that projects vectors to a low-dimensional space using PCA.

    >>> data = [Vectors.sparse(5, [(1, 1.0), (3, 7.0)]),
    ...     Vectors.dense([2.0, 0.0, 3.0, 4.0, 5.0]),
    ...     Vectors.dense([4.0, 0.0, 0.0, 6.0, 7.0])]
    >>> model = PCA(2).fit(sc.parallelize(data))
    >>> pcArray = model.transform(Vectors.sparse(5, [(1, 1.0), (3, 7.0)])).toArray()
    >>> pcArray[0]
    1.648...
    >>> pcArray[1]
    -4.013...
    """
    def __init__(self, k):
        """
        :param k: number of principal components.
        """
        self.k = int(k)

    def fit(self, data):
        """
        Computes a [[PCAModel]] that contains the principal components of the input vectors.
        :param data: source vectors
        """
        jmodel = callMLlibFunc("fitPCA", self.k, data)
        return PCAModel(jmodel)


class HashingTF(object):
    """
    .. note:: Experimental

    Maps a sequence of terms to their term frequencies using the hashing
    trick.

    Note: the terms must be hashable (can not be dict/set/list...).

    :param numFeatures: number of features (default: 2^20)

    >>> htf = HashingTF(100)
    >>> doc = "a a b b c d".split(" ")