text: # -------------------------------------------------------------
  #
  # 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.
  #
  # -------------------------------------------------------------
  
  __all__ = [
      'LinearRegression',
      'LogisticRegression',
      'SVM',
      'NaiveBayes',
      'Caffe2DML',
      'Keras2DML']
  
  import numpy as np
  from pyspark.ml import Estimator
  from pyspark.ml.feature import VectorAssembler
  from pyspark.sql import DataFrame
  import sklearn as sk
  from sklearn.metrics import accuracy_score, r2_score
  from py4j.protocol import Py4JError
  import traceback
  from sklearn.preprocessing import LabelEncoder
  import threading, time, math, os
  
  from ..converters import *
  from ..classloader import *
  
  
  def assemble(sparkSession, pdf, inputCols, outputCol):
      tmpDF = sparkSession.createDataFrame(pdf, list(pdf.columns))
      assembler = VectorAssembler(inputCols=list(inputCols), outputCol=outputCol)
      return assembler.transform(tmpDF)
  
  
  class BaseSystemMLEstimator(Estimator):
      features_col = 'features'
      label_col = 'label'
  
      def set_features_col(self, colName):
          """
          Sets the default column name for features of PySpark DataFrame.
  
          Parameters
          ----------
          colName: column name for features (default: 'features')
          """
          self.features_col = colName
  
      def set_label_col(self, colName):
          """
          Sets the default column name for features of PySpark DataFrame.
  
          Parameters
          ----------
          colName: column name for features (default: 'label')
          """
          self.label_col = colName
  
      def setGPU(self, enable):
          """
          Whether or not to enable GPU.
  
          Parameters
          ----------
          enable: boolean
          """
          self.estimator.setGPU(enable)
          return self
  
      def setForceGPU(self, enable):
          """
          Whether or not to force the usage of GPU operators.
  
          Parameters
          ----------
          enable: boolean
          """
          self.estimator.setForceGPU(enable)
          return self
  
      def setExplain(self, explain):
          """
          Explanation about the program. Mainly intended for developers.
  
          Parameters
          ----------
          explain: boolean
          """
          self.estimator.setExplain(explain)
          return self
  
      def setExplainLevel(self, explainLevel):
          """
          Set explain level. Mainly intended for developers.
  
          Parameters
          ----------
          explainLevel: string
              Can be one of "hops", "runtime", "recompile_hops", "recompile_runtime"
              or in the above in upper case.
          """
          self.estimator.setExplainLevel(explainLevel)
          return self
  
      def setStatistics(self, statistics):
          """
          Whether or not to output statistics (such as execution time, elapsed time)
          about script executions.
  
          Parameters
          ----------
          statistics: boolean
          """
          self.estimator.setStatistics(statistics)
          return self
  
      def setStatisticsMaxHeavyHitters(self, maxHeavyHitters):
          """
          The maximum number of heavy hitters that are printed as part of the statistics.
  
          Parameters
          ----------
          maxHeavyHitters: int
          """
          self.estimator.setStatisticsMaxHeavyHitters(maxHeavyHitters)
          return self
  
      def setConfigProperty(self, propertyName, propertyValue):
          """
          Set configuration property, such as setConfigProperty("sysml.localtmpdir", "/tmp/systemml").
  
          Parameters
          ----------
          propertyName: String
          propertyValue: String
          """
          self.estimator.setConfigProperty(propertyName, propertyValue)
          return self
  
      def getConfigProperty(self, propertyName):
          """
          Get configuration property, such as getConfigProperty("sysml.localtmpdir").
  
          Parameters
          ----------
          propertyName: String
          """
          return self.estimator.getConfigProperty(propertyName)
  
      def _fit_df(self):
          global default_jvm_stdout, default_jvm_stdout_parallel_flush
          try:
              if default_jvm_stdout:
                  with jvm_stdout(parallel_flush=default_jvm_stdout_parallel_flush):
                      self.model = self.estimator.fit(self.X._jdf)
              else:
                  self.model = self.estimator.fit(self.X._jdf)
          except Py4JError:
              traceback.print_exc()
  
      def fit_df(self, X):
          self.X = X
          self._fit_df()
          self.X = None
          return self
  
      def _fit_numpy(self):
          global default_jvm_stdout, default_jvm_stdout_parallel_flush
          try:
              if isinstance(self.y, np.ndarray) and len(self.y.shape) == 1:
                  # Since we know that mllearn always needs a column vector
                  self.y = np.matrix(self.y).T
              y_mb = convertToMatrixBlock(self.sc, self.y)
              if default_jvm_stdout:
                  with jvm_stdout(parallel_flush=default_jvm_stdout_parallel_flush):
                      self.model = self.estimator.fit(
                          convertToMatrixBlock(self.sc, self.X), y_mb)
              else:
                  self.model = self.estimator.fit(
                      convertToMatrixBlock(self.sc, self.X), y_mb)
          except Py4JError:
              traceback.print_exc()
  
      def fit_numpy(self, X, y):
          self.X = X
          self.y = y
          self._fit_numpy()
          self.X = None
          self.y = None
          return self
  
      def fit_file(self, X_file, y_file):
          global default_jvm_stdout, default_jvm_stdout_parallel_flush
          try:
              if default_jvm_stdout:
                  with jvm_stdout(parallel_flush=default_jvm_stdout_parallel_flush):
                      self.model = self.estimator.fit(X_file, y_file)
              else:
                  self.model = self.estimator.fit(X_file, y_file)
          except Py4JError:
              traceback.print_exc()
          return self
  
      # Returns a model after calling fit(df) on Estimator object on JVM
      def _fit(self, X):
          """
          Invokes the fit method on Estimator object on JVM if X is PySpark DataFrame
  
          Parameters
          ----------
          X: PySpark DataFrame that contain the columns features_col (default: 'features') and label_col (default: 'label')
          """
          if hasattr(
                  X, '_jdf') and self.features_col in X.columns and self.label_col in X.columns:
              return self.fit_df(X)
          else:
              raise Exception(
                  'Incorrect usage: Expected dataframe as input with features/label as columns')
  
      def fit(self, X, y=None, params=None):
          """
          Invokes the fit method on Estimator object on JVM if X and y are on of the supported data types
  
          Parameters
          ----------
          X: NumPy ndarray, Pandas DataFrame, scipy sparse matrix, Spark DataFrame, file path
          y: NumPy ndarray, Pandas DataFrame, scipy sparse matrix, file path
          """
          if y is None:
              return self._fit(X)
          elif isinstance(X, str) and isinstance(y, str):
              return self.fit_file(X, y)
          elif isinstance(X, SUPPORTED_TYPES) and isinstance(y, SUPPORTED_TYPES):
              # Donot encode if y is a numpy matrix => useful for segmentation
              skipEncodingY = len(
                  y.shape) == 2 and y.shape[0] != 1 and y.shape[1] != 1
              y = y if skipEncodingY else self.encode(y)
              if self.transferUsingDF:
                  pdfX = convertToPandasDF(X)
                  pdfY = convertToPandasDF(y)
                  if getNumCols(pdfY) != 1 and not skipEncodingY:
                      raise Exception('y should be a column vector')
                  if pdfX.shape[0] != pdfY.shape[0]:
                      raise Exception('Number of rows of X and y should match')
                  colNames = pdfX.columns
                  pdfX[self.label_col] = pdfY[pdfY.columns[0]]
                  df = assemble(
                      self.sparkSession,
                      pdfX,
                      colNames,
                      self.features_col).select(
                      self.features_col,
                      self.label_col)
                  self.fit_df(df)
              else:
                  numColsy = getNumCols(y)
                  if numColsy != 1 and not skipEncodingY:
                      raise Exception('Expected y to be a column vector')
                  self.fit_numpy(X, y)
              if self.setOutputRawPredictionsToFalse:
                  self.model.setOutputRawPredictions(False)
              return self
          else:
              raise Exception('Unsupported input type')
  
      def transform(self, X):
          return self.predict(X)
  
      def _convertPythonXToJavaObject(self, X):
          """
          Converts the input python object X to a java-side object (either MatrixBlock or Java DataFrame)
  
          Parameters
          ----------
          X: NumPy ndarray, Pandas DataFrame, scipy sparse matrix or PySpark DataFrame
          """
          if isinstance(X, SUPPORTED_TYPES) and self.transferUsingDF:
              pdfX = convertToPandasDF(X)
              df = assemble(
                  self.sparkSession,
                  pdfX,
                  pdfX.columns,
                  self.features_col).select(
                  self.features_col)
              return df._jdf
          elif isinstance(X, SUPPORTED_TYPES):
              return convertToMatrixBlock(self.sc, X)
          elif hasattr(X, '_jdf') and self.features_col in X.columns:
              # No need to assemble as input DF is likely coming via MLPipeline
              return X._jdf
          elif hasattr(X, '_jdf'):
              assembler = VectorAssembler(
                  inputCols=X.columns, outputCol=self.features_col)
              df = assembler.transform(X)
              return df._jdf
          else:
              raise Exception('Unsupported input type')
  
      def _convertJavaOutputToPythonObject(self, X, output):
          """
          Converts the a java-side object output (either MatrixBlock or Java DataFrame) to a python object (based on the type of X).
  
          Parameters
          ----------
          X: NumPy ndarray, Pandas DataFrame, scipy sparse matrix or PySpark DataFrame
          output: a java-side object (either MatrixBlock or Java DataFrame)
          """
          if isinstance(X, SUPPORTED_TYPES) and self.transferUsingDF:
              retDF = DataFrame(output, self.sparkSession._wrapped)
              retPDF = retDF.sort('__INDEX').select('prediction').toPandas()
              return retPDF.as_matrix().flatten() if isinstance(X, np.ndarray) else retPDF
          elif isinstance(X, SUPPORTED_TYPES):
              return convertToNumPyArr(self.sc, output)
          elif hasattr(X, '_jdf'):
              retDF = DataFrame(output, self.sparkSession)
              # Return DF
              return retDF.sort('__INDEX')
          else:
              raise Exception('Unsupported input type')
  
      def predict_proba(self, X):
          """
          Invokes the transform_probability method on Estimator object on JVM if X and y are on of the supported data types
          Return predicted class probabilities for X.
  
          Parameters
          ----------
          X: NumPy ndarray, Pandas DataFrame, scipy sparse matrix or PySpark DataFrame
          """
          global default_jvm_stdout, default_jvm_stdout_parallel_flush
          if hasattr(X, '_jdf'):
              return self.predict(X)
          elif self.transferUsingDF:
              raise ValueError(
                  'The parameter transferUsingDF is not valid for the method predict_proba')
          try:
              if self.estimator is not None and self.model is not None:
                  self.estimator.copyProperties(self.model)
          except AttributeError:
              pass
          try:
              if isinstance(X, str):
                  return self.model.transform_probability(X)
              jX = self._convertPythonXToJavaObject(X)
              if default_jvm_stdout:
                  with jvm_stdout(parallel_flush=default_jvm_stdout_parallel_flush):
                      return self._convertJavaOutputToPythonObject(
                          X, self.model.transform_probability(jX))
              else:
                  return self._convertJavaOutputToPythonObject(
                      X, self.model.transform_probability(jX))
          except Py4JError:
              traceback.print_exc()
  
      # Returns either a DataFrame or MatrixBlock after calling
      # transform(X:MatrixBlock, y:MatrixBlock) on Model object on JVM
      def predict(self, X):
          """
          Invokes the transform method on Estimator object on JVM if X and y are on of the supported data types
  
          Parameters
          ----------
          X: NumPy ndarray, Pandas DataFrame, scipy sparse matrix or PySpark DataFrame or file path
          """
          global default_jvm_stdout, default_jvm_stdout_parallel_flush
          try:
              if self.estimator is not None and self.model is not None:
                  self.estimator.copyProperties(self.model)
          except AttributeError:
              pass
          try:
              if isinstance(X, str):
                  return self.model.transform(X)
              jX = self._convertPythonXToJavaObject(X)
              if default_jvm_stdout:
                  with jvm_stdout(parallel_flush=default_jvm_stdout_parallel_flush):
                      ret = self._convertJavaOutputToPythonObject(
                          X, self.model.transform(jX))
              else:
                  ret = self._convertJavaOutputToPythonObject(
                      X, self.model.transform(jX))
              return self.decode(ret) if isinstance(X, SUPPORTED_TYPES) else ret
          except Py4JError:
              traceback.print_exc()
  
  
  class BaseSystemMLClassifier(BaseSystemMLEstimator):
  
      def encode(self, y):
          self.le = LabelEncoder()
          self.le.fit(y)
          return self.le.transform(y) + 1
  
      def decode(self, y):
          if not hasattr(self, 'le'):
              self.le = None
          if not hasattr(self, 'labelMap'):
              self.labelMap = None
          if self.le is not None:
              return self.le.inverse_transform(np.asarray(y - 1, dtype=int))
          elif self.labelMap is not None:
              return [self.labelMap[int(i)] for i in y]
          else:
              return y
  
      def predict(self, X):
          predictions = super(BaseSystemMLClassifier, self).predict(X)
          from pyspark.sql.dataframe import DataFrame as df
          if isinstance(predictions, df):
              return predictions
          else:
              try:
                  return np.asarray(predictions, dtype='double')
              except ValueError:
                  print(type(predictions))
                  return np.asarray(predictions, dtype='str')
  
      def score(self, X, y):
          """
          Scores the predicted value with ground truth 'y'
  
          Parameters
          ----------
          X: NumPy ndarray, Pandas DataFrame, scipy sparse matrix
          y: NumPy ndarray, Pandas DataFrame, scipy sparse matrix
          """
          predictions = np.asarray(self.predict(X))
          if np.issubdtype(predictions.dtype.type, np.number):
              return accuracy_score(y, predictions)
          else:
              return accuracy_score(np.asarray(
                  y, dtype='str'), np.asarray(predictions, dtype='str'))
  
      def loadLabels(self, file_path):
          createJavaObject(self.sc, 'dummy')
          utilObj = self.sc._jvm.org.apache.sysml.api.ml.Utils()
          if utilObj.checkIfFileExists(file_path):
              df = self.sparkSession.read.csv(file_path, header=False).toPandas()
              keys = np.asarray(df._c0, dtype='int')
              values = np.asarray(df._c1, dtype='str')
              self.labelMap = {}
              for i in range(len(keys)):
                  self.labelMap[int(keys[i])] = values[i]
              # self.encode(classes) # Giving incorrect results
  
      def load(self, weights, sep='/', eager=False):
          """
          Load a pretrained model.
  
          Parameters
          ----------
          weights: directory whether learned weights are stored
          sep: seperator to use (default: '/')
          eager: load the model eagerly. This flag should be only used for debugging purposes. (default: False)
          """
          self.weights = weights
          global default_jvm_stdout, default_jvm_stdout_parallel_flush
          if default_jvm_stdout:
              with jvm_stdout(parallel_flush=default_jvm_stdout_parallel_flush):
                  self.model.load(self.sc._jsc, weights, sep, eager)
          else:
              self.model.load(self.sc._jsc, weights, sep, eager)
          self.loadLabels(weights + '/labels.txt')
  
      def save(self, outputDir, format='binary', sep='/'):
          """
          Save a trained model.
  
          Parameters
          ----------
          outputDir: Directory to save the model to
          format: optional format (default: 'binary')
          sep: seperator to use (default: '/')
          """
          global default_jvm_stdout, default_jvm_stdout_parallel_flush
          if self.model is not None:
              if default_jvm_stdout:
                  with jvm_stdout(parallel_flush=default_jvm_stdout_parallel_flush):
                      self.model.save(self.sc._jsc, outputDir, format, sep)
              else:
                  self.model.save(self.sc._jsc, outputDir, format, sep)
  
              labelMapping = None
              if hasattr(self, 'le') and self.le is not None:
                  labelMapping = dict(enumerate(list(self.le.classes_), 1))
              elif hasattr(self, 'labelMap') and self.labelMap is not None:
                  labelMapping = self.labelMap
  
              if labelMapping is not None:
                  lStr = [[int(k), str(labelMapping[k])] for k in labelMapping]
                  df = self.sparkSession.createDataFrame(lStr)
                  df.write.csv(
                      outputDir + sep + 'labels.txt',
                      mode='overwrite',
                      header=False)
          else:
              raise Exception(
                  'Cannot save as you need to train the model first using fit')
          return self
  
  
  class BaseSystemMLRegressor(BaseSystemMLEstimator):
  
      def encode(self, y):
          return y
  
      def decode(self, y):
          return y
  
      def score(self, X, y):
          """
          Scores the predicted value with ground truth 'y'
  
          Parameters
          ----------
          X: NumPy ndarray, Pandas DataFrame, scipy sparse matrix
          y: NumPy ndarray, Pandas DataFrame, scipy sparse matrix
          """
          return r2_score(y, self.predict(X), multioutput='variance_weighted')
  
      def load(self, weights=None, sep='/', eager=False):
          """
          Load a pretrained model.
  
          Parameters
          ----------
          weights: directory whether learned weights are stored (default: None)
          sep: seperator to use (default: '/')
          eager: load the model eagerly (default: False)
          """
          self.weights = weights
          global default_jvm_stdout, default_jvm_stdout_parallel_flush
          if default_jvm_stdout:
              with jvm_stdout(parallel_flush=default_jvm_stdout_parallel_flush):
                  self.model.load(self.sc._jsc, weights, sep, eager)
          else:
              self.model.load(self.sc._jsc, weights, sep, eager)
  
      def save(self, outputDir, format='binary', sep='/'):
          """
          Save a trained model.
  
          Parameters
          ----------
          outputDir: Directory to save the model to
          format: optional format (default: 'binary')
          sep: seperator to use (default: '/')
          """
          global default_jvm_stdout, default_jvm_stdout_parallel_flush
          if self.model is not None:
              if default_jvm_stdout:
                  with jvm_stdout(parallel_flush=default_jvm_stdout_parallel_flush):
                      self.model.save(outputDir, format, sep)
              else:
                  self.model.save(outputDir, format, sep)
          else:
              raise Exception(
                  'Cannot save as you need to train the model first using fit')
          return self
  
  
  class LogisticRegression(BaseSystemMLClassifier):
      """
      Performs both binomial and multinomial logistic regression.
  
      Examples
      --------
  
      Scikit-learn way
  
      >>> from sklearn import datasets, neighbors
      >>> from systemml.mllearn import LogisticRegression
      >>> from pyspark.sql import SparkSession
      >>> sparkSession = SparkSession.builder.getOrCreate()
      >>> digits = datasets.load_digits()
      >>> X_digits = digits.data
      >>> y_digits = digits.target + 1
      >>> n_samples = len(X_digits)
      >>> X_train = X_digits[:.9 * n_samples]
      >>> y_train = y_digits[:.9 * n_samples]
      >>> X_test = X_digits[.9 * n_samples:]
      >>> y_test = y_digits[.9 * n_samples:]
      >>> logistic = LogisticRegression(sparkSession)
      >>> print('LogisticRegression score: %f' % logistic.fit(X_train, y_train).score(X_test, y_test))
  
      MLPipeline way
  
      >>> from pyspark.ml import Pipeline
      >>> from systemml.mllearn import LogisticRegression
      >>> from pyspark.ml.feature import HashingTF, Tokenizer
      >>> from pyspark.sql import SparkSession
      >>> sparkSession = SparkSession.builder.getOrCreate()
      >>> training = sparkSession.createDataFrame([
      >>>     (0L, "a b c d e spark", 1.0),
      >>>     (1L, "b d", 2.0),
      >>>     (2L, "spark f g h", 1.0),
      >>>     (3L, "hadoop mapreduce", 2.0),
      >>>     (4L, "b spark who", 1.0),
      >>>     (5L, "g d a y", 2.0),
      >>>     (6L, "spark fly", 1.0),
      >>>     (7L, "was mapreduce", 2.0),
      >>>     (8L, "e spark program", 1.0),
      >>>     (9L, "a e c l", 2.0),
      >>>     (10L, "spark compile", 1.0),
      >>>     (11L, "hadoop software", 2.0)
      >>> ], ["id", "text", "label"])
      >>> tokenizer = Tokenizer(inputCol="text", outputCol="words")
      >>> hashingTF = HashingTF(inputCol="words", outputCol="features", numFeatures=20)
      >>> lr = LogisticRegression(sparkSession)
      >>> pipeline = Pipeline(stages=[tokenizer, hashingTF, lr])
      >>> model = pipeline.fit(training)
      >>> test = sparkSession.createDataFrame([
      >>>     (12L, "spark i j k"),
      >>>     (13L, "l m n"),
      >>>     (14L, "mapreduce spark"),
      >>>     (15L, "apache hadoop")], ["id", "text"])
      >>> prediction = model.transform(test)
      >>> prediction.show()
  
      """
  
      def __init__(self, sparkSession, penalty='l2', fit_intercept=True, normalize=False, max_iter=100,
                   max_inner_iter=0, tol=0.000001, C=1.0, solver='newton-cg', transferUsingDF=False):
          """
          Performs both binomial and multinomial logistic regression.
  
          Parameters
          ----------
          sparkSession: PySpark SparkSession
          penalty: Only 'l2' supported
          fit_intercept: Specifies whether to add intercept or not (default: True)
          normalize: This parameter is ignored when fit_intercept is set to False. (default: False)
          max_iter: Maximum number of outer (Fisher scoring) iterations (default: 100)
          max_inner_iter: Maximum number of inner (conjugate gradient) iterations, or 0 if no maximum limit provided (default: 0)
          tol: Tolerance used in the convergence criterion (default: 0.000001)
          C: 1/regularization parameter (default: 1.0 similar to scikit-learn. To disable regularization, please use float("inf"))
          solver: Only 'newton-cg' solver supported
          """
          self.sparkSession = sparkSession
          self.sc = sparkSession._sc
          createJavaObject(self.sc, 'dummy')
          self.uid = "logReg"
          self.estimator = self.sc._jvm.org.apache.sysml.api.ml.LogisticRegression(
              self.uid, self.sc._jsc.sc())
          self.estimator.setMaxOuterIter(max_iter)
          self.estimator.setMaxInnerIter(max_inner_iter)
          reg = 0.0 if C == float("inf") else 1.0 / C
          icpt = 2 if fit_intercept == True and normalize == True else int(
              fit_intercept)
          self.estimator.setRegParam(reg)
          self.estimator.setTol(tol)
          self.estimator.setIcpt(icpt)
          self.transferUsingDF = transferUsingDF
          self.setOutputRawPredictionsToFalse = True
          self.model = self.sc._jvm.org.apache.sysml.api.ml.LogisticRegressionModel(
              self.estimator)
          if penalty != 'l2':
              raise Exception('Only l2 penalty is supported')
          if solver != 'newton-cg':
              raise Exception('Only newton-cg solver supported')
  
  
  class LinearRegression(BaseSystemMLRegressor):
      """
      Performs linear regression to model the relationship between one numerical response variable and one or more explanatory (feature) variables.
  
      Examples
      --------
  
      >>> import numpy as np
      >>> from sklearn import datasets
      >>> from systemml.mllearn import LinearRegression
      >>> from pyspark.sql import SparkSession
      >>> # Load the diabetes dataset
      >>> diabetes = datasets.load_diabetes()
      >>> # Use only one feature
      >>> diabetes_X = diabetes.data[:, np.newaxis, 2]
      >>> # Split the data into training/testing sets
      >>> diabetes_X_train = diabetes_X[:-20]
      >>> diabetes_X_test = diabetes_X[-20:]
      >>> # Split the targets into training/testing sets
      >>> diabetes_y_train = diabetes.target[:-20]
      >>> diabetes_y_test = diabetes.target[-20:]
      >>> # Create linear regression object
      >>> regr = LinearRegression(sparkSession, solver='newton-cg')
      >>> # Train the model using the training sets
      >>> regr.fit(diabetes_X_train, diabetes_y_train)
      >>> # The mean square error
      >>> print("Residual sum of squares: %.2f" % np.mean((regr.predict(diabetes_X_test) - diabetes_y_test) ** 2))
  
      """
  
      def __init__(self, sparkSession, fit_intercept=True, normalize=False, max_iter=100,
                   tol=0.000001, C=float("inf"), solver='newton-cg', transferUsingDF=False):
          """
          Performs linear regression to model the relationship between one numerical response variable and one or more explanatory (feature) variables.
  
          Parameters
          ----------
          sparkSession: PySpark SparkSession
          fit_intercept: Specifies whether to add intercept or not (default: True)
          normalize: If True, the regressors X will be normalized before regression. This parameter is ignored when fit_intercept is set to False. (default: False)
          max_iter: Maximum number of conjugate gradient iterations, or 0 if no maximum limit provided (default: 100)
          tol: Tolerance used in the convergence criterion (default: 0.000001)
          C: 1/regularization parameter (default: float("inf") as scikit learn doesnot support regularization by default)
          solver: Supports either 'newton-cg' or 'direct-solve' (default: 'newton-cg').
          Depending on the size and the sparsity of the feature matrix, one or the other solver may be more efficient.
          'direct-solve' solver is more efficient when the number of features is relatively small (m < 1000) and
          input matrix X is either tall or fairly dense; otherwise 'newton-cg' solver is more efficient.
          """
          self.sparkSession = sparkSession
          self.sc = sparkSession._sc
          createJavaObject(self.sc, 'dummy')
          self.uid = "lr"
          if solver == 'newton-cg' or solver == 'direct-solve':
              self.estimator = self.sc._jvm.org.apache.sysml.api.ml.LinearRegression(
                  self.uid, self.sc._jsc.sc(), solver)
          else:
              raise Exception('Only newton-cg solver supported')
          self.estimator.setMaxIter(max_iter)
          reg = 0.0 if C == float("inf") else 1.0 / C
          icpt = 2 if fit_intercept == True and normalize == True else int(
              fit_intercept)
          self.estimator.setRegParam(reg)
          self.estimator.setTol(tol)
          self.estimator.setIcpt(icpt)
          self.transferUsingDF = transferUsingDF
          self.setOutputRawPredictionsToFalse = False
          self.model = self.sc._jvm.org.apache.sysml.api.ml.LinearRegressionModel(
              self.estimator)
  
  
  class SVM(BaseSystemMLClassifier):
      """
      Performs both binary-class and multiclass SVM (Support Vector Machines).
  
      Examples
      --------
  
      >>> from sklearn import datasets, neighbors
      >>> from systemml.mllearn import SVM
      >>> from pyspark.sql import SparkSession
      >>> sparkSession = SparkSession.builder.getOrCreate()
      >>> digits = datasets.load_digits()
      >>> X_digits = digits.data
      >>> y_digits = digits.target
      >>> n_samples = len(X_digits)
      >>> X_train = X_digits[:.9 * n_samples]
      >>> y_train = y_digits[:.9 * n_samples]
      >>> X_test = X_digits[.9 * n_samples:]
      >>> y_test = y_digits[.9 * n_samples:]
      >>> svm = SVM(sparkSession, is_multi_class=True)
      >>> print('LogisticRegression score: %f' % svm.fit(X_train, y_train).score(X_test, y_test))
  
      """
  
      def __init__(self, sparkSession, fit_intercept=True, normalize=False, max_iter=100,
                   tol=0.000001, C=1.0, is_multi_class=False, transferUsingDF=False):
          """
          Performs both binary-class and multiclass SVM (Support Vector Machines).
  
          Parameters
          ----------
          sparkSession: PySpark SparkSession
          fit_intercept: Specifies whether to add intercept or not (default: True)
          normalize: This parameter is ignored when fit_intercept is set to False. (default: False)
          max_iter: Maximum number iterations (default: 100)
          tol: Tolerance used in the convergence criterion (default: 0.000001)
          C: 1/regularization parameter (default: 1.0 similar to scikit-learn. To disable regularization, please use float("inf"))
          is_multi_class: Specifies whether to use binary-class SVM or multi-class SVM algorithm (default: False)
          """
          self.sparkSession = sparkSession
          self.sc = sparkSession._sc
          self.uid = "svm"
          createJavaObject(self.sc, 'dummy')
          self.is_multi_class = is_multi_class
          self.estimator = self.sc._jvm.org.apache.sysml.api.ml.SVM(
              self.uid, self.sc._jsc.sc(), is_multi_class)
          self.estimator.setMaxIter(max_iter)
          if C <= 0:
              raise Exception('C has to be positive')
          reg = 0.0 if C == float("inf") else 1.0 / C
          icpt = 2 if fit_intercept == True and normalize == True else int(
              fit_intercept)
          self.estimator.setRegParam(reg)
          self.estimator.setTol(tol)
          self.estimator.setIcpt(icpt)
          self.transferUsingDF = transferUsingDF
          self.setOutputRawPredictionsToFalse = False
          self.model = self.sc._jvm.org.apache.sysml.api.ml.SVMModel(
              self.estimator, self.is_multi_class)
  
  
  class NaiveBayes(BaseSystemMLClassifier):
      """
      Performs Naive Bayes.
  
      Examples
      --------
  
      >>> from sklearn.datasets import fetch_20newsgroups
      >>> from sklearn.feature_extraction.text import TfidfVectorizer
      >>> from systemml.mllearn import NaiveBayes
      >>> from sklearn import metrics
      >>> from pyspark.sql import SparkSession
      >>> sparkSession = SparkSession.builder.getOrCreate(sc)
      >>> categories = ['alt.atheism', 'talk.religion.misc', 'comp.graphics', 'sci.space']
      >>> newsgroups_train = fetch_20newsgroups(subset='train', categories=categories)
      >>> newsgroups_test = fetch_20newsgroups(subset='test', categories=categories)
      >>> vectorizer = TfidfVectorizer()
      >>> # Both vectors and vectors_test are SciPy CSR matrix
      >>> vectors = vectorizer.fit_transform(newsgroups_train.data)
      >>> vectors_test = vectorizer.transform(newsgroups_test.data)
      >>> nb = NaiveBayes(sparkSession)
      >>> nb.fit(vectors, newsgroups_train.target)
      >>> pred = nb.predict(vectors_test)
      >>> metrics.f1_score(newsgroups_test.target, pred, average='weighted')
  
      """
  
      def __init__(self, sparkSession, laplace=1.0, transferUsingDF=False):
          """
          Performs Naive Bayes.
  
          Parameters
          ----------
          sparkSession: PySpark SparkSession
          laplace: Laplace smoothing specified by the user to avoid creation of 0 probabilities (default: 1.0)
          """
          self.sparkSession = sparkSession
          self.sc = sparkSession._sc
          self.uid = "nb"
          createJavaObject(self.sc, 'dummy')
          self.estimator = self.sc._jvm.org.apache.sysml.api.ml.NaiveBayes(
              self.uid, self.sc._jsc.sc())
          self.estimator.setLaplace(laplace)
          self.transferUsingDF = transferUsingDF
          self.setOutputRawPredictionsToFalse = False
          self.model = self.sc._jvm.org.apache.sysml.api.ml.NaiveBayesModel(
              self.estimator)
  
  
  class Caffe2DML(BaseSystemMLClassifier):
      """
      Performs training/prediction for a given caffe network.
  
      Examples
      --------
  
      >>> from systemml.mllearn import Caffe2DML
      >>> from mlxtend.data import mnist_data
      >>> import numpy as np
      >>> from sklearn.utils import shuffle
      >>> X, y = mnist_data()
      >>> X, y = shuffle(X, y)
      >>> imgShape = (1, 28, 28)
      >>> import urllib
      >>> urllib.urlretrieve('https://raw.githubusercontent.com/niketanpansare/model_zoo/master/caffe/vision/lenet/mnist/lenet.proto', 'lenet.proto')
      >>> urllib.urlretrieve('https://raw.githubusercontent.com/niketanpansare/model_zoo/master/caffe/vision/lenet/mnist/lenet_solver.proto', 'lenet_solver.proto')
      >>> caffe2DML = Caffe2DML(spark, 'lenet_solver.proto').set(max_iter=500)
      >>> caffe2DML.fit(X, y)
      """
  
      def __init__(self, sparkSession, solver, input_shape,
                   transferUsingDF=False):
          """
          Performs training/prediction for a given caffe network.
  
          Parameters
          ----------
          sparkSession: PySpark SparkSession
          solver: caffe solver file path
          input_shape: 3-element list (number of channels, input height, input width)
          transferUsingDF: whether to pass the input dataset via PySpark DataFrame (default: False)
          """
          self.sparkSession = sparkSession
          self.sc = sparkSession._sc
          createJavaObject(self.sc, 'dummy')
          self.uid = "Caffe2DML"
          self.model = None
          if len(input_shape) != 3:
              raise ValueError('Expected input_shape as list of 3 element')
          solver = self.sc._jvm.org.apache.sysml.api.dl.Utils.readCaffeSolver(
              solver)
          self.estimator = self.sc._jvm.org.apache.sysml.api.dl.Caffe2DML(
              self.sc._jsc.sc(), solver, str(
                  input_shape[0]), str(
                  input_shape[1]), str(
                  input_shape[2]))
          self.transferUsingDF = transferUsingDF
          self.setOutputRawPredictionsToFalse = False
  
      def load(self, weights=None, sep='/', ignore_weights=None, eager=False):
          """
          Load a pretrained model.
  
          Parameters
          ----------
          weights: directory whether learned weights are stored (default: None)
          sep: seperator to use (default: '/')
          ignore_weights: names of layers to not read from the weights directory (list of string, default:None)
          eager: load the model eagerly (default: False)
          """
          global default_jvm_stdout, default_jvm_stdout_parallel_flush
          self.weights = weights
          self.estimator.setInput("$weights", str(weights))
          self.model = self.sc._jvm.org.apache.sysml.api.dl.Caffe2DMLModel(
              self.estimator)
          if default_jvm_stdout:
              with jvm_stdout(parallel_flush=default_jvm_stdout_parallel_flush):
                  self.model.load(self.sc._jsc, weights, sep, eager)
          else:
              self.model.load(self.sc._jsc, weights, sep, eager)
          self.loadLabels(weights + '/labels.txt')
          if ignore_weights is not None:
              self.estimator.setWeightsToIgnore(ignore_weights)
  
      def set(self, debug=None, train_algo=None, test_algo=None, parallel_batches=None,
              output_activations=None, perform_one_hot_encoding=None, parfor_parameters=None, inline_nn_library=None, use_builtin_lstm_fn=None,
              perform_fused_backward_update=None, weight_parallel_batches=None):
          """
          Set input to Caffe2DML
  
          Parameters
          ----------
          debug: to add debugging DML code such as classification report, print DML script, etc (default: False)
          train_algo: can be minibatch, batch, allreduce_parallel_batches, looped_minibatch or allreduce (default: minibatch)
          test_algo: can be minibatch, batch, allreduce_parallel_batches, looped_minibatch or allreduce (default: minibatch)
          parallel_batches: number of parallel batches (required for allreduce_parallel_batches or looped_minibatch)
          output_activations: (developer flag) directory to output activations of each layer as csv while prediction. To be used only in batch mode (default: None)
          perform_one_hot_encoding: should perform one-hot encoding in DML using table function (default: True)
          parfor_parameters: dictionary for parfor parameters when using allreduce-style algorithms (default: "")
          inline_nn_library: whether to inline the NN library when generating DML using Caffe2DML (default: False)
          use_builtin_lstm_fn: whether to use builtin lstm function for LSTM layer (default: True)
          perform_fused_backward_update: whether to perform update immediately after backward pass at the script level. Supported for minibatch and batch algorithms. (default: True)
          weight_parallel_batches: whether to multiply 1/parallel_batches to gradients before performing SGD update (default: True)
          """
          if debug is not None:
              self.estimator.setInput("$debug", str(debug).upper())
          if inline_nn_library is not None:
              self.estimator.setInput("$inline_nn_library", str(inline_nn_library).upper())
          if train_algo is not None:
              self.estimator.setInput("$train_algo", str(train_algo).lower())
          if test_algo is not None:
              self.estimator.setInput("$test_algo", str(test_algo).lower())
          if parallel_batches is not None:
              self.estimator.setInput("$parallel_batches", str(parallel_batches))
          if use_builtin_lstm_fn is not None:
              self.estimator.setInput("$use_builtin_lstm_fn", str(use_builtin_lstm_fn).upper())
          if perform_fused_backward_update is not None:
              self.estimator.setInput("$perform_fused_backward_update", str(perform_fused_backward_update).upper())
          if weight_parallel_batches is not None:
              self.estimator.setInput("$weight_parallel_batches", str(weight_parallel_batches).upper())
          if output_activations is not None:
              self.estimator.setInput(
                  "$output_activations",
                  str(output_activations))
          if perform_one_hot_encoding is not None:
              self.estimator.setInput(
                  "$perform_one_hot_encoding",
                  str(perform_one_hot_encoding).lower())
          if parfor_parameters is not None:
              if isinstance(parfor_parameters, dict):
                  # Convert dictionary to comma-separated list
                  parfor_parameters = ''.join(
                      [
                          ', ' +
                          str(k) +
                          '=' +
                          str(v) for k,
                          v in parfor_parameters.items()]) if len(parfor_parameters) > 0 else ''
                  self.estimator.setInput(
                      "$parfor_parameters", parfor_parameters)
              else:
                  raise TypeError("parfor_parameters should be a dictionary")
          return self
  
      def get_training_script(self):
          """
          Return the training DML script
          """
          return self.estimator.get_training_script()
          
      def get_prediction_script(self):
          """
          Return the prediction DML script
          """
          return self.estimator.get_prediction_script()
      
      def summary(self):
          """
          Print the summary of the network
          """
          import pyspark
          global default_jvm_stdout, default_jvm_stdout_parallel_flush
          if isinstance(self.sparkSession, pyspark.sql.session.SparkSession):
              if default_jvm_stdout:
                  with jvm_stdout(parallel_flush=default_jvm_stdout_parallel_flush):
                      self.estimator.summary(self.sparkSession._jsparkSession)
              else:
                  self.estimator.summary(self.sparkSession._jsparkSession)
          else:
              raise TypeError(
                  'Please use spark session of type pyspark.sql.session.SparkSession in the constructor')
  
  
  class Keras2DML(Caffe2DML):
      """
      Peforms training/prediction for a given keras model.
  
  
      """
  
      def __init__(self, sparkSession, keras_model, input_shape=None, transferUsingDF=False, load_keras_weights=True, weights=None, labels=None,
                   lr_policy="step", weight_decay=0, regularization_type="L2"):
          """
          Performs training/prediction for a given keras model.
  
          Parameters
          ----------
          sparkSession: PySpark SparkSession
          keras_model: keras model
          input_shape: 3-element list (number of channels, input height, input width). If not provided, it is inferred from the input shape of the first layer.
          transferUsingDF: whether to pass the input dataset via PySpark DataFrame (default: False)
          load_keras_weights: whether to load weights from the keras_model. If False, the weights will be initialized to random value using NN libraries' init method  (default: True)
          weights: directory whether learned weights are stored (default: None)
          labels: file containing mapping between index and string labels (default: None)
          lr_policy: learning rate policy for caffe solver (default: "step")
          weight_decay: regularation strength (default: 0, recommended: 5e-4)
          regularization_type: regularization type (default: "L2")
          """
          from .keras2caffe import convertKerasToCaffeNetwork, convertKerasToCaffeSolver, convertKerasToSystemMLModel
          import tempfile, keras
          if keras.backend.image_data_format() != 'channels_first':
              raise Exception('The data format ' + str(keras.backend.image_data_format())
                              + ' is not supported. Please use keras.backend.set_image_data_format("channels_first")')
          if isinstance(keras_model, keras.models.Sequential):
              # Convert the sequential model to functional model
              if keras_model.model is None:
                  keras_model.build()
              keras_model = keras_model.model
          if input_shape is None:
              keras_shape = keras_model.layers[0].input_shape
              input_shape = [1, 1, 1]
              if len(keras_shape) > 4 or len(keras_shape) <= 1:
                  raise Exception('Input shape ' + str(keras_shape) + ' is not supported.')
              for i in range(len(keras_shape)-1):
                  input_shape[i] = keras_shape[i+1] # Ignore batch size
          elif len(input_shape) > 3 or len(input_shape) == 0:
              raise Exception('Input shape ' + str(input_shape) + ' is not supported.')
          self.name = keras_model.name
          createJavaObject(sparkSession._sc, 'dummy')
          if not hasattr(keras_model, 'optimizer'):
              raise Exception('Please compile the model before passing it to Keras2DML')
          # Default values for Caffe solver. These will be overriden by the fit method.
          batch_size, max_iter, test_iter, test_interval, display = 32, 2000, 0, 500, 100
          convertKerasToCaffeNetwork(
              keras_model,
              self.name + ".proto",
              int(batch_size))
          convertKerasToCaffeSolver(
              keras_model,
              self.name + ".proto",
              self.name + "_solver.proto",
              int(max_iter),
              int(test_iter),
              int(test_interval),
              int(display),
              lr_policy,
              weight_decay,
              regularization_type)
          self.weights = tempfile.mkdtemp() if weights is None else weights
          self.keras_model = keras_model
          if load_keras_weights:
              convertKerasToSystemMLModel(
                  sparkSession, keras_model, self.weights)
          if labels is not None and (labels.startswith(
                  'https:') or labels.startswith('http:')):
              import urllib
              urllib.urlretrieve(labels, os.path.join(weights, 'labels.txt'))
          elif labels is not None:
              from shutil import copyfile
              copyfile(labels, os.path.join(weights, 'labels.txt'))
          super(
              Keras2DML,
              self).__init__(
              sparkSession,
              self.name +
              "_solver.proto",
              input_shape,
              transferUsingDF)
          if load_keras_weights:
              self.load(self.weights)
  
      def close(self):
          import shutil
          shutil.rmtree(weights)
  
      def fit(self, X, y=None, batch_size=32, epochs=1, validation_split=0.0, validation_data=None):
          # verbose, callbacks flags are not supported
          self.estimator.modifySolverParam("batch_size", str(batch_size))
          self.estimator.modifySolverParam("epochs", str(epochs))
          self.estimator.modifySolverParam("validation_split", str(validation_split))
          if batch_size <= 0 or epochs <= 0 or validation_split < 0:
              raise ValueError('Incorrect parameters to fit method: batch_size=' +  str(batch_size) + ', epochs=' +
                               str(epochs) + ', validation_split=' + str(validation_split))
          if validation_data is not None:
              X_val, y_val = validation_data[0], validation_data[1]
              if isinstance(y_val, np.ndarray) and len(y_val.shape) == 1:
                  # Since we know that mllearn always needs a column vector
                  y_val = np.matrix(y_val).T
                  y_val = convertToMatrixBlock(self.sc, y_val)
                  self.estimator.setValidationData(convertToMatrixBlock(self.sc, X_val), y_val)
          if y is not None:
              super(Keras2DML, self).fit(X, y)
          else:
              super(Keras2DML, self).fit(X)
  
  