# Copyright 1999-2020 Alibaba Group Holding Ltd.
#
# 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 functools
import gzip
import struct
import zlib
from collections import namedtuple
from enum import Enum
from io import BytesIO

try:
    import numpy as np
except ImportError:  # pragma: no cover
    np = None
try:
    import scipy.sparse as sps
except ImportError:  # pragma: no cover
    sps = None

from ..lib.sparse import SparseNDArray

try:
    import pyarrow
except ImportError:  # pragma: no cover
    pyarrow = None


BUFFER_SIZE = 256 * 1024


class DummyCompress(object):
    def __enter__(self):  # pragma: no cover
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):  # pragma: no cover
        pass

    @staticmethod
    def compress(data):  # pragma: no cover
        return data

    @staticmethod
    def decompress(data):  # pragma: no cover
        return data


try:
    import lz4.frame
    try:
        lz4.frame._compression.BUFFER_SIZE = BUFFER_SIZE
    except AttributeError:  # pragma: no cover
        pass
    lz4_open = functools.partial(lz4.frame.open, block_size=lz4.frame.BLOCKSIZE_MAX1MB)
    lz4_compress = functools.partial(lz4.frame.compress, block_size=lz4.frame.BLOCKSIZE_MAX1MB)
    lz4_compressobj = lz4.frame.LZ4FrameCompressor
    lz4_decompress = lz4.frame.decompress
    lz4_decompressobj = lz4.frame.LZ4FrameDecompressor
except ImportError:  # pragma: no cover
    lz4_open = None
    lz4_compress, lz4_compressobj = None, None
    lz4_decompress, lz4_decompressobj = None, None

gz_open = gzip.open
gz_compressobj = functools.partial(
    lambda level=-1: zlib.compressobj(level, zlib.DEFLATED, 16 + zlib.MAX_WBITS, zlib.DEF_MEM_LEVEL, 0)
)
gz_decompressobj = functools.partial(lambda: zlib.decompressobj(16 + zlib.MAX_WBITS))
gz_compress = gzip.compress
gz_decompress = gzip.decompress

SERIAL_VERSION = 0


class CompressType(Enum):
    NONE = 'none'
    LZ4 = 'lz4'
    GZIP = 'gzip'

    @property
    def tag(self):
        return _compress_tags[self]

    @classmethod
    def from_tag(cls, tag):
        return _tag_to_compress[tag]

    def __gt__(self, other):
        return _compress_tags[self] > _compress_tags[other]


_compress_tags = {
    CompressType.NONE: 0,
    CompressType.LZ4: 1,
    CompressType.GZIP: 2,
}
_tag_to_compress = {
    0: CompressType.NONE,
    1: CompressType.LZ4,
    2: CompressType.GZIP,
}

compressors = {
    CompressType.LZ4: lz4_compress,
    CompressType.GZIP: gz_compress,
}
compressobjs = {
    CompressType.NONE: DummyCompress,
    CompressType.LZ4: lz4_compressobj,
    CompressType.GZIP: gz_compressobj,
}
decompressors = {
    CompressType.LZ4: lz4_decompress,
    CompressType.GZIP: gz_decompress,
}
decompressobjs = {
    CompressType.NONE: DummyCompress,
    CompressType.LZ4: lz4_decompressobj,
    CompressType.GZIP: gz_decompressobj,
}
compress_openers = {
    CompressType.LZ4: lz4_open,
    CompressType.GZIP: gz_open,
}


def get_supported_compressions():
    return set(k for k, v in decompressors.items() if v is not None)


def get_compressobj(compress):
    return compressobjs[compress]()


def get_decompressobj(compress):
    return decompressobjs[compress]()


def open_compression_file(file, compress):
    if compress != CompressType.NONE:
        file = compress_openers[compress](file, 'wb')
    return file


def open_decompression_file(file, compress):
    if compress != CompressType.NONE:
        file = compress_openers[compress](file, 'rb')
    return file


