from pyspark.ml.pipeline import Estimator, Model
from pyspark.ml.param.shared import *
from bigdl.util.common import *

if sys.version >= '3':
    long = int
    unicode = str

class HasBatchSize(Params):
    """
    Mixin for param batchSize: batch size.
    """

    # a placeholder to make it appear in the generated doc
    batchSize = Param(Params._dummy(), "batchSize", "batchSize")

    def __init__(self):
        super(HasBatchSize, self).__init__()
        #: param for batch size.
        self.batchSize = Param(self, "batchSize", "batchSize")
        self._setDefault(batchSize=1)

    def setBatchSize(self, val):
        """
        Sets the value of :py:attr:`batchSize`.
        """
        self._paramMap[self.batchSize] = val
        pythonBigDL_method_name = "setBatchSize" + self.__class__.__name__
        callBigDlFunc(self.bigdl_type, pythonBigDL_method_name, self.value, val)
        return self

    def getBatchSize(self):
        """
        Gets the value of batchSize or its default value.
        """
        return self.getOrDefault(self.batchSize)

class HasMaxEpoch(Params):
    maxEpoch = Param(Params._dummy(), "maxEpoch", "number of max Epoch")

    def __init__(self):
        super(HasMaxEpoch, self).__init__()
        self.maxEpoch = Param(self, "maxEpoch", "maxEpoch")
        self._setDefault(maxEpoch=100)

    def setMaxEpoch(self, val):
        self._paramMap[self.maxEpoch] = val
        pythonBigDL_method_name = "setMaxEpoch" + self.__class__.__name__
        callBigDlFunc(self.bigdl_type, pythonBigDL_method_name, self.value, val)
        return self

    def getMaxEpoch(self):
        """
        Gets the value of maxEpoch or its default value.
        """
        return self.getOrDefault(self.maxEpoch)

class HasFeatureSize(Params):
    featureSize = Param(Params._dummy(), "featureSize", "size of the feature")

    def __init__(self):
        super(HasFeatureSize, self).__init__()
        self.featureSize = Param(self, "featureSize", "featureSize")
        self._setDefault(featureSize=None)

    def setFeatureSize(self, val):
        self._paramMap[self.featureSize] = val
        pythonBigDL_mehtod_name = "setFeatureSize" + self.__class__.__name__
        callBigDlFunc(self.bigdl_type, pythonBigDL_mehtod_name, self.value, val)
        return self

    def getFeatureSize(self):
        return self.getOrDefault(self.featureSize)

class HasLearningRate(Params):
    learningRate = Param(Params._dummy(), "learningRate", "learning rate")

    def __init__(self):
        super(HasLearningRate, self).__init__()
        self.learningRate = Param(self, "learningRate", "learning rate")
        self._setDefault(learningRate=100)

    def setLearningRate(self, val):
        self._paramMap[self.learningRate] = val
        pythonBigDL_method_name = "setLearningRate" + self.__class__.__name__
        callBigDlFunc(self.bigdl_type, pythonBigDL_method_name, self.value, val)
        return self

    def getLearningRate(self):
        """
        Gets the value of maxEpoch or its default value.
        """
        return self.getOrDefault(self.learningRate)

class DLEstimator(Estimator, HasFeaturesCol, HasLabelCol, HasPredictionCol, HasBatchSize, HasMaxEpoch, HasLearningRate, JavaValue):
    """
    .. note:: Deprecated in 0.5.0. `DLEstimator` has been migrated to package
     `bigdl.dlframes`. This will be removed in BigDL 0.6.

    """
    def __init__(self,  model, criterion, feature_size, label_size, jvalue=None, bigdl_type="float"):
        super(DLEstimator, self).__init__()
        self.value = jvalue if jvalue else callBigDlFunc(
            bigdl_type, self.jvm_class_constructor(), model, criterion, feature_size, label_size)
        self.bigdl_type = bigdl_type
        self.featureSize = feature_size

    def _fit(self, dataset):
        #self._transfer_params_to_java()
        jmodel = callBigDlFunc(self.bigdl_type, "fitEstimator", self.value, dataset)
        model = DLModel.of(jmodel, self.featureSize, self.bigdl_type)
        return model

class DLModel(Model, HasFeaturesCol, HasPredictionCol, HasBatchSize, HasFeatureSize, JavaValue):
    """
    .. note:: Deprecated in 0.5.0. `DLModel` has been migrated to package
     `bigdl.dlframes`. This will be removed in BigDL 0.6.

    """
    def __init__(self,  model, featureSize, jvalue=None, bigdl_type="float"):
        super(DLModel, self).__init__()
        self.value = jvalue if jvalue else callBigDlFunc(
            bigdl_type, self.jvm_class_constructor(), model, featureSize)
        self.bigdl_type = bigdl_type
        self.setFeatureSize(featureSize)

    def _transform(self, dataset):
        return callBigDlFunc(self.bigdl_type, "dlModelTransform", self.value, dataset)

    @classmethod
    def of(self, jvalue, feature_size=None, bigdl_type="float"):
        model = DLModel(model=None, featureSize=feature_size, jvalue=jvalue, bigdl_type=bigdl_type)
        return model

class DLClassifier(DLEstimator):
    """
    .. note:: Deprecated in 0.5.0. `DLClassifier` has been migrated to package
     `bigdl.dlframes`. This will be removed in BigDL 0.6.

    """
    def __init__(self,  model, criterion, feature_size,  bigdl_type="float"):
        super(DLClassifier, self).__init__(model, criterion, feature_size, [1], None, bigdl_type)

    def _fit(self, dataset):
        jmodel = callBigDlFunc(self.bigdl_type, "fitClassifier", self.value, dataset)
        model = DLClassifierModel.of(jmodel, self.featureSize, self.bigdl_type)
        return model

class DLClassifierModel(DLModel):
    """
    .. note:: Deprecated in 0.5.0. `DLClassifierModel` has been migrated to package
     `bigdl.dlframes`. This will be removed in BigDL 0.6.

    """
    def __init__(self,  model, featureSize, jvalue=None, bigdl_type="float"):
        super(DLClassifierModel, self).__init__(model, featureSize, jvalue, bigdl_type)

    def _transform(self, dataset):
        return callBigDlFunc(self.bigdl_type, "dlClassifierModelTransform", self.value, dataset)

    @classmethod
    def of(self, jvalue, feature_size=None, bigdl_type="float"):
        model = DLClassifierModel(model=None, featureSize=feature_size, jvalue=jvalue, bigdl_type=bigdl_type)
        return model


#
# Copyright 2016 The BigDL Authors.
#
# Licensed 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.
#

import os
import sys
import gc
from tempfile import NamedTemporaryFile

from pyspark.cloudpickle import print_exec
from pyspark.broadcast import Broadcast
from pyspark.broadcast import _from_id
from bigdl.nn.layer import Model

def _from_id_and_type(bid, bigdl_type):
    result = _from_id(bid)
    return ModelBroadcast(path=result._path, bigdl_type=bigdl_type)

def broadcast_model(sc, layer):
    return ModelBroadcast(sc, layer, sc._pickled_broadcast_vars)

class ModelBroadcast(Broadcast):

    def __init__(self, sc=None, layer=None, pickle_registry=None, path=None, bigdl_type="float"):
        """
        Should not be called directly by users -- use L{SparkContext.broadcast()}
        instead.
        """
        if layer is not None:
            self.bigdl_type = layer.bigdl_type
        else:
            self.bigdl_type = bigdl_type
        super(ModelBroadcast, self).__init__(sc, layer, pickle_registry, path)

    def dump(self, value, f):
        try:
            value.saveModel(f.name, over_write=True)
        except Exception as e:
            msg = "Could not serialize broadcast: %s" % e.__class__.__name__
            print_exec(sys.stderr)
            raise ValueError(msg)
        f.close()
        return f.name

    def _load(self, path):
        return Model.loadModel(path, bigdl_type=self.bigdl_type)

    @property
    def value(self):
        """ Return the broadcasted value
        """
        if not hasattr(self, "_value") and self._path is not None:
            self._value = self._load(self._path)
        return self._value

    def __reduce__(self):
        if self._jbroadcast is None:
            raise Exception("Broadcast can only be serialized in driver")
        self._pickle_registry.add(self)
        return _from_id_and_type, (self._jbroadcast.id(), self.bigdl_type)


#
# Copyright 2016 The BigDL Authors.
#
# Licensed 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.
#

import os
import sys
import six
from py4j.protocol import Py4JJavaError
from py4j.java_gateway import JavaObject
from py4j.java_collections import ListConverter, JavaArray, JavaList, JavaMap, MapConverter
from py4j.java_gateway import JavaGateway, GatewayClient

from pyspark import RDD, SparkContext
from pyspark.serializers import PickleSerializer, AutoBatchedSerializer
from pyspark.sql import DataFrame, SQLContext
from pyspark.mllib.common import callJavaFunc
from pyspark import SparkConf
from pyspark.files import SparkFiles
import numpy as np
import threading
import tempfile
import traceback
from bigdl.util.engine import get_bigdl_classpath, is_spark_below_2_2

