# 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.

import collections
import warnings

import numpy as np

import pyarrow as pa
from pyarrow.lib import SerializationContext, py_buffer, builtin_pickle

try:
    import cloudpickle
except ImportError:
    cloudpickle = builtin_pickle


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 _deprecate_serialization(name):
    msg = (
        "'pyarrow.{}' is deprecated as of 2.0.0 and will be removed in a "
        "future version. Use pickle or the pyarrow IPC functionality instead."
    ).format(name)
    warnings.warn(msg, FutureWarning, stacklevel=3)


# ----------------------------------------------------------------------
# Set up serialization for numpy with dtype object (primitive types are
# handled efficiently with Arrow's Tensor facilities, see
# python_to_arrow.cc)

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]))


def _serialize_numpy_matrix(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.A)
        return obj.A.view('uint8'), np.lib.format.dtype_to_descr(obj.dtype)
    else:
        return obj.A.tolist(), np.lib.format.dtype_to_descr(obj.dtype)


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


# ----------------------------------------------------------------------
# pyarrow.RecordBatch-specific serialization matters

def _serialize_pyarrow_recordbatch(batch):
    output_stream = pa.BufferOutputStream()
    with pa.RecordBatchStreamWriter(output_stream, schema=batch.schema) as wr:
        wr.write_batch(batch)
    return output_stream.getvalue()  # This will also close the stream.


def _deserialize_pyarrow_recordbatch(buf):
    with pa.RecordBatchStreamReader(buf) as reader:
        return reader.read_next_batch()


# ----------------------------------------------------------------------
# pyarrow.Array-specific serialization matters

def _serialize_pyarrow_array(array):
    # TODO(suquark): implement more effcient array serialization.
    batch = pa.RecordBatch.from_arrays([array], [''])
    return _serialize_pyarrow_recordbatch(batch)


def _deserialize_pyarrow_array(buf):
    # TODO(suquark): implement more effcient array deserialization.
    batch = _deserialize_pyarrow_recordbatch(buf)
    return batch.columns[0]


# ----------------------------------------------------------------------
# pyarrow.Table-specific serialization matters

def _serialize_pyarrow_table(table):
    output_stream = pa.BufferOutputStream()
    with pa.RecordBatchStreamWriter(output_stream, schema=table.schema) as wr:
        wr.write_table(table)
    return output_stream.getvalue()  # This will also close the stream.


def _deserialize_pyarrow_table(buf):
    with pa.RecordBatchStreamReader(buf) as reader:
        return reader.read_all()


def _pickle_to_buffer(x):
    pickled = builtin_pickle.dumps(x, protocol=builtin_pickle.HIGHEST_PROTOCOL)
    return py_buffer(pickled)


def _load_pickle_from_buffer(data):
    as_memoryview = memoryview(data)
    return builtin_pickle.loads(as_memoryview)


# ----------------------------------------------------------------------
# pandas-specific serialization matters

def _register_custom_pandas_handlers(context):
    # ARROW-1784, faster path for pandas-only visibility

    try:
        import pandas as pd
    except ImportError:
        return

    import pyarrow.pandas_compat as pdcompat

    sparse_type_error_msg = (
        '{0} serialization is not supported.\n'
        'Note that {0} is planned to be deprecated '
        'in pandas future releases.\n'
        'See https://github.com/pandas-dev/pandas/issues/19239 '
        'for more information.'
    )

    def _serialize_pandas_dataframe(obj):
        if (pdcompat._pandas_api.has_sparse and
                isinstance(obj, pd.SparseDataFrame)):
            raise NotImplementedError(
                sparse_type_error_msg.format('SparseDataFrame')
            )

        return pdcompat.dataframe_to_serialized_dict(obj)

    def _deserialize_pandas_dataframe(data):
        return pdcompat.serialized_dict_to_dataframe(data)

    def _serialize_pandas_series(obj):
        if (pdcompat._pandas_api.has_sparse and
                isinstance(obj, pd.SparseSeries)):
            raise NotImplementedError(
                sparse_type_error_msg.format('SparseSeries')
            )

        return _serialize_pandas_dataframe(pd.DataFrame({obj.name: obj}))

    def _deserialize_pandas_series(data):
        deserialized = _deserialize_pandas_dataframe(data)
        return deserialized[deserialized.columns[0]]

    context.register_type(
        pd.Series, 'pd.Series',
        custom_serializer=_serialize_pandas_series,
        custom_deserializer=_deserialize_pandas_series)

    context.register_type(
        pd.Index, 'pd.Index',
        custom_serializer=_pickle_to_buffer,
        custom_deserializer=_load_pickle_from_buffer)

    if hasattr(pd.core, 'arrays'):
        if hasattr(pd.core.arrays, 'interval'):
            context.register_type(
                pd.core.arrays.interval.IntervalArray,
                'pd.core.arrays.interval.IntervalArray',
                custom_serializer=_pickle_to_buffer,
                custom_deserializer=_load_pickle_from_buffer)

        if hasattr(pd.core.arrays, 'period'):
            context.register_type(
                pd.core.arrays.period.PeriodArray,
                'pd.core.arrays.period.PeriodArray',
                custom_serializer=_pickle_to_buffer,
                custom_deserializer=_load_pickle_from_buffer)

        if hasattr(pd.core.arrays, 'datetimes'):
            context.register_type(
                pd.core.arrays.datetimes.DatetimeArray,
                'pd.core.arrays.datetimes.DatetimeArray',
                custom_serializer=_pickle_to_buffer,
                custom_deserializer=_load_pickle_from_buffer)

    context.register_type(
        pd.DataFrame, 'pd.DataFrame',
        custom_serializer=_serialize_pandas_dataframe,
        custom_deserializer=_deserialize_pandas_dataframe)


