#!/usr/bin/python
#-------------------------------------------------------------
#
# 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.
#
#-------------------------------------------------------------

from py4j.protocol import Py4JJavaError, Py4JError
import traceback
import os
from pyspark.sql import DataFrame, SQLContext
from pyspark.rdd import RDD


class MLContext(object):

    """
    Simple wrapper class for MLContext in SystemML.jar
    ...
    Attributes
    ----------
    ml : MLContext
        A reference to the java MLContext
    sc : SparkContext
        The SparkContext that has been specified during initialization
    """

    def __init__(self, sc, *args):
        """
        If initialized with a SparkContext, will connect to the Java MLContext
        class.
        args:
            sc (SparkContext): the current SparkContext
            monitor (boolean=False): Whether to monitor the performance
            forceSpark (boolean=False): Whether to force execution on spark
        returns:
            MLContext: Instance of MLContext
        """

        try:
            monitorPerformance = (args[0] if len(args) > 0 else False)
            setForcedSparkExecType = (args[1] if len(args) > 1 else False)
            self.sc = sc
            self.ml = sc._jvm.org.apache.sysml.api.MLContext(sc._jsc, monitorPerformance, setForcedSparkExecType)
        except Py4JError:
            traceback.print_exc()

    def reset(self):
        """
        Call this method of you want to clear any RDDs set via
        registerInput or registerOutput
        """
        try:
            self.ml.reset()
        except Py4JJavaError:
            traceback.print_exc()

    def execute(self, dmlScriptFilePath, *args):
        """
        Executes the script in spark-mode by passing the arguments to the
        MLContext java class.
        Returns:
            MLOutput: an instance of the MLOutput-class
        """
        numArgs = len(args) + 1
        try:
            if numArgs == 1:
                jmlOut = self.ml.execute(dmlScriptFilePath)
                mlOut = MLOutput(jmlOut, self.sc)
                return mlOut
            elif numArgs == 2:
                jmlOut = self.ml.execute(dmlScriptFilePath, args[0])
                mlOut = MLOutput(jmlOut, self.sc)
                return mlOut
            elif numArgs == 3:
                jmlOut = self.ml.execute(dmlScriptFilePath, args[0], args[1])
                mlOut = MLOutput(jmlOut, self.sc)
                return mlOut
            elif numArgs == 4:
                jmlOut = self.ml.execute(dmlScriptFilePath, args[0], args[1], args[2])
                mlOut = MLOutput(jmlOut, self.sc)
                return mlOut
            else:
                raise TypeError('Arguments do not match MLContext-API')
        except Py4JJavaError:
            traceback.print_exc()

    def executeScript(self, dmlScript, nargs=None, outputs=None, isPyDML=False, configFilePath=None):
        """
        Executes the script in spark-mode by passing the arguments to the
        MLContext java class.
        Returns:
            MLOutput: an instance of the MLOutput-class
        """
        try:
            # Register inputs as needed
            if nargs is not None:
                for key, value in nargs.items():
                    if isinstance(value, DataFrame):
                        self.registerInput(key, value)
                        del nargs[key]
                    else:
                        nargs[key] = str(value)
            else:
                nargs = {}

            # Register outputs as needed
            if outputs is not None:
                for out in outputs:
                    self.registerOutput(out)

            # Execute script
            jml_out = self.ml.executeScript(dmlScript, nargs, isPyDML, configFilePath)
            ml_out = MLOutput(jml_out, self.sc)
            return ml_out
        except Py4JJavaError:
            traceback.print_exc()

    def registerInput(self, varName, src, *args):
        """
        Method to register inputs used by the DML script.
        Supported format:
        1. DataFrame
        2. CSV/Text (as JavaRDD<String> or JavaPairRDD<LongWritable, Text>)
        3. Binary blocked RDD (JavaPairRDD<MatrixIndexes,MatrixBlock>))
        Also overloaded to support metadata information such as format, rlen, clen, ...
        Please note the variable names given below in quotes correspond to the variables in DML script.
        These variables need to have corresponding read/write associated in DML script.
        Currently, only matrix variables are supported through registerInput/registerOutput interface.
        To pass scalar variables, use named/positional arguments (described later) or wrap them into matrix variable.
        """
        numArgs = len(args) + 2

        if hasattr(src, '_jdf'):
            rdd = src._jdf
        elif hasattr(src, '_jrdd'):
            rdd = src._jrdd
        else:
            rdd = src

        try:
            if numArgs == 2:
                self.ml.registerInput(varName, rdd)
            elif numArgs == 3:
                self.ml.registerInput(varName, rdd, args[0])
            elif numArgs == 4:
                self.ml.registerInput(varName, rdd, args[0], args[1])
            elif numArgs == 5:
                self.ml.registerInput(varName, rdd, args[0], args[1], args[2])
            elif numArgs == 6:
                self.ml.registerInput(varName, rdd, args[0], args[1], args[2], args[3])
            elif numArgs == 7:
                self.ml.registerInput(varName, rdd, args[0], args[1], args[2], args[3], args[4])
            elif numArgs == 10:
                self.ml.registerInput(varName, rdd, args[0], args[1], args[2], args[3], args[4], args[5], args[6], args[7])
            else:
                raise TypeError('Arguments do not match MLContext-API')
        except Py4JJavaError:

            traceback.print_exc()

    def registerOutput(self, varName):
        """
        Register output variables used in the DML script
        args:
            varName: (String) The name used in the DML script
        """

        try:
            self.ml.registerOutput(varName)
        except Py4JJavaError:
            traceback.print_exc()

    def getDmlJson(self):
        try:
            return self.ml.getMonitoringUtil().getRuntimeInfoInJSONFormat()
        except Py4JJavaError:
            traceback.print_exc()


class MLOutput(object):

    """
    This is a simple wrapper object that returns the output of execute from MLContext
    ...
    Attributes
    ----------
    jmlOut MLContext:
        A reference to the MLOutput object through py4j
    """

    def __init__(self, jmlOut, sc):
        self.jmlOut = jmlOut
        self.sc = sc

    def getBinaryBlockedRDD(self, varName):
        raise Exception('Not supported in Python MLContext')
        #try:
        #    rdd = RDD(self.jmlOut.getBinaryBlockedRDD(varName), self.sc)
        #    return rdd
        #except Py4JJavaError:
        #    traceback.print_exc()

    def getMatrixCharacteristics(self, varName):
        raise Exception('Not supported in Python MLContext')
        #try:
        #    chars = self.jmlOut.getMatrixCharacteristics(varName)
        #    return chars
        #except Py4JJavaError:
        #    traceback.print_exc()

    def getDF(self, sqlContext, varName):
        try:
            jdf = self.jmlOut.getDF(sqlContext._ssql_ctx, varName)
            df = DataFrame(jdf, sqlContext)
            return df
        except Py4JJavaError:
            traceback.print_exc()

    def getMLMatrix(self, sqlContext, varName):
        raise Exception('Not supported in Python MLContext')
        #try:
        #    mlm = self.jmlOut.getMLMatrix(sqlContext._scala_SQLContext, varName)
        #    return mlm
        #except Py4JJavaError:
        #    traceback.print_exc()

    def getStringRDD(self, varName, format):
        raise Exception('Not supported in Python MLContext')
        #try:
        #    rdd = RDD(self.jmlOut.getStringRDD(varName, format), self.sc)
        #    return rdd
        #except Py4JJavaError:
        #    traceback.print_exc()