INTMAX = 2147483647
INTMIN = -2147483648
DOUBLEMAX = 1.7976931348623157E308

if sys.version >= '3':
    long = int
    unicode = str

class SingletonMixin(object):
    _lock = threading.RLock()
    _instance = None

    @classmethod
    def instance(cls,
                 bigdl_type, *args):
        if not cls._instance:
            with cls._lock:
                if not cls._instance:
                    cls._instance = cls(bigdl_type, *args)
        return cls._instance

class GatewayWrapper(SingletonMixin):

    def __init__(self, bigdl_type, port=25333):
        self.value = JavaGateway(GatewayClient(port=port), auto_convert=True)

class JavaCreator(SingletonMixin):
    __creator_class=["com.intel.analytics.bigdl.python.api.PythonBigDLKeras"]

    @classmethod
    def add_creator_class(cls, jinvoker):
        with JavaCreator._lock:
            JavaCreator.__creator_class.append(jinvoker)
            JavaCreator._instance = None

    @classmethod
    def get_creator_class(cls):
        with JavaCreator._lock:
            return JavaCreator.__creator_class

    @classmethod
    def set_creator_class(cls, cclass):
        if isinstance(cclass, six.string_types):
            cclass = [cclass]
        with JavaCreator._lock:
            JavaCreator.__creator_class = cclass
            JavaCreator._instance = None

    def __init__(self, bigdl_type, gateway):
        self.value = []
        for creator_class in JavaCreator.get_creator_class():
            jclass = getattr(gateway.jvm, creator_class)
            if bigdl_type == "float":
                self.value.append(getattr(jclass, "ofFloat")())
            elif bigdl_type == "double":
                self.value.append(getattr(jclass, "ofDouble")())
            else:
                raise Exception("Not supported bigdl_type: %s" % bigdl_type)

class JavaValue(object):
    def jvm_class_constructor(self):
        name = "create" + self.__class__.__name__
        print("creating: " + name)
        return name

    def __init__(self, jvalue, bigdl_type, *args):
        self.value = jvalue if jvalue else callBigDlFunc(
            bigdl_type, self.jvm_class_constructor(), *args)
        self.bigdl_type = bigdl_type

    def __str__(self):
        return self.value.toString()

class EvaluatedResult():
    """
    A testing result used to benchmark the model quality.
    """
    def __init__(self, result, total_num, method):
        """

        :param result: the validation result. i.e: top1 accuracy percentage.
        :param total_num: the total processed records.
        :param method: the validation method. i.e: Top1Accuracy
        """
        self.result = result
        self.total_num = total_num
        self.method = method

    def __reduce__(self):
        return (EvaluatedResult, (self.result, self.total_num, self.method))

    def __str__(self):
        return "Evaluated result: %s, total_num: %s, method: %s" % (
            self.result, self.total_num, self.method)

def get_dtype(bigdl_type):
    # Always return float32 for now
    return "float32"

class JActivity(object):

    def __init__(self, value):
        self.value = value

class JTensor(object):
    """
    A wrapper to easy our work when need to pass or return Tensor to/from Scala.

    >>> import numpy as np
    >>> from bigdl.util.common import JTensor
    >>> np.random.seed(123)
    >>>
    """
    def __init__(self, storage, shape, bigdl_type="float", indices=None):
        """

        :param storage: values in this tensor
        :param shape: shape of this tensor
        :param bigdl_type: numeric type
        :param indices: if indices is provided, means this is a SparseTensor;
                        if not provided, means this is a DenseTensor
        """
        if isinstance(storage, bytes) and isinstance(shape, bytes):
            self.storage = np.frombuffer(storage, dtype=get_dtype(bigdl_type))
            self.shape = np.frombuffer(shape, dtype=np.int32)
        else:
            self.storage = np.array(storage, dtype=get_dtype(bigdl_type))
            self.shape = np.array(shape, dtype=np.int32)
        if indices is None:
            self.indices = None
        elif isinstance(indices, bytes):
            self.indices = np.frombuffer(indices, dtype=np.int32)
        else:
            assert isinstance(indices, np.ndarray), \
            "indices should be a np.ndarray, not %s, %s" % (type(a_ndarray), str(indices))
            self.indices = np.array(indices, dtype=np.int32)
        self.bigdl_type = bigdl_type

    @classmethod
    def from_ndarray(cls, a_ndarray, bigdl_type="float"):
        """
        Convert a ndarray to a DenseTensor which would be used in Java side.

        >>> import numpy as np
        >>> from bigdl.util.common import JTensor
        >>> from bigdl.util.common import callBigDlFunc
        >>> np.random.seed(123)
        >>> data = np.random.uniform(0, 1, (2, 3)).astype("float32")
        >>> result = JTensor.from_ndarray(data)
        >>> expected_storage = np.array([[0.69646919, 0.28613934, 0.22685145], [0.55131477, 0.71946895, 0.42310646]])
        >>> expected_shape = np.array([2, 3])
        >>> np.testing.assert_allclose(result.storage, expected_storage, rtol=1e-6, atol=1e-6)
        >>> np.testing.assert_allclose(result.shape, expected_shape)
        >>> data_back = result.to_ndarray()
        >>> (data == data_back).all()
        True
        >>> tensor1 = callBigDlFunc("float", "testTensor", JTensor.from_ndarray(data))  # noqa
        >>> array_from_tensor = tensor1.to_ndarray()
        >>> (array_from_tensor == data).all()
        True
        """
        if a_ndarray is None:
            return None
        assert isinstance(a_ndarray, np.ndarray), \
            "input should be a np.ndarray, not %s" % type(a_ndarray)
        return cls(a_ndarray,
                   a_ndarray.shape if a_ndarray.shape else (a_ndarray.size),
                   bigdl_type)

    @classmethod
    def sparse(cls, a_ndarray, i_ndarray, shape, bigdl_type="float"):
        """
        Convert a three ndarray to SparseTensor which would be used in Java side.
        For example:
        a_ndarray = [1, 3, 2, 4]
        i_ndarray = [[0, 0, 1, 2],
                     [0, 3, 2, 1]]
        shape = [3, 4]
        Present a dense tensor
        [[ 1,  0,  0,  3],
         [ 0,  0,  2,  0],
         [ 0,  4,  0,  0]]

        :param a_ndarray non-zero elements in this SparseTensor
        :param i_ndarray zero-based indices for non-zero element
                         i_ndarray's shape should be (shape.size, a_ndarray.size)
                         And the i-th non-zero elements indices is i_ndarray[:, 1],
                         should be zero-based and ascending;
        :param shape     shape as a DenseTensor.

        >>> import numpy as np
        >>> from bigdl.util.common import JTensor
        >>> from bigdl.util.common import callBigDlFunc
        >>> np.random.seed(123)
        >>> data = np.arange(1, 7).astype("float32")
        >>> indices = np.arange(1, 7)
        >>> shape = np.array([10])
        >>> result = JTensor.sparse(data, indices, shape)
        >>> expected_storage = np.array([1., 2., 3., 4., 5., 6.])
        >>> expected_shape = np.array([10])
        >>> expected_indices = np.array([1, 2, 3, 4, 5, 6])
        >>> np.testing.assert_allclose(result.storage, expected_storage)
        >>> np.testing.assert_allclose(result.shape, expected_shape)
        >>> np.testing.assert_allclose(result.indices, expected_indices)
        >>> tensor1 = callBigDlFunc("float", "testTensor", result)  # noqa
        >>> array_from_tensor = tensor1.to_ndarray()
        >>> expected_ndarray = np.array([0, 1, 2, 3, 4, 5, 6, 0, 0, 0])
        >>> (array_from_tensor == expected_ndarray).all()
        True
        """
        if a_ndarray is None:
            return None
        assert isinstance(a_ndarray, np.ndarray), \
            "values array should be a np.ndarray, not %s" % type(a_ndarray)
        assert isinstance(i_ndarray, np.ndarray), \
            "indices array should be a np.ndarray, not %s" % type(a_ndarray)
        assert i_ndarray.size == a_ndarray.size * shape.size, \
            "size of values and indices should match."
        return cls(a_ndarray,
                   shape,
                   bigdl_type,
                   i_ndarray)

    def to_ndarray(self):
        """
        Transfer JTensor to ndarray.
        As SparseTensor may generate an very big ndarray, so we don't support this function for SparseTensor.
        :return: a ndarray
        """
        assert self.indices is None, "sparseTensor to ndarray is not supported"
        return np.array(self.storage, dtype=get_dtype(self.bigdl_type)).reshape(self.shape)  # noqa

    def __reduce__(self):
        if self.indices is None:
            return JTensor, (self.storage.tostring(), self.shape.tostring(), self.bigdl_type)
        else:
            return JTensor, (self.storage.tostring(), self.shape.tostring(), self.bigdl_type, self.indices.tostring())

    def __str__(self):
        return self.__repr__()

    def __repr__(self):
        indices = "" if self.indices is None else " ,indices %s" % str(self.indices)
        return "JTensor: storage: %s, shape: %s%s, %s" % (str(self.storage), str(self.shape), indices, self.bigdl_type)