def register_torch_serialization_handlers(serialization_context):
    # ----------------------------------------------------------------------
    # Set up serialization for pytorch tensors
    _deprecate_serialization("register_torch_serialization_handlers")

    try:
        import torch

        def _serialize_torch_tensor(obj):
            if obj.is_sparse:
                return pa.SparseCOOTensor.from_numpy(
                    obj._values().detach().numpy(),
                    obj._indices().detach().numpy().T,
                    shape=list(obj.shape))
            else:
                return obj.detach().numpy()

        def _deserialize_torch_tensor(data):
            if isinstance(data, pa.SparseCOOTensor):
                return torch.sparse_coo_tensor(
                    indices=data.to_numpy()[1].T,
                    values=data.to_numpy()[0][:, 0],
                    size=data.shape)
            else:
                return torch.from_numpy(data)

        for t in [torch.FloatTensor, torch.DoubleTensor, torch.HalfTensor,
                  torch.ByteTensor, torch.CharTensor, torch.ShortTensor,
                  torch.IntTensor, torch.LongTensor, torch.Tensor]:
            serialization_context.register_type(
                t, "torch." + t.__name__,
                custom_serializer=_serialize_torch_tensor,
                custom_deserializer=_deserialize_torch_tensor)
    except ImportError:
        # no torch
        pass


def _register_collections_serialization_handlers(serialization_context):
    def _serialize_deque(obj):
        return list(obj)

    def _deserialize_deque(data):
        return collections.deque(data)

    serialization_context.register_type(
        collections.deque, "collections.deque",
        custom_serializer=_serialize_deque,
        custom_deserializer=_deserialize_deque)

    def _serialize_ordered_dict(obj):
        return list(obj.keys()), list(obj.values())

    def _deserialize_ordered_dict(data):
        return collections.OrderedDict(zip(data[0], data[1]))

    serialization_context.register_type(
        collections.OrderedDict, "collections.OrderedDict",
        custom_serializer=_serialize_ordered_dict,
        custom_deserializer=_deserialize_ordered_dict)

    def _serialize_default_dict(obj):
        return list(obj.keys()), list(obj.values()), obj.default_factory

    def _deserialize_default_dict(data):
        return collections.defaultdict(data[2], zip(data[0], data[1]))

    serialization_context.register_type(
        collections.defaultdict, "collections.defaultdict",
        custom_serializer=_serialize_default_dict,
        custom_deserializer=_deserialize_default_dict)

    def _serialize_counter(obj):
        return list(obj.keys()), list(obj.values())

    def _deserialize_counter(data):
        return collections.Counter(dict(zip(data[0], data[1])))

    serialization_context.register_type(
        collections.Counter, "collections.Counter",
        custom_serializer=_serialize_counter,
        custom_deserializer=_deserialize_counter)


# ----------------------------------------------------------------------
# Set up serialization for scipy sparse matrices. Primitive types are handled
# efficiently with Arrow's SparseTensor facilities, see numpy_convert.cc)

