# -------------------------------------------------------------
#
# 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__ = [
    'getNumCols',
    'convertToMatrixBlock',
    'convert_caffemodel',
    'convert_lmdb_to_jpeg',
    'convertToNumPyArr',
    'convertToPandasDF',
    'SUPPORTED_TYPES',
    'convertToLabeledDF',
    'convertImageToNumPyArr',
    'getDatasetMean']

import numpy as np
import pandas as pd
import os
import math

from pyspark.context import SparkContext
from scipy.sparse import coo_matrix, spmatrix, csr_matrix
from .classloader import *

SUPPORTED_TYPES = (np.ndarray, pd.DataFrame, spmatrix)

DATASET_MEAN = {'VGG_ILSVRC_19_2014': [103.939, 116.779, 123.68]}


def getNumCols(numPyArr):
    if numPyArr.ndim == 1:
        return 1
    else:
        return numPyArr.shape[1]


def get_pretty_str(key, value):
    return '\t"' + key + '": ' + str(value) + ',\n'


def save_tensor_csv(tensor, file_path, shouldTranspose):
    w = w.reshape(w.shape[0], -1)
    if shouldTranspose:
        w = w.T
    np.savetxt(file_path, w, delimiter=',')
    with open(file_path + '.mtd', 'w') as file:
        file.write('{\n\t"data_type": "matrix",\n\t"value_type": "double",\n')
        file.write(get_pretty_str('rows', w.shape[0]))
        file.write(get_pretty_str('cols', w.shape[1]))
        file.write(get_pretty_str('nnz', np.count_nonzero(w)))
        file.write(
            '\t"format": "csv",\n\t"description": {\n\t\t"author": "SystemML"\n\t}\n}\n')


def convert_caffemodel(sc, deploy_file, caffemodel_file,
                       output_dir, format="binary", is_caffe_installed=False):
    """
    Saves the weights and bias in the caffemodel file to output_dir in the specified format.
    This method does not requires caffe to be installed.

    Parameters
    ----------
    sc: SparkContext
        SparkContext

    deploy_file: string
        Path to the input network file

    caffemodel_file: string
        Path to the input caffemodel file

    output_dir: string
        Path to the output directory

    format: string
        Format of the weights and bias (can be binary, csv or text)

    is_caffe_installed: bool
        True if caffe is installed
    """
    if is_caffe_installed:
        if format != 'csv':
            raise ValueError(
                'The format ' +
                str(format) +
                ' is not supported when caffe is installed. Hint: Please specify format=csv')
        import caffe
        net = caffe.Net(deploy_file, caffemodel_file, caffe.TEST)
        for layerName in net.params.keys():
            num_parameters = len(net.params[layerName])
            if num_parameters == 0:
                continue
            elif num_parameters == 2:
                # Weights and Biases
                layerType = net.layers[list(
                    net._layer_names).index(layerName)].type
                shouldTranspose = True if layerType == 'InnerProduct' else False
                save_tensor_csv(
                    net.params[layerName][0].data,
                    os.path.join(
                        output_dir,
                        layerName +
                        '_weight.mtx'),
                    shouldTranspose)
                save_tensor_csv(
                    net.params[layerName][1].data,
                    os.path.join(
                        output_dir,
                        layerName +
                        '_bias.mtx'),
                    shouldTranspose)
            elif num_parameters == 1:
                # Only Weight
                layerType = net.layers[list(
                    net._layer_names).index(layerName)].type
                shouldTranspose = True if layerType == 'InnerProduct' else False
                save_tensor_csv(
                    net.params[layerName][0].data,
                    os.path.join(
                        output_dir,
                        layerName +
                        '_weight.mtx'),
                    shouldTranspose)
            else:
                raise ValueError(
                    'Unsupported number of parameters:' +
                    str(num_parameters))
    else:
        createJavaObject(sc, 'dummy')
        utilObj = sc._jvm.org.apache.sysml.api.dl.Utils()
        utilObj.saveCaffeModelFile(
            sc._jsc,
            deploy_file,
            caffemodel_file,
            output_dir,
            format)


