# Copyright (c) Facebook, Inc. and its affiliates.
from typing import List, Optional, cast

# Skipping analyzing 'numpy': found module but no type hints or library stubs
import numpy as np  # type: ignore
import numpy.ma as ma  # type: ignore

# Skipping analyzing 'pandas': found module but no type hints or library stubs
import pandas as pd  # type: ignore
import pyarrow as pa  # type: ignore
import torcharrow.dtypes as dt
from torcharrow import Scope


def from_arrow_table(
    table,
    dtype: Optional[dt.DType] = None,
    columns: Optional[List[str]] = None,
    scope=None,
    device="",
):
    """ "
    Convert arrow table to a torcharrow dataframe.
    """
    scope = scope or Scope.default
    device = device or scope.device
    assert isinstance(table, pa.Table)
    if dtype is not None:
        assert dt.is_struct(dtype)
        dtype = cast(dt.Struct, dtype)
        res = {}
        for f in dtype.fields:
            chunked_array = table.column(f.name)
            pydata = chunked_array.to_pylist()
            res[f.name] = scope.Column(pydata, f.dtype)
        return scope.DataFrame(res, device=device)
    else:
        res = {}
        table = table.select(columns) if columns is not None else table
        for n in table.column_names:
            chunked_array = table.column(n)
            pydata = chunked_array.to_pylist()
            res[n] = scope.Column(
                pydata,
                dtype=_arrowtype_to_dtype(
                    table.schema.field(n).type, table.column(n).null_count > 0
                ),
            )
        return scope.DataFrame(res, device=device)


def from_pandas_dataframe(
    df,
    dtype: Optional[dt.DType] = None,
    columns: Optional[List[str]] = None,
    scope=None,
    device="",
):
    """
    Convert pandas dataframe to torcharrow dataframe (drops indices).

    Parameters
    ----------
    df : Pandas dataframe

    dtype : dtype, default None
    Data type to force, if None will automatically infer.

    columns : array-like
    List of column names to extract from df.

    scope : Scope or None
    Scope to use, or None for default scope.

    device : str or ""
    Device to use, or default if blank.

    Examples
    --------
    >>> import pandas as pd
    >>> import torcharrow as ta
    >>> pdf = pd.DataFrame({'a': [0, 1, 2, 3],'b': [0.1, 0.2, None, 0.3]})
    >>> gdf = ta.from_pandas_dataframe(pdf)
    >>> gdf
      index    a    b
    -------  ---  ---
          0    0  0.1
          1    1  0.2
          2    2
          3    3  0.3
    dtype: Struct([Field('a', int64), Field('b', Float64(nullable=True))]), count: 4, null_count: 0

    """
    scope = scope or Scope.default
    device = device or scope.device

    if dtype is not None:
        assert dt.is_struct(dtype)
        dtype = cast(dt.Struct, dtype)
        res = {}
        for f in dtype.fields:
            # this shows that Column shoud also construct Dataframes!
            res[f.name] = from_pandas_series(
                pd.Series(df[f.name]), f.dtype, scope=scope
            )
        return scope.Frame(res, dtype=dtype, device=device)
    else:
        res = {}
        for n in df.columns:
            if columns is None or n in columns:
                res[n] = from_pandas_series(pd.Series(df[n]), scope=scope)
        return scope.Frame(res, device=device)


def from_arrow_array(array, dtype=None, scope=None, device=""):
    """ "
    Convert arrow array to a torcharrow column.
    """
    scope = scope or Scope.default
    device = device or scope.device
    assert isinstance(array, pa.Array)
    pydata = _arrow_scalar_to_py(array)
    if dtype is not None:
        assert not dt.is_struct(dtype)
        return scope.Column(pydata, dtype, device=device)
    else:
        return scope.Column(
            pydata,
            dtype=_arrowtype_to_dtype(array.type, array.null_count > 0),
            device=device,
        )


def from_pandas_series(series, dtype=None, scope=None, device=""):
    """ "
    Convert pandas series array to a torcharrow column (drops indices).
    """
    scope = scope or Scope.default
    device = device or scope.device

    return from_numpy(series.to_numpy(), dtype, scope, device)


def from_numpy(array, dtype, scope=None, device=""):
    """
    Convert 1dim numpy array to a torcharrow column (zero copy).
    """
    scope = scope or Scope.default
    device = device or scope.device

    if isinstance(array, ma.core.MaskedArray) and array.ndim == 1:
        return _from_numpy_ma(array.data, array.mask, dtype, scope, device)
    elif isinstance(array, np.ndarray) and array.ndim == 1:
        return _from_numpy_nd(array, dtype, scope, device)
    else:
        raise TypeError(f"cannot convert numpy array of type {array.dtype}")