class Sample(object):
    def __init__(self, features, labels, bigdl_type="float"):
        """
        User should always use Sample.from_ndarray to construct Sample.
        :param features: a list of JTensors
        :param labels: a list of JTensors
        :param bigdl_type: "double" or "float"
        """
        self.feature = features[0]
        self.features = features
        self.label = labels[0]
        self.bigdl_type = bigdl_type
        self.labels = labels

    @classmethod
    def from_ndarray(cls, features, labels, bigdl_type="float"):
        """
        Convert a ndarray of features and labels to Sample, which would be used in Java side.
        :param features: an ndarray or a list of ndarrays
        :param labels: an ndarray or a list of ndarrays or a scalar
        :param bigdl_type: "double" or "float"

        >>> import numpy as np
        >>> from bigdl.util.common import callBigDlFunc
        >>> from numpy.testing import assert_allclose
        >>> np.random.seed(123)
        >>> sample = Sample.from_ndarray(np.random.random((2,3)), np.random.random((2,3)))
        >>> sample_back = callBigDlFunc("float", "testSample", sample)
        >>> assert_allclose(sample.features[0].to_ndarray(), sample_back.features[0].to_ndarray())
        >>> assert_allclose(sample.label.to_ndarray(), sample_back.label.to_ndarray())
        >>> expected_feature_storage = np.array(([[0.69646919, 0.28613934, 0.22685145], [0.55131477, 0.71946895, 0.42310646]]))
        >>> expected_feature_shape = np.array([2, 3])
        >>> expected_label_storage = np.array(([[0.98076421, 0.68482971, 0.48093191], [0.39211753, 0.343178, 0.72904968]]))
        >>> expected_label_shape = np.array([2, 3])
        >>> assert_allclose(sample.features[0].storage, expected_feature_storage, rtol=1e-6, atol=1e-6)
        >>> assert_allclose(sample.features[0].shape, expected_feature_shape)
        >>> assert_allclose(sample.labels[0].storage, expected_label_storage, rtol=1e-6, atol=1e-6)
        >>> assert_allclose(sample.labels[0].shape, expected_label_shape)
        """
        if isinstance(features, np.ndarray):
            features = [features]
        else:
            assert all(isinstance(feature, np.ndarray) for feature in features), \
                "features should be a list of np.ndarray, not %s" % type(features)
        if np.isscalar(labels):  # in case labels is a scalar.
            labels = [np.array(labels)]
        elif isinstance(labels, np.ndarray):
            labels = [labels]
        else:
            assert all(isinstance(label, np.ndarray) for label in labels), \
                "labels should be a list of np.ndarray, not %s" % type(labels)
        return cls(
            features=[JTensor.from_ndarray(feature) for feature in features],
            labels=[JTensor.from_ndarray(label) for label in labels],
            bigdl_type=bigdl_type)

    @classmethod
    def from_jtensor(cls, features, labels, bigdl_type="float"):
        """
        Convert a sequence of JTensor to Sample, which would be used in Java side.
        :param features: an JTensor or a list of JTensor
        :param labels: an JTensor or a list of JTensor or a scalar
        :param bigdl_type: "double" or "float"

        >>> import numpy as np
        >>> data = np.random.uniform(0, 1, (6)).astype("float32")
        >>> indices = np.arange(1, 7)
        >>> shape = np.array([10])
        >>> feature0 = JTensor.sparse(data, indices, shape)
        >>> feature1 = JTensor.from_ndarray(np.random.uniform(0, 1, (2, 3)).astype("float32"))
        >>> sample = Sample.from_jtensor([feature0, feature1], 1)
        """
        if isinstance(features, JTensor):
            features = [features]
        else:
            assert all(isinstance(feature, JTensor) for feature in features), \
                "features should be a list of JTensor, not %s" % type(features)
        if np.isscalar(labels):  # in case labels is a scalar.
            labels = [JTensor.from_ndarray(np.array(labels))]
        elif isinstance(labels, JTensor):
            labels = [labels]
        else:
            assert all(isinstance(label, JTensor) for label in labels), \
                "labels should be a list of np.ndarray, not %s" % type(labels)
        return cls(
            features=features,
            labels=labels,
            bigdl_type=bigdl_type)

    def __reduce__(self):
        return Sample, (self.features, self.labels, self.bigdl_type)

    def __str__(self):
        return "Sample: features: %s, labels: %s," % (self.features, self.labels)

    def __repr__(self):
        return "Sample: features: %s, labels: %s" % (self.features, self.labels)

class RNG():
    """
    generate tensor data with seed
    """
    def __init__(self, bigdl_type="float"):
        self.bigdl_type = bigdl_type

    def set_seed(self, seed):
        callBigDlFunc(self.bigdl_type, "setModelSeed", seed)

    def uniform(self, a, b, size):
        return callBigDlFunc(self.bigdl_type, "uniform", a, b, size).to_ndarray() # noqa

_picklable_classes = [
    'LinkedList',
    'SparseVector',
    'DenseVector',
    'DenseMatrix',
    'Rating',
    'LabeledPoint',
    'Sample',
    'EvaluatedResult',
    'JTensor',
    'JActivity'
]

def init_engine(bigdl_type="float"):
    callBigDlFunc(bigdl_type, "initEngine")
    # Spark context is supposed to have been created when init_engine is called
    get_spark_context()._jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.initialize()

def init_executor_gateway(sc, bigdl_type="float"):
    callBigDlFunc(bigdl_type, "initExecutorGateway", sc, sc._gateway._gateway_client.port)

def get_node_and_core_number(bigdl_type="float"):
    result = callBigDlFunc(bigdl_type, "getNodeAndCoreNumber")
    return result[0], result[1]

def redire_spark_logs(bigdl_type="float", log_path=os.getcwd()+"/bigdl.log"):
    """
    Redirect spark logs to the specified path.
    :param bigdl_type: "double" or "float"
    :param log_path: the file path to be redirected to; the default file is under the current workspace named `bigdl.log`.
    """
    callBigDlFunc(bigdl_type, "redirectSparkLogs", log_path)

def show_bigdl_info_logs(bigdl_type="float"):
    """
    Set BigDL log level to INFO.
    :param bigdl_type: "double" or "float"
    """
    callBigDlFunc(bigdl_type, "showBigDlInfoLogs")

def get_bigdl_conf():
    bigdl_conf_file = "spark-bigdl.conf"
    bigdl_python_wrapper = "python-api.zip"

    def load_conf(conf_str):
        return dict(line.split() for line in conf_str.split("\n") if
                    "#" not in line and line.strip())

    for p in sys.path:
        if bigdl_conf_file in p and os.path.isfile(p):
            with open(p) if sys.version_info < (3,) else open(p, encoding='latin-1') as conf_file: # noqa
                return load_conf(conf_file.read())
        if bigdl_python_wrapper in p and os.path.isfile(p):
            import zipfile
            with zipfile.ZipFile(p, 'r') as zip_conf:
                if bigdl_conf_file  in zip_conf.namelist():
                    content = zip_conf.read(bigdl_conf_file)
                    if sys.version_info >= (3,):
                        content = str(content, 'latin-1')
                    return load_conf(content)
    return {}

def to_list(a):
    if type(a) is list:
        return a
    return [a]

def to_sample_rdd(x, y, numSlices=None):
    """
    Conver x and y into RDD[Sample]
    :param x: ndarray and the first dimension should be batch
    :param y: ndarray and the first dimension should be batch
    :param numSlices:
    :return:
    """
    sc = get_spark_context()
    from bigdl.util.common import Sample
    x_rdd = sc.parallelize(x, numSlices)
    y_rdd = sc.parallelize(y, numSlices)
    return x_rdd.zip(y_rdd).map(lambda item: Sample.from_ndarray(item[0], item[1]))

def extend_spark_driver_cp(sparkConf, path):
    original_driver_classpath = ":" + sparkConf.get("spark.driver.extraClassPath") \
        if sparkConf.contains("spark.driver.extraClassPath") else ""
    sparkConf.set("spark.driver.extraClassPath", path + original_driver_classpath)

def create_spark_conf():
    bigdl_conf = get_bigdl_conf()
    sparkConf = SparkConf()
    sparkConf.setAll(bigdl_conf.items())
    if os.environ.get("BIGDL_JARS", None) and not is_spark_below_2_2():
        for jar in os.environ["BIGDL_JARS"].split(":"):
            extend_spark_driver_cp(sparkConf, jar)

    # add content in PYSPARK_FILES in spark.submit.pyFiles
    # This is a workaround for current Spark on k8s
    python_lib = os.environ.get('PYSPARK_FILES', None)
    if python_lib:
        existing_py_files = sparkConf.get("spark.submit.pyFiles")
        if existing_py_files:
            sparkConf.set(key="spark.submit.pyFiles", value="%s,%s" % (python_lib, existing_py_files))
        else:
            sparkConf.set(key="spark.submit.pyFiles", value=python_lib)

    return sparkConf