def convert_lmdb_to_jpeg(lmdb_img_file, output_dir):
    """
    Saves the images in the lmdb file as jpeg in the output_dir. This method requires caffe to be installed along with lmdb and cv2 package.
    To install cv2 package, do `pip install opencv-python`.

    Parameters
    ----------
    lmdb_img_file: string
        Path to the input lmdb file

    output_dir: string
        Output directory for images (local filesystem)
    """
    import lmdb
    import caffe
    import cv2
    lmdb_cursor = lmdb.open(lmdb_file, readonly=True).begin().cursor()
    datum = caffe.proto.caffe_pb2.Datum()
    i = 1
    for _, value in lmdb_cursor:
        datum.ParseFromString(value)
        data = caffe.io.datum_to_array(datum)
        output_file_path = os.path.join(output_dir, 'file_' + str(i) + '.jpg')
        image = np.transpose(data, (1, 2, 0))  # CxHxW to HxWxC in cv2
        cv2.imwrite(output_file_path, image)
        i = i + 1


def convertToLabeledDF(sparkSession, X, y=None):
    from pyspark.ml.feature import VectorAssembler
    if y is not None:
        pd1 = pd.DataFrame(X)
        pd2 = pd.DataFrame(y, columns=['label'])
        pdf = pd.concat([pd1, pd2], axis=1)
        inputColumns = ['C' + str(i) for i in pd1.columns]
        outputColumns = inputColumns + ['label']
    else:
        pdf = pd.DataFrame(X)
        inputColumns = ['C' + str(i) for i in pdf.columns]
        outputColumns = inputColumns
    assembler = VectorAssembler(inputCols=inputColumns, outputCol='features')
    out = assembler.transform(sparkSession.createDataFrame(pdf, outputColumns))
    if y is not None:
        return out.select('features', 'label')
    else:
        return out.select('features')


def _convertSPMatrixToMB(sc, src):
    src = coo_matrix(src, dtype=np.float64)
    numRows = src.shape[0]
    numCols = src.shape[1]
    data = src.data
    row = src.row.astype(np.int32)
    col = src.col.astype(np.int32)
    nnz = len(src.col)
    buf1 = bytearray(data.tostring())
    buf2 = bytearray(row.tostring())
    buf3 = bytearray(col.tostring())
    createJavaObject(sc, 'dummy')
    return sc._jvm.org.apache.sysml.runtime.instructions.spark.utils.RDDConverterUtilsExt.convertSciPyCOOToMB(
        buf1, buf2, buf3, numRows, numCols, nnz)


def _convertDenseMatrixToMB(sc, src):
    numCols = getNumCols(src)
    numRows = src.shape[0]
    src = np.asarray(src, dtype=np.float64) if not isinstance(src, np.ndarray) else src
    # data_type: 0: int, 1: float and 2: double
    if src.dtype is np.dtype(np.int32):
        arr = src.ravel().astype(np.int32)
        dataType = 0
    elif src.dtype is np.dtype(np.float32):
        arr = src.ravel().astype(np.float32)
        dataType = 1
    else:
        arr = src.ravel().astype(np.float64)
        dataType = 2 
    buf = bytearray(arr.tostring())
    createJavaObject(sc, 'dummy')
    return sc._jvm.org.apache.sysml.runtime.instructions.spark.utils.RDDConverterUtilsExt.convertPy4JArrayToMB(
        buf, numRows, numCols, dataType)


def _copyRowBlock(i, sc, ret, src, numRowsPerBlock, rlen, clen):
    rowIndex = int(i / numRowsPerBlock)
    tmp = src[i:min(i + numRowsPerBlock, rlen), ]
    mb = _convertSPMatrixToMB(
        sc,
        tmp) if isinstance(
        src,
        spmatrix) else _convertDenseMatrixToMB(
            sc,
        tmp)
    sc._jvm.org.apache.sysml.runtime.instructions.spark.utils.RDDConverterUtilsExt.copyRowBlocks(
        mb, rowIndex, ret, numRowsPerBlock, rlen, clen)
    return i


def convertToMatrixBlock(sc, src, maxSizeBlockInMB=128):
    if not isinstance(sc, SparkContext):
        raise TypeError('sc needs to be of type SparkContext')
    if isinstance(src, spmatrix):
        isSparse = True
    else:
        isSparse = False
        src = np.asarray(src, dtype=np.float64) if not isinstance(src, np.ndarray) else src
    if len(src.shape) != 2:
        src_type = str(type(src).__name__)
        raise TypeError('Expected 2-dimensional ' +
                        src_type +
                        ', instead passed ' +
                        str(len(src.shape)) +
                        '-dimensional ' +
                        src_type)
    worstCaseSizeInMB = (8*(src.getnnz()*3 if isSparse else src.shape[0]*src.shape[1])) / 1000000
    # Ignoring sparsity for computing numRowsPerBlock for now
    numRowsPerBlock = int(
        math.ceil((maxSizeBlockInMB * 1000000) / (src.shape[1] * 8)))
    if worstCaseSizeInMB <= maxSizeBlockInMB:
        return _convertSPMatrixToMB(
            sc, src) if isSparse else _convertDenseMatrixToMB(sc, src)
    else:
        # Since coo_matrix does not have range indexing
        src = csr_matrix(src) if isSparse else src
        rlen = int(src.shape[0])
        clen = int(src.shape[1])
        ret = sc._jvm.org.apache.sysml.runtime.instructions.spark.utils.RDDConverterUtilsExt.allocateDenseOrSparse(
            rlen, clen, isSparse)
        [_copyRowBlock(i, sc, ret, src, numRowsPerBlock, rlen, clen)
         for i in range(0, src.shape[0], numRowsPerBlock)]
        sc._jvm.org.apache.sysml.runtime.instructions.spark.utils.RDDConverterUtilsExt.postProcessAfterCopying(
            ret)
        return ret