file_header = namedtuple('FileHeader', 'version nbytes compress')
HEADER_LENGTH = 12


def read_file_header(file):
    if hasattr(file, 'read'):
        header_bytes = file.read(HEADER_LENGTH)
    else:
        header_bytes = file[:HEADER_LENGTH]
    version, = struct.unpack('<H', header_bytes[:2])
    nbytes, = struct.unpack('<Q', header_bytes[2:10])
    compress, = struct.unpack('<H', header_bytes[10:12])
    return file_header(version, nbytes, CompressType.from_tag(compress))


def write_file_header(file, header):
    file.write(struct.pack('<H', header.version))
    file.write(struct.pack('<Q', header.nbytes))
    file.write(struct.pack('<H', header.compress.tag))


def peek_file_header(file):
    pos = file.tell()
    try:
        return read_file_header(file)
    finally:
        file.seek(pos)


def load(file, raw=False):
    header = read_file_header(file)
    file = open_decompression_file(file, header.compress)

    try:
        buf = file.read()
    finally:
        if header.compress != CompressType.NONE:
            file.close()

    if raw:
        return buf
    else:
        return pyarrow.deserialize(memoryview(buf), mars_serialize_context())


def loads(buf, raw=False):
    mv = memoryview(buf)
    header = read_file_header(mv)
    compress = header.compress

    if compress == CompressType.NONE:
        data = buf[HEADER_LENGTH:]
    else:
        data = decompressors[compress](mv[HEADER_LENGTH:])
    if raw:
        return data
    else:
        return pyarrow.deserialize(memoryview(data), mars_serialize_context())


def dump(obj, file, compress=CompressType.NONE, raw=False):
    if raw:
        serialized = obj
        data_size = len(serialized)
    else:
        serialized = pyarrow.serialize(obj, mars_serialize_context())
        data_size = serialized.total_bytes

    write_file_header(file, file_header(SERIAL_VERSION, data_size, compress))
    file = open_compression_file(file, compress)
    try:
        if raw:
            file.write(serialized)
        else:
            serialized.write_to(file)
    finally:
        if compress != CompressType.NONE:
            file.close()
    return


def dumps(obj, compress=CompressType.NONE, raw=False):
    sio = BytesIO()
    dump(obj, sio, compress=compress, raw=raw)
    return sio.getvalue()


def serialize(data):
    return pyarrow.serialize(data, mars_serialize_context())


def deserialize(data):
    return pyarrow.deserialize(data, mars_serialize_context())


def _serialize_numpy_array_list(obj):
    if obj.dtype.str != '|O':
        # Make the array c_contiguous if necessary so that we can call change
        # the view.
        if not obj.flags.c_contiguous:
            obj = np.ascontiguousarray(obj)
        return obj.view('uint8'), obj.dtype.str
    else:
        return obj.tolist(), obj.dtype.str


def _deserialize_numpy_array_list(data):
    if data[1] != '|O':
        assert data[0].dtype == np.uint8
        return data[0].view(data[1])
    else:
        return np.array(data[0], dtype=np.dtype(data[1]))


def _serialize_sparse_csr_list(obj):
    sizes = []
    outputs = [None, None]
    for item in (obj.data, obj.indices, obj.indptr):
        serialized_items = _serialize_numpy_array_list(item)
        outputs.extend(serialized_items)
        sizes.append(len(serialized_items))
    outputs[0] = struct.pack('<' + 'B' * len(sizes), *sizes)
    outputs[1] = struct.pack('<' + 'L' * len(obj.shape), *obj.shape)
    return tuple(outputs)