def _register_scipy_handlers(serialization_context):
    try:
        from scipy.sparse import (csr_matrix, csc_matrix, coo_matrix,
                                  isspmatrix_coo, isspmatrix_csr,
                                  isspmatrix_csc, isspmatrix)

        def _serialize_scipy_sparse(obj):
            if isspmatrix_coo(obj):
                return 'coo', pa.SparseCOOTensor.from_scipy(obj)

            elif isspmatrix_csr(obj):
                return 'csr', pa.SparseCSRMatrix.from_scipy(obj)

            elif isspmatrix_csc(obj):
                return 'csc', pa.SparseCSCMatrix.from_scipy(obj)

            elif isspmatrix(obj):
                return 'csr', pa.SparseCOOTensor.from_scipy(obj.to_coo())

            else:
                raise NotImplementedError(
                    "Serialization of {} is not supported.".format(obj[0]))

        def _deserialize_scipy_sparse(data):
            if data[0] == 'coo':
                return data[1].to_scipy()

            elif data[0] == 'csr':
                return data[1].to_scipy()

            elif data[0] == 'csc':
                return data[1].to_scipy()

            else:
                return data[1].to_scipy()

        serialization_context.register_type(
            coo_matrix, 'scipy.sparse.coo.coo_matrix',
            custom_serializer=_serialize_scipy_sparse,
            custom_deserializer=_deserialize_scipy_sparse)

        serialization_context.register_type(
            csr_matrix, 'scipy.sparse.csr.csr_matrix',
            custom_serializer=_serialize_scipy_sparse,
            custom_deserializer=_deserialize_scipy_sparse)

        serialization_context.register_type(
            csc_matrix, 'scipy.sparse.csc.csc_matrix',
            custom_serializer=_serialize_scipy_sparse,
            custom_deserializer=_deserialize_scipy_sparse)

    except ImportError:
        # no scipy
        pass


# ----------------------------------------------------------------------
# Set up serialization for pydata/sparse tensors.

def _register_pydata_sparse_handlers(serialization_context):
    try:
        import sparse

        def _serialize_pydata_sparse(obj):
            if isinstance(obj, sparse.COO):
                return 'coo', pa.SparseCOOTensor.from_pydata_sparse(obj)
            else:
                raise NotImplementedError(
                    "Serialization of {} is not supported.".format(sparse.COO))

        def _deserialize_pydata_sparse(data):
            if data[0] == 'coo':
                data_array, coords = data[1].to_numpy()
                return sparse.COO(
                    data=data_array[:, 0],
                    coords=coords.T, shape=data[1].shape)

        serialization_context.register_type(
            sparse.COO, 'sparse.COO',
            custom_serializer=_serialize_pydata_sparse,
            custom_deserializer=_deserialize_pydata_sparse)

    except ImportError:
        # no pydata/sparse
        pass


def _register_default_serialization_handlers(serialization_context):

    # ----------------------------------------------------------------------
    # Set up serialization for primitive datatypes

    # TODO(pcm): This is currently a workaround until arrow supports
    # arbitrary precision integers. This is only called on long integers,
    # see the associated case in the append method in python_to_arrow.cc
    serialization_context.register_type(
        int, "int",
        custom_serializer=lambda obj: str(obj),
        custom_deserializer=lambda data: int(data))

    serialization_context.register_type(
        type(lambda: 0), "function",
        pickle=True)

    serialization_context.register_type(type, "type", pickle=True)

    serialization_context.register_type(
        np.matrix, 'np.matrix',
        custom_serializer=_serialize_numpy_matrix,
        custom_deserializer=_deserialize_numpy_matrix)

    serialization_context.register_type(
        np.ndarray, 'np.array',
        custom_serializer=_serialize_numpy_array_list,
        custom_deserializer=_deserialize_numpy_array_list)

    serialization_context.register_type(
        pa.Array, 'pyarrow.Array',
        custom_serializer=_serialize_pyarrow_array,
        custom_deserializer=_deserialize_pyarrow_array)

    serialization_context.register_type(
        pa.RecordBatch, 'pyarrow.RecordBatch',
        custom_serializer=_serialize_pyarrow_recordbatch,
        custom_deserializer=_deserialize_pyarrow_recordbatch)

    serialization_context.register_type(
        pa.Table, 'pyarrow.Table',
        custom_serializer=_serialize_pyarrow_table,
        custom_deserializer=_deserialize_pyarrow_table)

    _register_collections_serialization_handlers(serialization_context)
    _register_custom_pandas_handlers(serialization_context)
    _register_scipy_handlers(serialization_context)
    _register_pydata_sparse_handlers(serialization_context)


def register_default_serialization_handlers(serialization_context):
    _deprecate_serialization("register_default_serialization_handlers")
    _register_default_serialization_handlers(serialization_context)


def default_serialization_context():
    _deprecate_serialization("default_serialization_context")
    context = SerializationContext()
    _register_default_serialization_handlers(context)
    return context