def _is_not_str(s):
    return not isinstance(s, str)


def _from_numpy_ma(data, mask, dtype, scope=None, device=""):
    # adopt types
    if dtype is None:
        dtype = dt.typeof_np_dtype(data.dtype).with_null()
    else:
        assert dt.is_primitive_type(dtype)
        assert dtype == dt.typeof_np_dtype(data.dtype).with_null()
        # TODO if not, adopt the type or?
        # Something like ma.array
        # np.array([np.nan, np.nan,  3.]).astype(np.int64),
        # mask = np.isnan([np.nan, np.nan,  3.]))

    # create column, only zero copy supported
    if dt.is_boolean_or_numerical(dtype):
        assert not np.all(np.isnan(ma.array(data, mask).compressed()))

        return scope._FullColumn(data, dtype=dtype, mask=mask)
    elif dt.is_string(dtype) or dtype == "object":
        assert np.all(np.vectorize(_is_not_str)(ma.array(data, mask).compressed()))
        return scope._FullColumn(data, dtype=dtype, mask=mask)
    else:
        raise TypeError(f"cannot convert masked numpy array of type {data.dtype}")


def _from_numpy_nd(data, dtype, scope=None, device=""):
    # adopt types
    if dtype is None:
        dtype = dt.typeof_np_dtype(data.dtype)
        if dtype is None:
            dtype = dt.string
    else:
        assert dt.is_primitive(dtype)
        # TODO Check why teh following assert  isn't the case
        # assert dtype == dt.typeof_np_dtype(data.dtype)

    # create column, only zero copy supported
    if dt.is_boolean_or_numerical(dtype):
        mask = np.isnan(data)
        return scope._FullColumn(data, dtype=dtype, mask=mask)
    elif dt.is_string(dtype):
        mask = np.vectorize(_is_not_str)(data)
        if np.any(mask):
            dtype = dtype.with_null()
        return scope._FullColumn(data, dtype=dtype, mask=mask)
    else:
        raise TypeError("can not convert numpy array of type {data.dtype,}")


# def _column_without_nan(series, dtype):
#     if dtype is None or is_floating(dtype):
#         for i in series:
#             if isinstance(i, float) and np.isnan(i):
#                 yield None
#             else:
#                 yield i
#     else:
#         for i in series:
#             yield i


def _arrow_scalar_to_py(array):
    for i in array:
        yield i.as_py()


def _pandatype_to_dtype(t, nullable):
    return dt.typeof_nptype(t, nullable)


def _arrowtype_to_dtype(t, nullable):
    if pa.types.is_boolean(t):
        return dt.Boolean(nullable)
    if pa.types.is_int8(t):
        return dt.Int8(nullable)
    if pa.types.is_int16(t):
        return dt.Int16(nullable)
    if pa.types.is_int32(t):
        return dt.Int32(nullable)
    if pa.types.is_int64(t):
        return dt.Int64(nullable)
    if pa.types.is_float32(t):
        return dt.Float32(nullable)
    if pa.types.is_float64(t):
        return dt.Float64(nullable)
    if pa.types.is_list(t):
        return List(t.value_type, nullable)
    if pa.types.is_struct(t):
        return _pandatype_to_dtype(t.to_pandas_dtype(), True)
    if pa.types.is_null(t):
        return dt.Void()
    if pa.types.is_string(t):
        return dt.String(nullable)
    if pa.types.is_map(t):
        return dt.Map(t.item_type, t.key_type, nullable)
    raise NotImplementedError("unsupported case")

