from typing import TYPE_CHECKING, Tuple, Sequence, Optional, List

import numpy as np

if TYPE_CHECKING:
    from ..types import ArrayType
    from .. import Document, DocumentArray


def unravel(docs: Sequence['Document'], field: str) -> Optional['ArrayType']:
    _first = getattr(docs[0], field)
    if _first is None:
        # failed to unravel, return as a list
        r = [getattr(d, field) for d in docs]
        if any(_rr is not None for _rr in r):
            return r
        else:
            return None

    framework, is_sparse = get_array_type(_first)
    all_fields = [getattr(d, field) for d in docs]
    cls_type = type(_first)

    if framework == 'python':
        return cls_type(all_fields)

    elif framework == 'numpy':
        return np.stack(all_fields)

    elif framework == 'tensorflow':
        import tensorflow as tf

        return tf.stack(all_fields)

    elif framework == 'torch':
        import torch

        return torch.stack(all_fields)

    elif framework == 'paddle':
        import paddle

        return paddle.stack(all_fields)

    elif framework == 'scipy':
        import scipy.sparse

        return cls_type(scipy.sparse.vstack(all_fields))


def ravel(value: 'ArrayType', docs: 'DocumentArray', field: str) -> None:
    """Ravel :attr:`value` into ``doc.field`` of each documents

    :param docs: the docs to set
    :param field: the field of the doc to set
    :param value: the value to be set on ``doc.field``
    """
    use_get_row = False
    if hasattr(value, 'getformat'):
        # for scipy only
        sp_format = value.getformat()
        if sp_format in {'bsr', 'coo'}:
            # for BSR and COO, they dont implement [j, ...] in scipy
            # but they offer get_row() API which implicitly translate the
            # sparse row into CSR format, hence needs to convert back
            # not very efficient, but this is the best we can do.
            use_get_row = True

    if use_get_row:
        emb_shape0 = value.shape[0]
        for d, j in zip(docs, range(emb_shape0)):
            row = getattr(value.getrow(j), f'to{sp_format}')()
            docs[d.id, field] = row
    elif isinstance(value, (list, tuple)):
        for d, j in zip(docs, value):
            docs[d.id, field] = j
    else:

        emb_shape0 = value.shape[0]
        for d, j in zip(docs, range(emb_shape0)):
            docs[d.id, field] = value[j, ...]


def get_array_type(array: 'ArrayType') -> Tuple[str, bool]:
    """Get the type of ndarray without importing the framework

    :param array: any array, scipy, numpy, tf, torch, etc.
    :return: a tuple where the first element represents the framework, the second represents if it is sparse array
    """
    module_tags = array.__class__.__module__.split('.')
    class_name = array.__class__.__name__

    if isinstance(array, (list, tuple)):
        return 'python', False

    if 'numpy' in module_tags:
        return 'numpy', False

    if 'docarray' in module_tags:
        if class_name == 'NdArray':
            return 'docarray', False  # sparse or not is irrelevant

    if 'docarray_pb2' in module_tags:
        if class_name == 'NdArrayProto':
            return 'docarray_proto', False  # sparse or not is irrelevant

    if 'tensorflow' in module_tags:
        if class_name == 'SparseTensor':
            return 'tensorflow', True
        if class_name == 'Tensor' or class_name == 'EagerTensor':
            return 'tensorflow', False

    if 'torch' in module_tags and class_name == 'Tensor':
        return 'torch', array.is_sparse

    if 'paddle' in module_tags and class_name == 'Tensor':
        # Paddle does not support sparse tensor on 11/8/2021
        # https://github.com/PaddlePaddle/Paddle/issues/36697
        return 'paddle', False

    if 'scipy' in module_tags and 'sparse' in module_tags:
        return 'scipy', True

    raise TypeError(f'can not determine the array type: {module_tags}.{class_name}')


def to_numpy_array(value) -> 'np.ndarray':
    """Return the value always in :class:`numpy.ndarray` regardless the framework type.

    :return: the value in :class:`numpy.ndarray`.
    """
    v = value
    framework, is_sparse = get_array_type(value)
    if is_sparse:
        if hasattr(v, 'todense'):
            v = v.todense()
        elif hasattr(v, 'to_dense'):
            v = v.to_dense()
        elif framework == 'tensorflow':
            import tensorflow as tf

            if isinstance(v, tf.SparseTensor):
                v = tf.sparse.to_dense(v)

    if hasattr(v, 'numpy'):
        v = v.numpy()
    if framework == 'python':
        v = np.array(v)
    return v


def to_list(value) -> List[float]:
    r = to_numpy_array(value)
    if isinstance(r, np.ndarray):
        return r.tolist()
    elif isinstance(r, list):
        return r
    else:
        raise TypeError(f'{r} can not be converted into list')


def get_array_rows(array: 'ArrayType') -> Tuple[int, int]:
    """Get the number of rows of the ndarray without importing all frameworks

    :param array: input array
    :return: (num_rows, ndim)

    Examples

    >>> get_array_rows([1,2,3])
    1, 1
    >>> get_array_rows([[1,2,3], [4,5,6]])
    2, 2
    >>> get_array_rows([[1,2,3], [4,5,6], [7,8,9]])
    3, 2
    >>> get_array_rows(np.array([[1,2,3], [4,5,6], [7,8,9]]))
    3, 2
    """
    array_type, _ = get_array_type(array)

    if array_type == 'python':
        first_element_list_like = isinstance(array[0], (list, tuple))
        num_rows = len(array) if first_element_list_like else 1
        ndim = 2 if first_element_list_like else 1
    elif array_type in ('numpy', 'tensorflow', 'torch', 'paddle', 'scipy'):
        ndim = array.ndim
        if ndim == 1:
            num_rows = 1
        else:
            num_rows = array.shape[0]
    else:
        raise ValueError

    return num_rows, ndim