def get_spark_context(conf=None):
    """
    Get the current active spark context and create one if no active instance
    :param conf: combining bigdl configs into spark conf
    :return: SparkContext
    """
    if hasattr(SparkContext, "getOrCreate"):
        with SparkContext._lock:
            if SparkContext._active_spark_context is None:
                spark_conf = create_spark_conf() if conf is None else conf
                return SparkContext.getOrCreate(spark_conf)
            else:
                return SparkContext.getOrCreate()

    else:
        # Might have threading issue but we cann't add _lock here
        # as it's not RLock in spark1.5;
        if SparkContext._active_spark_context is None:
            spark_conf = create_spark_conf() if conf is None else conf
            return SparkContext(conf=spark_conf)
        else:
            return SparkContext._active_spark_context

def get_spark_sql_context(sc):
    if "getOrCreate" in SQLContext.__dict__:
        return SQLContext.getOrCreate(sc)
    else:
        return SQLContext(sc)  # Compatible with Spark1.5.1

def _get_port():
    root_dir = SparkFiles.getRootDirectory()
    path = os.path.join(root_dir, "gateway_port")
    try:
        with open(path) as f:
            port = int(f.readline())
    except IOError as e:
        traceback.print_exc()
        raise RuntimeError("Could not open the file %s, which contains the listening port of"
                           " local Java Gateway, please make sure the init_executor_gateway()"
                           " function is called before any call of java function on the"
                           " executor side." % e.filename)
    return port

def _get_gateway():
    if SparkFiles._is_running_on_worker:
        gateway_port = _get_port()
        gateway = GatewayWrapper.instance(None, gateway_port).value
    else:
        sc = get_spark_context()
        gateway = sc._gateway
    return gateway

def callBigDlFunc(bigdl_type, name, *args):
    """ Call API in PythonBigDL """
    gateway = _get_gateway()
    error = Exception("Cannot find function: %s" % name)
    for jinvoker in JavaCreator.instance(bigdl_type, gateway).value:
        # hasattr(jinvoker, name) always return true here,
        # so you need to invoke the method to check if it exist or not
        try:
            api = getattr(jinvoker, name)
            result = callJavaFunc(api, *args)
        except Exception as e:
            error = e
            if "does not exist" not in str(e):
                raise e
        else:
            return result
    raise error

def _java2py(gateway, r, encoding="bytes"):
    if isinstance(r, JavaObject):
        clsName = r.getClass().getSimpleName()
        # convert RDD into JavaRDD
        if clsName != 'JavaRDD' and clsName.endswith("RDD"):
            r = r.toJavaRDD()
            clsName = 'JavaRDD'

        if clsName == 'JavaRDD':
            jrdd = gateway.jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.javaToPython(r)
            return RDD(jrdd, get_spark_context())

        if clsName == 'DataFrame':
            return DataFrame(r, get_spark_sql_context(get_spark_context()))

        if clsName == 'Dataset':
            return DataFrame(r, get_spark_sql_context(get_spark_context()))

        if clsName == "ImageFrame[]":
            return r

        if clsName in _picklable_classes:
            r = gateway.jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.dumps(r)
        elif isinstance(r, (JavaArray, JavaList, JavaMap)):
            try:
                r = gateway.jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.dumps(
                    r)
            except Py4JJavaError:
                pass  # not pickable

    if isinstance(r, (bytearray, bytes)):
        r = PickleSerializer().loads(bytes(r), encoding=encoding)
    return r

def callJavaFunc(func, *args):
    """ Call Java Function """
    gateway = _get_gateway()
    args = [_py2java(gateway, a) for a in args]
    result = func(*args)
    return _java2py(gateway, result)

def _to_java_object_rdd(rdd):
    """ Return a JavaRDD of Object by unpickling

    It will convert each Python object into Java object by Pyrolite, whenever
    the RDD is serialized in batch or not.
    """
    rdd = rdd._reserialize(AutoBatchedSerializer(PickleSerializer()))
    return \
        rdd.ctx._jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.pythonToJava(
            rdd._jrdd, True)

def _py2java(gateway, obj):
    """ Convert Python object into Java """
    if isinstance(obj, RDD):
        obj = _to_java_object_rdd(obj)
    elif isinstance(obj, DataFrame):
        obj = obj._jdf
    elif isinstance(obj, SparkContext):
        obj = obj._jsc
    elif isinstance(obj, (list, tuple)):
        obj = ListConverter().convert([_py2java(gateway, x) for x in obj],
                                      gateway._gateway_client)
    elif isinstance(obj, dict):
        result = {}
        for (key, value) in obj.items():
            result[key] = _py2java(gateway, value)
        obj = MapConverter().convert(result, gateway._gateway_client)
    elif isinstance(obj, JavaValue):
        obj = obj.value
    elif isinstance(obj, JavaObject):
        pass
    elif isinstance(obj, (int, long, float, bool, bytes, unicode)):
        pass
    else:
        data = bytearray(PickleSerializer().dumps(obj))
        obj = gateway.jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.loads(data)
    return obj

def create_tmp_path():
    tmp_file = tempfile.NamedTemporaryFile(prefix="bigdl")
    tmp_file.close()
    return tmp_file.name

def text_from_path(path):
    sc = get_spark_context()
    return sc.textFile(path).collect()[0]

def get_local_file(a_path):
    if not is_distributed(a_path):
        return a_path
    path, data = get_spark_context().binaryFiles(a_path).collect()[0]
    local_file_path = create_tmp_path()
    with open(local_file_path, 'w') as local_file:
        local_file.write(data)
    return local_file_path

def is_distributed(path):
    return "://" in path

def get_activation_by_name(activation_name, activation_id=None):
    """ Convert to a bigdl activation layer
        given the name of the activation as a string  """
    import bigdl.nn.layer as BLayer
    activation = None
    activation_name = activation_name.lower()
    if activation_name == "tanh":
        activation = BLayer.Tanh()
    elif activation_name == "sigmoid":
        activation = BLayer.Sigmoid()
    elif activation_name == "hard_sigmoid":
        activation = BLayer.HardSigmoid()
    elif activation_name == "relu":
        activation = BLayer.ReLU()
    elif activation_name == "softmax":
        activation = BLayer.SoftMax()
    elif activation_name == "softplus":
        activation = BLayer.SoftPlus(beta=1.0)
    elif activation_name == "softsign":
        activation = BLayer.SoftSign()
    elif activation_name == "linear":
        activation = BLayer.Identity()
    else:
        raise Exception("Unsupported activation type: %s" % activation_name)
    if not activation_id:
        activation.set_name(activation_id)
    return activation

def _test():
    import doctest
    from pyspark import SparkContext
    from bigdl.nn import layer
    globs = layer.__dict__.copy()
    sc = SparkContext(master="local[2]", appName="test common utility")
    globs['sc'] = sc
    (failure_count, test_count) = doctest.testmod(globs=globs,
                                                  optionflags=doctest.ELLIPSIS)
    if failure_count:
        exit(-1)

if __name__ == "__main__":
    _test()


#
# Copyright 2016 The BigDL Authors.
#
# Licensed 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.
#

import sys
from bigdl.util.common import JavaValue
from bigdl.util.common import callBigDlFunc
from bigdl.util.common import *

if sys.version >= '3':
    long = int
    unicode = str

class FeatureTransformer(JavaValue):
    """
    FeatureTransformer is a transformer that transform ImageFeature
    """

    def __init__(self, bigdl_type="float", *args):
        self.value = callBigDlFunc(
                bigdl_type, JavaValue.jvm_class_constructor(self), *args)

    def transform(self, image_feature, bigdl_type="float"):
        """
        transform ImageFeature
        """
        callBigDlFunc(bigdl_type, "transformImageFeature", self.value, image_feature)
        return image_feature

    def __call__(self, image_frame, bigdl_type="float"):
        """
        transform ImageFrame
        """
        jframe = callBigDlFunc(bigdl_type,
                             "transformImageFrame", self.value, image_frame)
        return ImageFrame(jvalue=jframe)

class Pipeline(FeatureTransformer):
    """
    Pipeline of FeatureTransformer
    """

    def __init__(self, transformers, bigdl_type="float"):
        for transfomer in transformers:
            assert transfomer.__class__.__bases__[0].__name__ == "FeatureTransformer", "the transformer should be " \
                                                                                "subclass of FeatureTransformer"
        super(Pipeline, self).__init__(bigdl_type, transformers)

class ImageFeature(JavaValue):
    """
    Each ImageFeature keeps information about single image,
    it can include various status of an image,
    e.g. original bytes read from image file, an opencv mat,
    pixels in float array, image label, meta data and so on.
    it uses HashMap to store all these data,
    the key is string that identify the corresponding value
    """

    def __init__(self, image=None, label=None, path=None, bigdl_type="float"):
        image_tensor = JTensor.from_ndarray(image) if image is not None else None
        label_tensor = JTensor.from_ndarray(label) if label is not None else None
        self.bigdl_type = bigdl_type
        self.value = callBigDlFunc(
            bigdl_type, JavaValue.jvm_class_constructor(self), image_tensor, label_tensor, path)

    def get_image(self, float_key="floats", to_chw=True):
        """
        get image as ndarray from ImageFeature
        """
        tensor = callBigDlFunc(self.bigdl_type, "imageFeatureToImageTensor", self.value,
                               float_key, to_chw)
        return tensor.to_ndarray()

    def get_label(self):
        """
        get label as ndarray from ImageFeature
        """
        label = callBigDlFunc(self.bigdl_type, "imageFeatureToLabelTensor", self.value)
        return label.to_ndarray()

    def keys(self):
        """
        get key set from ImageFeature
        """
        return callBigDlFunc(self.bigdl_type, "imageFeatureGetKeys", self.value)

