# Path: splearn/utils/validation.py
# def check_rdd(rdd, expected_dtype):
#     """Wrapper function to check_rdd_dtype. Raises TypeError in case of dtype
#     mismatch.
# 
#     Parameters:
#     -----------
#     rdd: splearn.BlockRDD
#         The RDD to check
#     expected_dtype: {type, list of types, tuple of types, dict of types}
#         Expected type(s). If the RDD is a DictRDD the parameter type is
#         restricted to dict.
#     """
#     if not check_rdd_dtype(rdd, expected_dtype):
#         raise TypeError("{0} dtype mismatch.".format(rdd))
# 
# Path: splearn/linear_model/base.py
# class SparkLinearModelMixin(object):
# 
#     def __add__(self, other):
#         """Add method for Linear models with coef and intercept attributes.
# 
#         Parameters
#         ----------
#         other : fitted sklearn linear model
#             Model to add.
# 
#         Returns
#         -------
#         model : Linear model
#             Model with updated coefficients.
#         """
#         model = copy.deepcopy(self)
#         model.coef_ += other.coef_
#         model.intercept_ += other.intercept_
#         return model
# 
#     def __radd__(self, other):
#         """Reverse add method for Linear models.
# 
#         Parameters
#         ----------
#         other : fitted sklearn linear model
#             Model to add.
# 
#         Returns
#         -------
#         model : Linear model
#             Model with updated coefficients.
#         """
#         return self if other == 0 else self.__add__(other)
# 
#     def __div__(self, other):
#         """Division method for Linear models. Used for averaging.
# 
#         Parameters
#         ----------
#         other : integer
#             Integer to divide with.
# 
#         Returns
#         -------
#         model : Linear model
#             Model with updated coefficients.
#         """
#         self.coef_ /= other
#         self.intercept_ /= other
#         return self
# 
#     __truediv__ = __div__
# 
#     def _spark_fit(self, cls, Z, *args, **kwargs):
#         """Wraps a Scikit-learn Linear model's fit method to use with RDD
#         input.
# 
#         Parameters
#         ----------
#         cls : class object
#             The sklearn linear model's class to wrap.
#         Z : TupleRDD or DictRDD
#             The distributed train data in a DictRDD.
# 
#         Returns
#         -------
#         self: the wrapped class
#         """
#         mapper = lambda X_y: super(cls, self).fit(
#             X_y[0], X_y[1], *args, **kwargs
#         )
#         models = Z.map(mapper)
#         avg = models.reduce(operator.add) / models.count()
#         self.__dict__.update(avg.__dict__)
#         return self
# 
#     def _spark_predict(self, cls, X, *args, **kwargs):
#         """Wraps a Scikit-learn Linear model's predict method to use with RDD
#         input.
# 
#         Parameters
#         ----------
#         cls : class object
#             The sklearn linear model's class to wrap.
#         Z : ArrayRDD
#             The distributed data to predict in a DictRDD.
# 
#         Returns
#         -------
#         self: the wrapped class
#         """
#         return X.map(lambda X: super(cls, self).predict(X, *args, **kwargs))
# 
#     def _to_scikit(self, cls):
#         new = cls()
#         new.__dict__ = self.__dict__
#         return new