def _deserialize_sparse_csr_list(data):
    sizes = struct.unpack('<' + 'B' * len(data[0]), data[0])
    shape = struct.unpack('<' + 'L' * (len(data[1]) // 4), data[1])
    data_parts = []
    pos = 2
    for size in sizes:
        data_parts.append(_deserialize_numpy_array_list(data[pos:pos + size]))
        pos += size

    empty_arr = np.zeros(0, dtype=data_parts[0].dtype)
    if len(shape) == 1:
        sps_shape = (1, shape[0])
    else:
        sps_shape = shape
    target_csr = sps.coo_matrix((empty_arr, (empty_arr,) * 2), shape=sps_shape,
                                dtype=data_parts[0].dtype).tocsr()
    target_csr.data, target_csr.indices, target_csr.indptr = data_parts
    return SparseNDArray(target_csr, shape=shape)


_serialize_context = None


def _apply_pyarrow_serialization_patch(serialization_context):
    """
    Fix the bug about dtype serialization in pyarrow (pyarrow#4953).

    From the JIRA of arrow, the fixed version of this bug is 1.0, so we only apply
    the patch for pyarrow less than version 1.0.
    """
    from distutils.version import LooseVersion
    import numpy as np
    import pyarrow

    try:
        # This function is available after numpy-0.16.0.
        # See also: https://github.com/numpy/numpy/blob/master/numpy/lib/format.py
        from numpy.lib.format import descr_to_dtype
    except ImportError:
        def descr_to_dtype(descr):
            '''
            descr may be stored as dtype.descr, which is a list of
            (name, format, [shape]) tuples where format may be a str or a tuple.
            Offsets are not explicitly saved, rather empty fields with
            name, format == '', '|Vn' are added as padding.
            This function reverses the process, eliminating the empty padding fields.
            '''
            if isinstance(descr, str):
                # No padding removal needed
                return np.dtype(descr)
            elif isinstance(descr, tuple):
                # subtype, will always have a shape descr[1]
                dt = descr_to_dtype(descr[0])
                return np.dtype((dt, descr[1]))
            fields = []
            offset = 0
            for field in descr:
                if len(field) == 2:
                    name, descr_str = field
                    dt = descr_to_dtype(descr_str)
                else:
                    name, descr_str, shape = field
                    dt = np.dtype((descr_to_dtype(descr_str), shape))

                # Ignore padding bytes, which will be void bytes with '' as name
                # Once support for blank names is removed, only "if name == ''" needed)
                is_pad = (name == '' and dt.type is np.void and dt.names is None)
                if not is_pad:
                    fields.append((name, dt, offset))

                offset += dt.itemsize

            names, formats, offsets = zip(*fields)
            # names may be (title, names) tuples
            nametups = (n if isinstance(n, tuple) else (None, n) for n in names)
            titles, names = zip(*nametups)
            return np.dtype({'names': names, 'formats': formats, 'titles': titles,
                             'offsets': offsets, 'itemsize': offset})

    def _serialize_numpy_array_list(obj):
        if obj.dtype.str != '|O':
            # Make the array c_contiguous if necessary so that we can call change
            # the view.
            if not obj.flags.c_contiguous:
                obj = np.ascontiguousarray(obj)
            return obj.view('uint8'), np.lib.format.dtype_to_descr(obj.dtype)
        else:
            return obj.tolist(), np.lib.format.dtype_to_descr(obj.dtype)

    def _deserialize_numpy_array_list(data):
        if data[1] != '|O':
            assert data[0].dtype == np.uint8
            return data[0].view(descr_to_dtype(data[1]))
        else:
            return np.array(data[0], dtype=np.dtype(data[1]))

    if LooseVersion(pyarrow.__version__) < LooseVersion('1.0'):
        serialization_context.register_type(
            np.ndarray, 'np.array',
            custom_serializer=_serialize_numpy_array_list,
            custom_deserializer=_deserialize_numpy_array_list)


def mars_serialize_context():
    global _serialize_context
    if _serialize_context is None:
        ctx = pyarrow.default_serialization_context()
        ctx.register_type(SparseNDArray, 'mars.SparseNDArray',
                          custom_serializer=_serialize_sparse_csr_list,
                          custom_deserializer=_deserialize_sparse_csr_list)
        _apply_pyarrow_serialization_patch(ctx)
        _serialize_context = ctx
    return _serialize_context