class ImageFrame(JavaValue):
    """
    ImageFrame wraps a set of ImageFeature
    """

    def __init__(self, jvalue, bigdl_type="float"):
        self.value = jvalue
        self.bigdl_type = bigdl_type
        if self.is_local():
            self.image_frame = LocalImageFrame(jvalue=self.value)
        else:
            self.image_frame = DistributedImageFrame(jvalue=self.value)

    @classmethod
    def read(cls, path, sc=None, min_partitions=1, bigdl_type="float"):
        """
        Read images as Image Frame
        if sc is defined, Read image as DistributedImageFrame from local file system or HDFS
        if sc is null, Read image as LocalImageFrame from local file system
        :param path path to read images
        if sc is defined, path can be local or HDFS. Wildcard character are supported.
        if sc is null, path is local directory/image file/image file with wildcard character
        :param sc SparkContext
        :param min_partitions A suggestion value of the minimal splitting number for input data.
        :return ImageFrame
        """
        return ImageFrame(jvalue=callBigDlFunc(bigdl_type, "read", path, sc, min_partitions))

    @classmethod
    def read_parquet(cls, path, sc, bigdl_type="float"):
        """
        Read parquet file as DistributedImageFrame
        """
        return DistributedImageFrame(jvalue=callBigDlFunc(bigdl_type, "readParquet", path, sc))

    @classmethod
    def write_parquet(cls, path, output, sc, partition_num = 1, bigdl_type="float"):
        """
        write ImageFrame as parquet file
        """
        return callBigDlFunc(bigdl_type, "writeParquet", path, output, sc, partition_num)

    def is_local(self):
        """
        whether this is a LocalImageFrame
        """
        return callBigDlFunc(self.bigdl_type, "isLocal", self.value)

    def is_distributed(self):
        """
        whether this is a DistributedImageFrame
        """
        return callBigDlFunc(self.bigdl_type, "isDistributed", self.value)

    def transform(self, transformer, bigdl_type="float"):
        """
        transformImageFrame
        """
        self.value = callBigDlFunc(bigdl_type,
                                 "transformImageFrame", transformer, self.value)
        return self

    def get_image(self, float_key="floats", to_chw=True):
        """
        get image from ImageFrame
        """
        return self.image_frame.get_image(float_key, to_chw)

    def get_label(self):
        """
        get label from ImageFrame
        """
        return self.image_frame.get_label()

    def get_predict(self, key="predict"):
        """
        get prediction from ImageFrame
        """
        return self.image_frame.get_predict(key)

    def get_sample(self):
        """
        get sample from ImageFrame
        """
        return self.image_frame.get_sample()

    def get_uri(self):
        """
        get uri from imageframe
        """
        return self.image_frame.get_uri()

    def set_label(self, label, bigdl_type="float"):
        """
        set label for imageframe
        """
        return callBigDlFunc(bigdl_type,
                             "setLabel", label, self.value)

    def random_split(self, weights):
        """
        Random split imageframes according to weights
        :param weights: weights for each ImageFrame
        :return:
        """
        jvalues =  self.image_frame.random_split(weights)
        return [ImageFrame(jvalue) for jvalue in jvalues]

class LocalImageFrame(ImageFrame):
    """
    LocalImageFrame wraps a list of ImageFeature
    """
    def __init__(self, image_list=None, label_list=None, jvalue=None, bigdl_type="float"):
        assert jvalue or image_list, "jvalue and image_list cannot be None in the same time"
        if jvalue:
            self.value = jvalue
        else:
            # init from image ndarray list and label rdd(optional)
            image_tensor_list = map(lambda image: JTensor.from_ndarray(image), image_list)
            label_tensor_list = map(lambda label: JTensor.from_ndarray(label), label_list) if label_list else None
            self.value = callBigDlFunc(bigdl_type, JavaValue.jvm_class_constructor(self),
                                       image_tensor_list, label_tensor_list)

        self.bigdl_type = bigdl_type

    def get_image(self, float_key="floats", to_chw=True):
        """
        get image list from ImageFrame
        """
        tensors = callBigDlFunc(self.bigdl_type,
                                   "localImageFrameToImageTensor", self.value, float_key, to_chw)
        return map(lambda tensor: tensor.to_ndarray(), tensors)

    def get_label(self):
        """
        get label list from ImageFrame
        """
        labels = callBigDlFunc(self.bigdl_type, "localImageFrameToLabelTensor", self.value)
        return map(lambda tensor: tensor.to_ndarray(), labels)

    def get_predict(self, key="predict"):
        """
        get prediction list from ImageFrame
        """
        predicts = callBigDlFunc(self.bigdl_type, "localImageFrameToPredict", self.value, key)
        return map(lambda predict: (predict[0], predict[1].to_ndarray()) if predict[1] else (predict[0], None), predicts)

    def get_sample(self,  key="sample"):
        return callBigDlFunc(self.bigdl_type, "localImageFrameToSample", self.value, key)

    def get_uri(self, key = "uri"):
        return callBigDlFunc(self.bigdl_type, "localImageFrameToUri", self.value, key)

    def random_split(self, weights):
        raise "random split not supported in LocalImageFrame"

class DistributedImageFrame(ImageFrame):
    """
    DistributedImageFrame wraps an RDD of ImageFeature
    """

    def __init__(self, image_rdd=None, label_rdd=None, jvalue=None, bigdl_type="float"):
        assert jvalue or image_rdd, "jvalue and image_rdd cannot be None in the same time"
        if jvalue:
            self.value = jvalue
        else:
            # init from image ndarray rdd and label rdd(optional)
            image_tensor_rdd = image_rdd.map(lambda image: JTensor.from_ndarray(image))
            label_tensor_rdd = label_rdd.map(lambda label: JTensor.from_ndarray(label)) if label_rdd else None
            self.value = callBigDlFunc(bigdl_type, JavaValue.jvm_class_constructor(self),
                                       image_tensor_rdd, label_tensor_rdd)

        self.bigdl_type = bigdl_type

    def get_image(self, float_key="floats", to_chw=True):
        """
        get image rdd from ImageFrame
        """
        tensor_rdd = callBigDlFunc(self.bigdl_type,
                               "distributedImageFrameToImageTensorRdd", self.value, float_key, to_chw)
        return tensor_rdd.map(lambda tensor: tensor.to_ndarray())

    def get_label(self):
        """
        get label rdd from ImageFrame
        """
        tensor_rdd = callBigDlFunc(self.bigdl_type, "distributedImageFrameToLabelTensorRdd", self.value)
        return tensor_rdd.map(lambda tensor: tensor.to_ndarray())

    def get_predict(self, key="predict"):
        """
        get prediction rdd from ImageFrame
        """
        predicts = callBigDlFunc(self.bigdl_type, "distributedImageFrameToPredict", self.value, key)
        return predicts.map(lambda predict: (predict[0], predict[1].to_ndarray()) if predict[1] else (predict[0], None))

    def get_sample(self,  key="sample"):
        return callBigDlFunc(self.bigdl_type, "distributedImageFrameToSample", self.value, key)

    def get_uri(self, key = "uri"):
        return callBigDlFunc(self.bigdl_type, "distributedImageFrameToUri", self.value, key)

    def random_split(self, weights):
        return callBigDlFunc(self.bigdl_type, "distributedImageFrameRandomSplit", self.value, weights)

class HFlip(FeatureTransformer):
    """
    Flip the image horizontally
    """

    def __init__(self, bigdl_type="float"):
            super(HFlip, self).__init__(bigdl_type)

class Resize(FeatureTransformer):
    """
    Resize image
    :param resize_h height after resize
    :param resize_w width after resize
    :param resize_mode if resizeMode = -1, random select a mode from (Imgproc.INTER_LINEAR,
     Imgproc.INTER_CUBIC, Imgproc.INTER_AREA, Imgproc.INTER_NEAREST, Imgproc.INTER_LANCZOS4)
    :param use_scale_factor if true, scale factor fx and fy is used, fx = fy = 0
    note that the result of the following are different
    Imgproc.resize(mat, mat, new Size(resizeWH, resizeWH), 0, 0, Imgproc.INTER_LINEAR)
    Imgproc.resize(mat, mat, new Size(resizeWH, resizeWH))
    """

    def __init__(self, resize_h, resize_w, resize_mode = 1, use_scale_factor=True,
                 bigdl_type="float"):
        super(Resize, self).__init__(bigdl_type, resize_h, resize_w, resize_mode, use_scale_factor)