[((176, 7, 176, 40), 'torcharrow.dtypes.is_boolean_or_numerical', 'dt.is_boolean_or_numerical', ({(176, 34, 176, 39): 'dtype'}, {}), '(dtype)', True, 'import torcharrow.dtypes as dt\n'), ((199, 7, 199, 40), 'torcharrow.dtypes.is_boolean_or_numerical', 'dt.is_boolean_or_numerical', ({(199, 34, 199, 39): 'dtype'}, {}), '(dtype)', True, 'import torcharrow.dtypes as dt\n'), ((229, 11, 229, 40), 'torcharrow.dtypes.typeof_nptype', 'dt.typeof_nptype', ({(229, 28, 229, 29): 't', (229, 31, 229, 39): 'nullable'}, {}), '(t, nullable)', True, 'import torcharrow.dtypes as dt\n'), ((233, 7, 233, 29), 'pyarrow.types.is_boolean', 'pa.types.is_boolean', ({(233, 27, 233, 28): 't'}, {}), '(t)', True, 'import pyarrow as pa\n'), ((235, 7, 235, 26), 'pyarrow.types.is_int8', 'pa.types.is_int8', ({(235, 24, 235, 25): 't'}, {}), '(t)', True, 'import pyarrow as pa\n'), ((237, 7, 237, 27), 'pyarrow.types.is_int16', 'pa.types.is_int16', ({(237, 25, 237, 26): 't'}, {}), '(t)', True, 'import pyarrow as pa\n'), ((239, 7, 239, 27), 'pyarrow.types.is_int32', 'pa.types.is_int32', ({(239, 25, 239, 26): 't'}, {}), '(t)', True, 'import pyarrow as pa\n'), ((241, 7, 241, 27), 'pyarrow.types.is_int64', 'pa.types.is_int64', ({(241, 25, 241, 26): 't'}, {}), '(t)', True, 'import pyarrow as pa\n'), ((243, 7, 243, 29), 'pyarrow.types.is_float32', 'pa.types.is_float32', ({(243, 27, 243, 28): 't'}, {}), '(t)', True, 'import pyarrow as pa\n'), ((245, 7, 245, 29), 'pyarrow.types.is_float64', 'pa.types.is_float64', ({(245, 27, 245, 28): 't'}, {}), '(t)', True, 'import pyarrow as pa\n'), ((247, 7, 247, 26), 'pyarrow.types.is_list', 'pa.types.is_list', ({(247, 24, 247, 25): 't'}, {}), '(t)', True, 'import pyarrow as pa\n'), ((249, 7, 249, 28), 'pyarrow.types.is_struct', 'pa.types.is_struct', ({(249, 26, 249, 27): 't'}, {}), '(t)', True, 'import pyarrow as pa\n'), ((251, 7, 251, 26), 'pyarrow.types.is_null', 'pa.types.is_null', ({(251, 24, 251, 25): 't'}, {}), '(t)', True, 'import pyarrow as pa\n'), ((253, 7, 253, 28), 'pyarrow.types.is_string', 'pa.types.is_string', ({(253, 26, 253, 27): 't'}, {}), '(t)', True, 'import pyarrow as pa\n'), ((255, 7, 255, 25), 'pyarrow.types.is_map', 'pa.types.is_map', ({(255, 23, 255, 24): 't'}, {}), '(t)', True, 'import pyarrow as pa\n'), ((29, 15, 29, 34), 'torcharrow.dtypes.is_struct', 'dt.is_struct', ({(29, 28, 29, 33): 'dtype'}, {}), '(dtype)', True, 'import torcharrow.dtypes as dt\n'), ((30, 16, 30, 38), 'typing.cast', 'cast', ({(30, 21, 30, 30): 'dt.Struct', (30, 32, 30, 37): 'dtype'}, {}), '(dt.Struct, dtype)', False, 'from typing import List, Optional, cast\n'), ((98, 15, 98, 34), 'torcharrow.dtypes.is_struct', 'dt.is_struct', ({(98, 28, 98, 33): 'dtype'}, {}), '(dtype)', True, 'import torcharrow.dtypes as dt\n'), ((99, 16, 99, 38), 'typing.cast', 'cast', ({(99, 21, 99, 30): 'dt.Struct', (99, 32, 99, 37): 'dtype'}, {}), '(dt.Struct, dtype)', False, 'from typing import List, Optional, cast\n'), ((168, 15, 168, 42), 'torcharrow.dtypes.is_primitive_type', 'dt.is_primitive_type', ({(168, 36, 168, 41): 'dtype'}, {}), '(dtype)', True, 'import torcharrow.dtypes as dt\n'), ((190, 16, 190, 46), 'torcharrow.dtypes.typeof_np_dtype', 'dt.typeof_np_dtype', ({(190, 35, 190, 45): 'data.dtype'}, {}), '(data.dtype)', True, 'import torcharrow.dtypes as dt\n'), ((194, 15, 194, 37), 'torcharrow.dtypes.is_primitive', 'dt.is_primitive', ({(194, 31, 194, 36): 'dtype'}, {}), '(dtype)', True, 'import torcharrow.dtypes as dt\n'), ((200, 15, 200, 29), 'numpy.isnan', 'np.isnan', ({(200, 24, 200, 28): 'data'}, {}), '(data)', True, 'import numpy as np\n'), ((202, 9, 202, 28), 'torcharrow.dtypes.is_string', 'dt.is_string', ({(202, 22, 202, 27): 'dtype'}, {}), '(dtype)', True, 'import torcharrow.dtypes as dt\n'), ((234, 15, 234, 35), 'torcharrow.dtypes.Boolean', 'dt.Boolean', ({(234, 26, 234, 34): 'nullable'}, {}), '(nullable)', True, 'import torcharrow.dtypes as dt\n'), ((236, 15, 236, 32), 'torcharrow.dtypes.Int8', 'dt.Int8', ({(236, 23, 236, 31): 'nullable'}, {}), '(nullable)', True, 'import torcharrow.dtypes as dt\n'), ((238, 15, 238, 33), 'torcharrow.dtypes.Int16', 'dt.Int16', ({(238, 24, 238, 32): 'nullable'}, {}), '(nullable)', True, 'import torcharrow.dtypes as dt\n'), ((240, 15, 240, 33), 'torcharrow.dtypes.Int32', 'dt.Int32', ({(240, 24, 240, 32): 'nullable'}, {}), '(nullable)', True, 'import torcharrow.dtypes as dt\n'), ((242, 15, 242, 33), 'torcharrow.dtypes.Int64', 'dt.Int64', ({(242, 24, 242, 32): 'nullable'}, {}), '(nullable)', True, 'import torcharrow.dtypes as dt\n'), ((244, 15, 244, 35), 'torcharrow.dtypes.Float32', 'dt.Float32', ({(244, 26, 244, 34): 'nullable'}, {}), '(nullable)', True, 'import torcharrow.dtypes as dt\n'), ((246, 15, 246, 35), 'torcharrow.dtypes.Float64', 'dt.Float64', ({(246, 26, 246, 34): 'nullable'}, {}), '(nullable)', True, 'import torcharrow.dtypes as dt\n'), ((248, 15, 248, 43), 'typing.List', 'List', ({(248, 20, 248, 32): 't.value_type', (248, 34, 248, 42): 'nullable'}, {}), '(t.value_type, nullable)', False, 'from typing import List, Optional, cast\n'), ((252, 15, 252, 24), 'torcharrow.dtypes.Void', 'dt.Void', ({}, {}), '()', True, 'import torcharrow.dtypes as dt\n'), ((254, 15, 254, 34), 'torcharrow.dtypes.String', 'dt.String', ({(254, 25, 254, 33): 'nullable'}, {}), '(nullable)', True, 'import torcharrow.dtypes as dt\n'), ((256, 15, 256, 56), 'torcharrow.dtypes.Map', 'dt.Map', ({(256, 22, 256, 33): 't.item_type', (256, 35, 256, 45): 't.key_type', (256, 47, 256, 55): 'nullable'}, {}), '(t.item_type, t.key_type, nullable)', True, 'import torcharrow.dtypes as dt\n'), ((124, 19, 124, 38), 'torcharrow.dtypes.is_struct', 'dt.is_struct', ({(124, 32, 124, 37): 'dtype'}, {}), '(dtype)', True, 'import torcharrow.dtypes as dt\n'), ((180, 9, 180, 28), 'torcharrow.dtypes.is_string', 'dt.is_string', ({(180, 22, 180, 27): 'dtype'}, {}), '(dtype)', True, 'import torcharrow.dtypes as dt\n'), ((204, 11, 204, 23), 'numpy.any', 'np.any', ({(204, 18, 204, 22): 'mask'}, {}), '(mask)', True, 'import numpy as np\n'), ((104, 16, 104, 37), 'pandas.Series', 'pd.Series', ({(104, 26, 104, 36): 'df[f.name]'}, {}), '(df[f.name])', True, 'import pandas as pd\n'), ((166, 16, 166, 46), 'torcharrow.dtypes.typeof_np_dtype', 'dt.typeof_np_dtype', ({(166, 35, 166, 45): 'data.dtype'}, {}), '(data.dtype)', True, 'import torcharrow.dtypes as dt\n'), ((203, 15, 203, 40), 'numpy.vectorize', 'np.vectorize', ({(203, 28, 203, 39): '_is_not_str'}, {}), '(_is_not_str)', True, 'import numpy as np\n'), ((111, 44, 111, 60), 'pandas.Series', 'pd.Series', ({(111, 54, 111, 59): 'df[n]'}, {}), '(df[n])', True, 'import pandas as pd\n'), ((169, 24, 169, 54), 'torcharrow.dtypes.typeof_np_dtype', 'dt.typeof_np_dtype', ({(169, 43, 169, 53): 'data.dtype'}, {}), '(data.dtype)', True, 'import torcharrow.dtypes as dt\n'), ((181, 22, 181, 47), 'numpy.vectorize', 'np.vectorize', ({(181, 35, 181, 46): '_is_not_str'}, {}), '(_is_not_str)', True, 'import numpy as np\n'), ((177, 35, 177, 55), 'numpy.ma.array', 'ma.array', ({(177, 44, 177, 48): 'data', (177, 50, 177, 54): 'mask'}, {}), '(data, mask)', True, 'import numpy.ma as ma\n'), ((181, 48, 181, 68), 'numpy.ma.array', 'ma.array', ({(181, 57, 181, 61): 'data', (181, 63, 181, 67): 'mask'}, {}), '(data, mask)', True, 'import numpy.ma as ma\n')]