def convertToNumPyArr(sc, mb):
    if isinstance(sc, SparkContext):
        numRows = mb.getNumRows()
        numCols = mb.getNumColumns()
        createJavaObject(sc, 'dummy')
        buf = sc._jvm.org.apache.sysml.runtime.instructions.spark.utils.RDDConverterUtilsExt.convertMBtoPy4JDenseArr(
            mb)
        return np.frombuffer(buf, count=numRows * numCols,
                             dtype=np.float64).reshape((numRows, numCols))
    else:
        # TODO: We can generalize this by creating py4j gateway ourselves
        raise TypeError('sc needs to be of type SparkContext')

# Returns the mean of a model if defined otherwise None


def getDatasetMean(dataset_name):
    """
    Parameters
    ----------
    dataset_name: Name of the dataset used to train model. This name is artificial name based on dataset used to train the model.

    Returns
    -------
    mean: Mean value of model if its defined in the list DATASET_MEAN else None.

    """

    try:
        mean = DATASET_MEAN[dataset_name.upper()]
    except BaseException:
        mean = None
    return mean


# Example usage: convertImageToNumPyArr(im, img_shape=(3, 224, 224), add_rotated_images=True, add_mirrored_images=True)
# The above call returns a numpy array of shape (6, 50176) in NCHW format
def convertImageToNumPyArr(im, img_shape=None, add_rotated_images=False, add_mirrored_images=False,
                           color_mode='RGB', mean=None):

    # Input Parameters

    # color_mode: In case of VGG models which expect image data in BGR format instead of RGB for other most models,
    # color_mode parameter is used to process image data in BGR format.

    # mean: mean value is used to subtract from input data from every pixel
    # value. By default value is None, so mean value not subtracted.

    if img_shape is not None:
        num_channels = img_shape[0]
        size = (img_shape[1], img_shape[2])
    else:
        num_channels = 1 if im.mode == 'L' else 3
        size = None
    if num_channels != 1 and num_channels != 3:
        raise ValueError('Expected the number of channels to be either 1 or 3')

    from PIL import Image

    if size is not None:
        im = im.resize(size, Image.LANCZOS)
    expected_mode = 'L' if num_channels == 1 else 'RGB'
    if expected_mode is not im.mode:
        im = im.convert(expected_mode)

    def _im2NumPy(im):
        if expected_mode == 'L':
            return np.asarray(im.getdata()).reshape((1, -1))
        else:
            im = (np.array(im).astype(np.float))

            # (H,W,C) -> (C,H,W)
            im = im.transpose(2, 0, 1)

            # RGB -> BGR
            if color_mode == 'BGR':
                im = im[..., ::-1]

            # Subtract Mean
            if mean is not None:
                for c in range(3):
                    im[:, :, c] = im[:, :, c] - mean[c]

            # (C,H,W) --> (1, C*H*W)
            return im.reshape((1, -1))

    ret = _im2NumPy(im)

    if add_rotated_images:
        ret = np.vstack(
            (ret, _im2NumPy(
                im.rotate(90)), _im2NumPy(
                im.rotate(180)), _im2NumPy(
                im.rotate(270))))
    if add_mirrored_images:
        ret = np.vstack(
            (ret, _im2NumPy(
                im.transpose(
                    Image.FLIP_LEFT_RIGHT)), _im2NumPy(
                im.transpose(
                    Image.FLIP_TOP_BOTTOM))))
    return ret


def convertToPandasDF(X):
    if not isinstance(X, pd.DataFrame):
        return pd.DataFrame(X, columns=['C' + str(i)
                                        for i in range(getNumCols(X))])
    return X