class Brightness(FeatureTransformer):
    """
    adjust the image brightness
    :param deltaLow brightness parameter: low bound
    :param deltaHigh brightness parameter: high bound
    """

    def __init__(self, delta_low, delta_high, bigdl_type="float"):
        super(Brightness, self).__init__(bigdl_type, delta_low, delta_high)

class ChannelOrder(FeatureTransformer):
    """
    random change the channel of an image
    """

    def __init__(self, bigdl_type="float"):
        super(ChannelOrder, self).__init__(bigdl_type)

class Contrast(FeatureTransformer):
    """
    Adjust the image contrast
    :param delta_low contrast parameter low bound
    :param delta_high contrast parameter high bound
    """

    def __init__(self, delta_low, delta_high, bigdl_type="float"):
        super(Contrast, self).__init__(bigdl_type, delta_low, delta_high)

class Saturation(FeatureTransformer):
    """
    Adjust image saturation
    """

    def __init__(self, delta_low, delta_high, bigdl_type="float"):
        super(Saturation, self).__init__(bigdl_type, delta_low, delta_high)

class Hue(FeatureTransformer):
    """
    Adjust image hue
    :param delta_low hue parameter: low bound
    :param delta_high hue parameter: high bound
    """

    def __init__(self, delta_low, delta_high, bigdl_type="float"):
        super(Hue, self).__init__(bigdl_type, delta_low, delta_high)

class ChannelNormalize(FeatureTransformer):
    """
    image channel normalize
    :param mean_r mean value in R channel
    :param mean_g mean value in G channel
    :param meanB_b mean value in B channel
    :param std_r std value in R channel
    :param std_g std value in G channel
    :param std_b std value in B channel
    """
    def __init__(self, mean_r, mean_g, mean_b, std_r=1.0, std_g=1.0, std_b=1.0, bigdl_type="float"):
        super(ChannelNormalize, self).__init__(bigdl_type, mean_r, mean_g, mean_b, std_r, std_g, std_b)

class PixelNormalize(FeatureTransformer):
    """
    Pixel level normalizer, data(i) = data(i) - mean(i)

    :param means pixel level mean, following H * W * C order
    """

    def __init__(self, means, bigdl_type="float"):
        super(PixelNormalize, self).__init__(bigdl_type, means)

class RandomCrop(FeatureTransformer):
    """
    Random crop a `cropWidth` x `cropHeight` patch from an image.
    The patch size should be less than the image size.

    :param crop_width width after crop
    :param crop_height height after crop
    :param is_clip whether to clip the roi to image boundaries
    """

    def __init__(self, crop_width, crop_height, is_clip=True, bigdl_type="float"):
        super(RandomCrop, self).__init__(bigdl_type, crop_width, crop_height, is_clip)

class CenterCrop(FeatureTransformer):
    """
    Crop a `cropWidth` x `cropHeight` patch from center of image.
    The patch size should be less than the image size.
    :param crop_width width after crop
    :param crop_height height after crop
    :param is_clip  clip cropping box boundary
    """

    def __init__(self, crop_width, crop_height, is_clip=True, bigdl_type="float"):
        super(CenterCrop, self).__init__(bigdl_type, crop_width, crop_height, is_clip)

class FixedCrop(FeatureTransformer):
    """
    Crop a fixed area of image

    :param x1 start in width
    :param y1 start in height
    :param x2 end in width
    :param y2 end in height
    :param normalized whether args are normalized, i.e. in range [0, 1]
    :param is_clip whether to clip the roi to image boundaries
    """

    def __init__(self, x1, y1, x2, y2, normalized=True, is_clip=True, bigdl_type="float"):
        super(FixedCrop, self).__init__(bigdl_type, x1, y1, x2, y2, normalized, is_clip)

class DetectionCrop(FeatureTransformer):
    """
    Crop from object detections, each image should has a tensor detection,
    which is stored in ImageFeature
    :param roi_key key that map a tensor detection
    :param normalized whether is detection is normalized, i.e. in range [0, 1]
    """

    def __init__(self, roi_key, normalized=True, bigdl_type="float"):
        super(DetectionCrop, self).__init__(bigdl_type, roi_key, normalized)

class Expand(FeatureTransformer):
    """
    expand image, fill the blank part with the meanR, meanG, meanB

    :param means_r means in R channel
    :param means_g means in G channel
    :param means_b means in B channel
    :param min_expand_ratio min expand ratio
    :param max_expand_ratio max expand ratio
    """

    def __init__(self, means_r=123, means_g=117, means_b=104,
                 min_expand_ratio=1.0,
                 max_expand_ratio=4.0, bigdl_type="float"):
        super(Expand, self).__init__(bigdl_type, means_r, means_g, means_b,
                                     min_expand_ratio, max_expand_ratio)

class Filler(FeatureTransformer):
    """
    Fill part of image with certain pixel value
    :param start_x start x ratio
    :param start_y start y ratio
    :param end_x end x ratio
    :param end_y end y ratio
    :param value filling value
    """

    def __init__(self, start_x, start_y, end_x, end_y, value = 255, bigdl_type="float"):
        super(Filler, self).__init__(bigdl_type, start_x,
                                     start_y,
                                     end_x,
                                     end_y,
                                     value)

class RandomTransformer(FeatureTransformer):
    """
    It is a wrapper for transformers to control the transform probability
    :param transformer transformer to apply randomness
    :param prob max prob
    """

    def __init__(self, transformer, prob, bigdl_type="float"):
        super(RandomTransformer, self).__init__(bigdl_type, transformer, prob)

class ColorJitter(FeatureTransformer):
    """
    Random adjust brightness, contrast, hue, saturation
    :param brightness_prob probability to adjust brightness
    :param brightness_delta brightness parameter
    :param contrast_prob probability to adjust contrast
    :param contrast_lower contrast lower parameter
    :param contrast_upper contrast upper parameter
    :param hue_prob probability to adjust hue
    :param hue_delta hue parameter
    :param saturation_prob probability to adjust saturation
    :param saturation_lower saturation lower parameter
    :param saturation_upper saturation upper parameter
    :param random_order_prob random order for different operation
    :param shuffle  shuffle the transformers
    """
    def __init__(self, brightness_prob = 0.5,
                 brightness_delta = 32.0,
                 contrast_prob = 0.5,
                 contrast_lower = 0.5,
                 contrast_upper = 1.5,
                 hue_prob = 0.5,
                 hue_delta = 18.0,
                 saturation_prob = 0.5,
                 saturation_lower = 0.5,
                 saturation_upper = 1.5,
                 random_order_prob = 0.0,
                 shuffle = False,
                 bigdl_type="float"):
        super(ColorJitter, self).__init__(bigdl_type, brightness_prob,
                                          brightness_delta,
                                          contrast_prob,
                                          contrast_lower,
                                          contrast_upper,
                                          hue_prob,
                                          hue_delta,
                                          saturation_prob,
                                          saturation_lower,
                                          saturation_upper,
                                          random_order_prob,
                                          shuffle)

class RandomSampler(FeatureTransformer):
    """
    Random sample a bounding box given some constraints and crop the image
    This is used in SSD training augmentation
    """

    def __init__(self):
        super(RandomSampler, self).__init__(bigdl_type)

class RoiProject(FeatureTransformer):
    """
    Project gt boxes onto the coordinate system defined by image boundary
    :param need_meet_center_constraint whether need to meet center constraint, i.e., the center of gt box need be within image boundary
    """

    def __init__(self, need_meet_center_constraint, bigdl_type="float"):
        super(RoiProject, self).__init__(bigdl_type, need_meet_center_constraint)

class RoiHFlip(FeatureTransformer):
    """
    horizontally flip the roi
    :param normalized whether the roi is normalized, i.e. in range [0, 1]
    """

    def __init__(self, normalized=True, bigdl_type="float"):
        super(RoiHFlip, self).__init__(bigdl_type, normalized)

class RoiResize(FeatureTransformer):
    """
    resize the roi according to scale
    :param normalized whether the roi is normalized, i.e. in range [0, 1]
    """
    def __init__(self, normalized=True, bigdl_type="float"):
        super(RoiResize, self).__init__(bigdl_type, normalized)

class RoiNormalize(FeatureTransformer):
    """
    Normalize Roi to [0, 1]
    """

    def __init__(self, bigdl_type="float"):
        super(RoiNormalize, self).__init__(bigdl_type)

class MatToFloats(FeatureTransformer):
    """
    Transform OpenCVMat to float array, note that in this transformer, the mat is released
    :param valid_height valid height in case the mat is invalid
    :param valid_width valid width in case the mat is invalid
    :param valid_channel valid channel in case the mat is invalid
    :param out_key key to store float array
    :param share_buffer share buffer of output
    """

    def __init__(self, valid_height=300, valid_width=300, valid_channel=300,
                 out_key = "floats", share_buffer=True, bigdl_type="float"):
        super(MatToFloats, self).__init__(bigdl_type, valid_height, valid_width, valid_channel,
                                          out_key, share_buffer)