import numpy as np
import scipy.sparse as sp
from sklearn.linear_model import SGDClassifier
from ..utils.validation import check_rdd
from .base import SparkLinearModelMixin
        Parameters
        ----------
        Z : DictRDD containing (X, y) pairs
            X - Training vector
            y - Target labels
        classes : iterable
            The set of available classes

        Returns
        -------
        self : object
            Returns self.
        """
        check_rdd(Z, {'X': (sp.spmatrix, np.ndarray)})
        self._classes_ = np.unique(classes)
        return self._spark_fit(SparkSGDClassifier, Z)

    def predict(self, X):
        """Distributed method to predict class labels for samples in X.

        Parameters
        ----------
        X : ArrayRDD containing {array-like, sparse matrix}
            Samples.

        Returns
        -------
        C : ArrayRDD
            Predicted class label per sample.
        """

# Path: splearn/utils/validation.py
# def check_rdd(rdd, expected_dtype):
#     """Wrapper function to check_rdd_dtype. Raises TypeError in case of dtype
#     mismatch.
# 
#     Parameters:
#     -----------
#     rdd: splearn.BlockRDD
#         The RDD to check
#     expected_dtype: {type, list of types, tuple of types, dict of types}
#         Expected type(s). If the RDD is a DictRDD the parameter type is
#         restricted to dict.
#     """
#     if not check_rdd_dtype(rdd, expected_dtype):
#         raise TypeError("{0} dtype mismatch.".format(rdd))
# 
# Path: splearn/linear_model/base.py
# class SparkLinearModelMixin(object):
# 
#     def __add__(self, other):
#         """Add method for Linear models with coef and intercept attributes.
# 
#         Parameters
#         ----------
#         other : fitted sklearn linear model
#             Model to add.
# 
#         Returns
#         -------
#         model : Linear model
#             Model with updated coefficients.
#         """
#         model = copy.deepcopy(self)
#         model.coef_ += other.coef_
#         model.intercept_ += other.intercept_
#         return model
# 
#     def __radd__(self, other):
#         """Reverse add method for Linear models.
# 
#         Parameters
#         ----------
#         other : fitted sklearn linear model
#             Model to add.
# 
#         Returns
#         -------
#         model : Linear model
#             Model with updated coefficients.
#         """
#         return self if other == 0 else self.__add__(other)
# 
#     def __div__(self, other):
#         """Division method for Linear models. Used for averaging.
# 
#         Parameters
#         ----------
#         other : integer
#             Integer to divide with.
# 
#         Returns
#         -------
#         model : Linear model
#             Model with updated coefficients.
#         """
#         self.coef_ /= other
#         self.intercept_ /= other
#         return self
# 
#     __truediv__ = __div__
# 
#     def _spark_fit(self, cls, Z, *args, **kwargs):
#         """Wraps a Scikit-learn Linear model's fit method to use with RDD
#         input.
# 
#         Parameters
#         ----------
#         cls : class object
#             The sklearn linear model's class to wrap.
#         Z : TupleRDD or DictRDD
#             The distributed train data in a DictRDD.
# 
#         Returns
#         -------
#         self: the wrapped class
#         """
#         mapper = lambda X_y: super(cls, self).fit(
#             X_y[0], X_y[1], *args, **kwargs
#         )
#         models = Z.map(mapper)
#         avg = models.reduce(operator.add) / models.count()
#         self.__dict__.update(avg.__dict__)
#         return self
# 
#     def _spark_predict(self, cls, X, *args, **kwargs):
#         """Wraps a Scikit-learn Linear model's predict method to use with RDD
#         input.
# 
#         Parameters
#         ----------
#         cls : class object
#             The sklearn linear model's class to wrap.
#         Z : ArrayRDD
#             The distributed data to predict in a DictRDD.
# 
#         Returns
#         -------
#         self: the wrapped class
#         """
#         return X.map(lambda X: super(cls, self).predict(X, *args, **kwargs))
# 
#     def _to_scikit(self, cls):
#         new = cls()
#         new.__dict__ = self.__dict__
#         return new


# Path: splearn/linear_model/stochastic_gradient.py
import numpy as np
import scipy.sparse as sp
from sklearn.linear_model import SGDClassifier
from ..utils.validation import check_rdd
from .base import SparkLinearModelMixin

        Parameters
        ----------
        Z : DictRDD containing (X, y) pairs
            X - Training vector
            y - Target labels
        classes : iterable
            The set of available classes

        Returns
        -------
        self : object
            Returns self.
        """
        check_rdd(Z, {'X': (sp.spmatrix, np.ndarray)})
        self._classes_ = np.unique(classes)
        return self._spark_fit(SparkSGDClassifier, Z)

    def predict(self, X):
        """Distributed method to predict class labels for samples in X.

        Parameters
        ----------
        X : ArrayRDD containing {array-like, sparse matrix}
            Samples.

        Returns
        -------
        C : ArrayRDD
            Predicted class label per sample.
        """

        check_rdd(X, (sp.spmatrix, np.ndarray))