class MatToTensor(FeatureTransformer):
    """
    transform opencv mat to tensor
    :param to_rgb BGR to RGB (default is BGR)
    :param tensor_key key to store transformed tensor
    """

    def __init__(self, to_rgb=False, tensor_key="imageTensor", bigdl_type="float"):
        super(MatToTensor, self).__init__(bigdl_type, to_rgb, tensor_key)

class AspectScale(FeatureTransformer):
    """
    Resize the image, keep the aspect ratio. scale according to the short edge
    :param min_size scale size, apply to short edge
    :param scale_multiple_of make the scaled size multiple of some value
    :param max_size max size after scale
    :param resize_mode if resizeMode = -1, random select a mode from
    (Imgproc.INTER_LINEAR, Imgproc.INTER_CUBIC, Imgproc.INTER_AREA,
    Imgproc.INTER_NEAREST, Imgproc.INTER_LANCZOS4)
    :param use_scale_factor if true, scale factor fx and fy is used, fx = fy = 0
    :aram min_scale control the minimum scale up for image
    """

    def __init__(self, min_size, scale_multiple_of = 1, max_size = 1000,
                 resize_mode = 1, use_scale_factor=True, min_scale=-1.0,
                 bigdl_type="float"):
        super(AspectScale, self).__init__(bigdl_type, min_size, scale_multiple_of, max_size,
                                          resize_mode, use_scale_factor, min_scale)

class RandomAspectScale(FeatureTransformer):
    """
    resize the image by randomly choosing a scale
    :param scales array of scale options that for random choice
    :param scaleMultipleOf Resize test images so that its width and height are multiples of
    :param maxSize Max pixel size of the longest side of a scaled input image
    """
    def __init__(self, scales, scale_multiple_of = 1, max_size = 1000, bigdl_type="float"):
        super(RandomAspectScale, self).__init__(bigdl_type, scales, scale_multiple_of, max_size)

class BytesToMat(FeatureTransformer):
    """
    Transform byte array(original image file in byte) to OpenCVMat
    :param byte_key key that maps byte array
    """
    def __init__(self, byte_key = "bytes", bigdl_type="float"):
        super(BytesToMat, self).__init__(bigdl_type, byte_key)

class ImageFrameToSample(FeatureTransformer):
    """
    transform imageframe to samples
    :param input_keys keys that maps inputs (each input should be a tensor)
    :param target_keys keys that maps targets (each target should be a tensor)
    :param sample_key key to store sample
    """
    def __init__(self, input_keys=["imageTensor"], target_keys=None,
                 sample_key="sample", bigdl_type="float"):
        super(ImageFrameToSample, self).__init__(bigdl_type, input_keys, target_keys, sample_key)

class PixelBytesToMat(FeatureTransformer):
    """
    Transform byte array(pixels in byte) to OpenCVMat
    :param byte_key key that maps byte array
    """
    def __init__(self, byte_key = "bytes", bigdl_type="float"):
        super(PixelBytesToMat, self).__init__(bigdl_type, byte_key)

class FixExpand(FeatureTransformer):
    """
    Expand image with given expandHeight and expandWidth,
    put the original image to the center of expanded image
    :param expand_height height expand to
    :param expand_width width expand to
    """
    def __init__(self, expand_height, expand_width, bigdl_type="float"):
        super(FixExpand, self).__init__(bigdl_type, expand_height, expand_width)

class ChannelScaledNormalizer(FeatureTransformer):
    """
    Scaled image at channel level with offset and scale
    :param mean_r : offset for R channel
    :param mean_g : offset for G channel
    :param mean_b: offset for B channel
    :param scale: scaling factor for all channels
    """
    def __init__(self, mean_r, mean_g, mean_b, scale, bigdl_type="float"):
        super(ChannelScaledNormalizer, self).__init__(bigdl_type, mean_r, mean_g, mean_b, scale)

class RandomAlterAspect(FeatureTransformer):
    """
    Apply random crop based on area ratio and resize to cropLenth size
    :param min_area_ratio  min area ratio
    :param max_area_ratio  max area ratio
    :param min_aspect_ratio_change factor applied to ratio area
    :param interp_mode   interp mode applied in resize
    :param crop_length final size resized to
    """
    def __init__(self, min_area_ratio,
                 max_area_ratio,
                 min_aspect_ratio_change,
                 interp_mode,
                 crop_length, bigdl_type="float"):
        super(RandomAlterAspect, self).__init__(bigdl_type, min_area_ratio,
                                                max_area_ratio,
                                                min_aspect_ratio_change,
                                                interp_mode,
                                                crop_length)

class RandomCropper(FeatureTransformer):
    """
    Random cropper on uniform distribution with fixed height & width
    :param crop_w  width cropped to
    :param crop_h height cropped to
    :param mirror   whether mirror
    :param cropper_method crop method
    :param channels total channels
    """
    def __init__(self, crop_w, crop_h, mirror, cropper_method, channels, bigdl_type="float"):
        super(RandomCropper, self).__init__(bigdl_type, crop_w, crop_h, mirror, cropper_method, channels)

class RandomResize(FeatureTransformer):
    """
    Random resize between minSize and maxSize and scale height and width to each other
    :param min_size min size to resize to
    :param max_size max size to resize to
    """
    def __init__(self, min_size, max_size, bigdl_type="float"):
        super(RandomResize, self).__init__(bigdl_type, min_size, max_size)

class SeqFileFolder(JavaValue):

    @classmethod
    def files_to_image_frame(cls,
                             url,
                             sc,
                             class_num,
                             partition_num=-1,
                             bigdl_type="float"):
        """
        Extract hadoop sequence files from an HDFS path as ImageFrame
        :param url: sequence files folder path
        :param sc: spark context
        :param class_num: class number of data
        :param partition_num: partition number, default: Engine.nodeNumber() * Engine.coreNumber()
        """
        jvalue = callBigDlFunc(bigdl_type,
                               "seqFilesToImageFrame",
                               url,
                               sc,
                               class_num,
                               partition_num)
        return ImageFrame(jvalue=jvalue)


#
# Copyright 2016 The BigDL Authors.
#
# Licensed 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 bigdl.keras.optimization import *
from bigdl.util.common import *

class KerasModelWrapper:

    def __init__(self, kmodel):
        redire_spark_logs()
        show_bigdl_info_logs()
        self.bmodel = DefinitionLoader.from_kmodel(kmodel)
        WeightLoader.load_weights_from_kmodel(self.bmodel, kmodel)  # share the same weight.
        self.criterion = OptimConverter.to_bigdl_criterion(kmodel.loss) if kmodel.loss else None
        self.optim_method =\
            OptimConverter.to_bigdl_optim_method(kmodel.optimizer) if kmodel.optimizer else None
        self.metrics = OptimConverter.to_bigdl_metrics(kmodel.metrics) if kmodel.metrics else None

    def evaluate(self, x, y, batch_size=32, sample_weight=None, is_distributed=False):
        """
        Evaluate a model by the given metrics.
        :param x: ndarray or list of ndarray for local mode.
                  RDD[Sample] for distributed mode
        :param y: ndarray or list of ndarray for local mode and would be None for cluster mode.
        :param batch_size
        :param is_distributed: run in local mode or distributed mode.
               NB: if is_distributed=true, x should be RDD[Sample] and y should be None
        :return:
        """
        if sample_weight:
            unsupport_exp("sample_weight")
        if is_distributed:
            if isinstance(x, np.ndarray):
                input = to_sample_rdd(x, y)
            elif isinstance(x, RDD):
                input = x
            if self.metrics:
                sc = get_spark_context()
                return [r.result for r in
                        self.bmodel.evaluate(input, batch_size, self.metrics)]
            else:
                raise Exception("No Metrics found.")
        else:
            raise Exception("We only support evaluation in distributed mode")

    def predict(self, x, batch_size=None, verbose=None, is_distributed=False):
        """Generates output predictions for the input samples,
        processing the samples in a batched way.

        # Arguments
            x: the input data, as a Numpy array or list of Numpy array for local mode.
               as RDD[Sample] for distributed mode
            is_distributed: used to control run in local or cluster. the default value is False
        # Returns
            A Numpy array or RDD[Sample] of predictions.
        """
        if batch_size or verbose:
            raise Exception("we don't support batch_size or verbose for now")
        if is_distributed:
            if isinstance(x, np.ndarray):
                input = to_sample_rdd(x, np.zeros([x.shape[0]]))
            #  np.asarray(self.bmodel.predict(x_rdd).collect())
            elif isinstance(x, RDD):
                input = x
            return self.bmodel.predict(input)
        else:
            if isinstance(x, np.ndarray):
                return self.bmodel.predict_local(x)
        raise Exception("not supported type: %s" % x)

    def fit(self, x, y=None, batch_size=32, nb_epoch=10, verbose=1, callbacks=None,
            validation_split=0., validation_data=None, shuffle=True,
            class_weight=None, sample_weight=None, initial_epoch=0, is_distributed=False):
        """Optimize the model by the given options

        :param x: ndarray or list of ndarray for local mode.
                  RDD[Sample] for distributed mode
        :param y: ndarray or list of ndarray for local mode and would be None for cluster mode.
            is_distributed: used to control run in local or cluster. the default value is False.
            NB: if is_distributed=true, x should be RDD[Sample] and y should be None
        :param is_distributed: Whether to train in local mode or distributed mode
        :return:
            A Numpy array or RDD[Sample] of predictions.
        """
        if callbacks:
            raise Exception("We don't support callbacks in fit for now")
        if class_weight:
            unsupport_exp("class_weight")
        if sample_weight:
            unsupport_exp("sample_weight")
        if initial_epoch != 0:
            unsupport_exp("initial_epoch")
        if shuffle != True:
            unsupport_exp("shuffle")
        if validation_split != 0.:
            unsupport_exp("validation_split")
        bopt = self.__create_optimizer(x=x,
                                       y=y,
                                       batch_size=batch_size,
                                       nb_epoch=nb_epoch,
                                       validation_data=validation_data,
                                       is_distributed=is_distributed)
        bopt.optimize()

    def __create_optimizer(self, x=None, y=None, batch_size=32, nb_epoch=10,
                           validation_data=None, is_distributed=False):
        if is_distributed:
            if isinstance(x, np.ndarray):
                input = to_sample_rdd(x, y)
                validation_data_rdd = to_sample_rdd(*validation_data)
            elif isinstance(x, RDD):
                input = x
                validation_data_rdd = validation_data
            return self.__create_distributed_optimizer(training_rdd=input,
                                                       batch_size=batch_size,
                                                       nb_epoch=nb_epoch,
                                                       validation_data=validation_data_rdd)
        else:
            if isinstance(x, np.ndarray):
                return self.__create_local_optimizer(x, y,
                                                     batch_size=batch_size,
                                                     nb_epoch=nb_epoch,
                                                     validation_data=validation_data)
        raise Exception("not supported type: %s" % x)

    def __create_local_optimizer(self, x, y, batch_size=32, nb_epoch=10, validation_data=None):
        if validation_data:
            raise unsupport_exp("validation_data")
        bopt = boptimizer.LocalOptimizer(
            X=x,
            Y=y,
            model=self.bmodel,
            criterion=self.criterion,
            end_trigger=boptimizer.MaxEpoch(nb_epoch),
            batch_size=batch_size,
            optim_method=self.optim_method,
            cores=None
        )
        # TODO: enable validation for local optimizer.
        return bopt

    def __create_distributed_optimizer(self, training_rdd,
                                       batch_size=32,
                                       nb_epoch=10,
                                       validation_data=None):
        sc = get_spark_context()
        bopt = boptimizer.Optimizer(
            model=self.bmodel,
            training_rdd=training_rdd,
            criterion=self.criterion,
            end_trigger=boptimizer.MaxEpoch(nb_epoch),
            batch_size=batch_size,
            optim_method=self.optim_method
        )
        if validation_data:
            bopt.set_validation(batch_size,
                                val_rdd=validation_data,
                                # TODO: check if keras use the same strategy
                                trigger=boptimizer.EveryEpoch(),
                                val_method=self.metrics)
        return bopt

def with_bigdl_backend(kmodel):
    init_engine()
    return KerasModelWrapper(kmodel)


#
# Copyright 2016 The BigDL Authors.
#
# Licensed 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.
#

import sys
from bigdl.util.common import JavaValue, callBigDlFunc
from pyspark.ml.param.shared import *
from pyspark.ml.wrapper import JavaTransformer

if sys.version >= '3':
    long = int
    unicode = str

class DLImageTransformer(JavaTransformer, HasInputCol, HasOutputCol, JavaValue):
    """
    Provides DataFrame-based API for image pre-processing and feature transformation.
    DLImageTransformer follows the Spark Transformer API pattern and can be used as one stage
    in Spark ML pipeline.

    The input column can be either DLImageSchema.byteSchema or DLImageSchema.floatSchema. If
    using DLImageReader, the default format is DLImageSchema.byteSchema
    The output column is always DLImageSchema.floatSchema.
    """

    def __init__(self,  transformer, jvalue=None, bigdl_type="float"):
        super(DLImageTransformer, self).__init__()
        self.value = jvalue if jvalue else callBigDlFunc(
            bigdl_type, self.jvm_class_constructor(), transformer)
        self._java_obj = self.value
        self.bigdl_type = bigdl_type

    def transform(self, dataset):
        """
        Apply the transformer to the images in "inputCol" and store the transformed result
        into "outputCols"
        """
        self._transfer_params_to_java()
        return callBigDlFunc(self.bigdl_type, "dlImageTransform", self.value, dataset)


#
# Copyright 2016 The BigDL Authors.
#
# Licensed 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 bigdl.util.common import *

def save_keras_definition(keras_model, path):
    """
    Save a Keras model definition to JSON with given path
    """
    model_json = keras_model.to_json()
    with open(path, "w") as json_file:
        json_file.write(model_json)

def dump_keras(keras_model, json_path=None, hdf5_path=None, dump_weights=False):
    tmp_path = create_tmp_path()
    if not json_path:
        json_path = tmp_path + ".json"
    if not hdf5_path:
        hdf5_path = tmp_path + ".hdf5"
    save_keras_definition(keras_model, json_path)
    print("json path: " + json_path)
    if dump_weights:
        keras_model.save(hdf5_path)
        print("hdf5 path: " + hdf5_path)
    return json_path, hdf5_path


#
# Copyright 2016 The BigDL Authors.
#
# Licensed 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.
#

# MNIST CNN Example on BigDL
# Reference: https://github.com/fchollet/keras/blob/1.2.2/examples/mnist_cnn.py
#            ../../models/lenet/lenet5.py
# The Keras version we support and test is Keras 1.2.2 with TensorFlow backend.
# See README.md for how to run this example.

from optparse import OptionParser
from bigdl.examples.keras.keras_utils import *

import keras.backend
if keras.backend.image_dim_ordering() == "th":
    input_shape = (1, 28, 28)
else:
    input_shape = (28, 28, 1)

def get_mnist(sc, data_type="train", location="/tmp/mnist"):
    """
    Download or load MNIST dataset to/from the specified path.
    Normalize and transform input data into an RDD of Sample
    """
    from bigdl.dataset import mnist
    from bigdl.dataset.transformer import normalizer
    (images, labels) = mnist.read_data_sets(location, data_type)
    images = images.reshape((images.shape[0], ) + input_shape)
    images = sc.parallelize(images)
    labels = sc.parallelize(labels + 1)  # Target start from 1 in BigDL
    record = images.zip(labels).map(lambda rec_tuple: (normalizer(rec_tuple[0], mnist.TRAIN_MEAN, mnist.TRAIN_STD),
                                    rec_tuple[1])) \
                               .map(lambda t: Sample.from_ndarray(t[0], t[1]))
    return record

def build_keras_model():
    """
    Define a convnet model in Keras 1.2.2
    """
    from keras.models import Sequential
    from keras.layers import Dense, Dropout, Activation, Flatten
    from keras.layers import Convolution2D, MaxPooling2D

    keras_model = Sequential()
    keras_model.add(Convolution2D(32, 3, 3, border_mode='valid',
                                  input_shape=input_shape))
    keras_model.add(Activation('relu'))
    keras_model.add(Convolution2D(32, 3, 3))
    keras_model.add(Activation('relu'))
    keras_model.add(MaxPooling2D(pool_size=(2, 2)))
    keras_model.add(Dropout(0.25))
    keras_model.add(Flatten())
    keras_model.add(Dense(128))
    keras_model.add(Activation('relu'))
    keras_model.add(Dropout(0.5))
    keras_model.add(Dense(10))
    keras_model.add(Activation('softmax'))
    return keras_model

if __name__ == "__main__":
    parser = OptionParser()
    parser.add_option("-b", "--batchSize", type=int, dest="batchSize", default="128")
    parser.add_option("-m", "--max_epoch", type=int, dest="max_epoch", default="12")
    parser.add_option("-d", "--dataPath", dest="dataPath", default="/tmp/mnist")
    (options, args) = parser.parse_args(sys.argv)

    keras_model = build_keras_model()
    json_path = "/tmp/lenet.json"
    save_keras_definition(keras_model, json_path)

    from bigdl.util.common import *
    from bigdl.nn.layer import *
    from bigdl.optim.optimizer import *
    from bigdl.nn.criterion import *

    # Load the JSON file to a BigDL model
    bigdl_model = Model.load_keras(json_path=json_path)

    sc = get_spark_context(conf=create_spark_conf())
    redire_spark_logs()
    show_bigdl_info_logs()
    init_engine()

    train_data = get_mnist(sc, "train", options.dataPath)
    test_data = get_mnist(sc, "test", options.dataPath)

    optimizer = Optimizer(
        model=bigdl_model,
        training_rdd=train_data,
        criterion=ClassNLLCriterion(logProbAsInput=False),
        optim_method=Adadelta(),
        end_trigger=MaxEpoch(options.max_epoch),
        batch_size=options.batchSize)
    optimizer.set_validation(
        batch_size=options.batchSize,
        val_rdd=test_data,
        trigger=EveryEpoch(),
        val_method=[Top1Accuracy()]
    )
    optimizer.optimize()