# 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")

[((5363, 5396), 'torcharrow.dtypes.is_boolean_or_numerical', 'dt.is_boolean_or_numerical', (['dtype'], {}), '(dtype)\n', (5389, 5396), True, 'import torcharrow.dtypes as dt\n'), ((6246, 6279), 'torcharrow.dtypes.is_boolean_or_numerical', 'dt.is_boolean_or_numerical', (['dtype'], {}), '(dtype)\n', (6272, 6279), True, 'import torcharrow.dtypes as dt\n'), ((7101, 7130), 'torcharrow.dtypes.typeof_nptype', 'dt.typeof_nptype', (['t', 'nullable'], {}), '(t, nullable)\n', (7117, 7130), True, 'import torcharrow.dtypes as dt\n'), ((7178, 7200), 'pyarrow.types.is_boolean', 'pa.types.is_boolean', (['t'], {}), '(t)\n', (7197, 7200), True, 'import pyarrow as pa\n'), ((7245, 7264), 'pyarrow.types.is_int8', 'pa.types.is_int8', (['t'], {}), '(t)\n', (7261, 7264), True, 'import pyarrow as pa\n'), ((7306, 7326), 'pyarrow.types.is_int16', 'pa.types.is_int16', (['t'], {}), '(t)\n', (7323, 7326), True, 'import pyarrow as pa\n'), ((7369, 7389), 'pyarrow.types.is_int32', 'pa.types.is_int32', (['t'], {}), '(t)\n', (7386, 7389), True, 'import pyarrow as pa\n'), ((7432, 7452), 'pyarrow.types.is_int64', 'pa.types.is_int64', (['t'], {}), '(t)\n', (7449, 7452), True, 'import pyarrow as pa\n'), ((7495, 7517), 'pyarrow.types.is_float32', 'pa.types.is_float32', (['t'], {}), '(t)\n', (7514, 7517), True, 'import pyarrow as pa\n'), ((7562, 7584), 'pyarrow.types.is_float64', 'pa.types.is_float64', (['t'], {}), '(t)\n', (7581, 7584), True, 'import pyarrow as pa\n'), ((7629, 7648), 'pyarrow.types.is_list', 'pa.types.is_list', (['t'], {}), '(t)\n', (7645, 7648), True, 'import pyarrow as pa\n'), ((7701, 7722), 'pyarrow.types.is_struct', 'pa.types.is_struct', (['t'], {}), '(t)\n', (7719, 7722), True, 'import pyarrow as pa\n'), ((7793, 7812), 'pyarrow.types.is_null', 'pa.types.is_null', (['t'], {}), '(t)\n', (7809, 7812), True, 'import pyarrow as pa\n'), ((7846, 7867), 'pyarrow.types.is_string', 'pa.types.is_string', (['t'], {}), '(t)\n', (7864, 7867), True, 'import pyarrow as pa\n'), ((7911, 7929), 'pyarrow.types.is_map', 'pa.types.is_map', (['t'], {}), '(t)\n', (7926, 7929), True, 'import pyarrow as pa\n'), ((824, 843), 'torcharrow.dtypes.is_struct', 'dt.is_struct', (['dtype'], {}), '(dtype)\n', (836, 843), True, 'import torcharrow.dtypes as dt\n'), ((860, 882), 'typing.cast', 'cast', (['dt.Struct', 'dtype'], {}), '(dt.Struct, dtype)\n', (864, 882), False, 'from typing import List, Optional, cast\n'), ((2776, 2795), 'torcharrow.dtypes.is_struct', 'dt.is_struct', (['dtype'], {}), '(dtype)\n', (2788, 2795), True, 'import torcharrow.dtypes as dt\n'), ((2812, 2834), 'typing.cast', 'cast', (['dt.Struct', 'dtype'], {}), '(dt.Struct, dtype)\n', (2816, 2834), False, 'from typing import List, Optional, cast\n'), ((5028, 5055), 'torcharrow.dtypes.is_primitive_type', 'dt.is_primitive_type', (['dtype'], {}), '(dtype)\n', (5048, 5055), True, 'import torcharrow.dtypes as dt\n'), ((5938, 5968), 'torcharrow.dtypes.typeof_np_dtype', 'dt.typeof_np_dtype', (['data.dtype'], {}), '(data.dtype)\n', (5956, 5968), True, 'import torcharrow.dtypes as dt\n'), ((6050, 6072), 'torcharrow.dtypes.is_primitive', 'dt.is_primitive', (['dtype'], {}), '(dtype)\n', (6065, 6072), True, 'import torcharrow.dtypes as dt\n'), ((6296, 6310), 'numpy.isnan', 'np.isnan', (['data'], {}), '(data)\n', (6304, 6310), True, 'import numpy as np\n'), ((6383, 6402), 'torcharrow.dtypes.is_string', 'dt.is_string', (['dtype'], {}), '(dtype)\n', (6395, 6402), True, 'import torcharrow.dtypes as dt\n'), ((7217, 7237), 'torcharrow.dtypes.Boolean', 'dt.Boolean', (['nullable'], {}), '(nullable)\n', (7227, 7237), True, 'import torcharrow.dtypes as dt\n'), ((7281, 7298), 'torcharrow.dtypes.Int8', 'dt.Int8', (['nullable'], {}), '(nullable)\n', (7288, 7298), True, 'import torcharrow.dtypes as dt\n'), ((7343, 7361), 'torcharrow.dtypes.Int16', 'dt.Int16', (['nullable'], {}), '(nullable)\n', (7351, 7361), True, 'import torcharrow.dtypes as dt\n'), ((7406, 7424), 'torcharrow.dtypes.Int32', 'dt.Int32', (['nullable'], {}), '(nullable)\n', (7414, 7424), True, 'import torcharrow.dtypes as dt\n'), ((7469, 7487), 'torcharrow.dtypes.Int64', 'dt.Int64', (['nullable'], {}), '(nullable)\n', (7477, 7487), True, 'import torcharrow.dtypes as dt\n'), ((7534, 7554), 'torcharrow.dtypes.Float32', 'dt.Float32', (['nullable'], {}), '(nullable)\n', (7544, 7554), True, 'import torcharrow.dtypes as dt\n'), ((7601, 7621), 'torcharrow.dtypes.Float64', 'dt.Float64', (['nullable'], {}), '(nullable)\n', (7611, 7621), True, 'import torcharrow.dtypes as dt\n'), ((7665, 7693), 'typing.List', 'List', (['t.value_type', 'nullable'], {}), '(t.value_type, nullable)\n', (7669, 7693), False, 'from typing import List, Optional, cast\n'), ((7829, 7838), 'torcharrow.dtypes.Void', 'dt.Void', ([], {}), '()\n', (7836, 7838), True, 'import torcharrow.dtypes as dt\n'), ((7884, 7903), 'torcharrow.dtypes.String', 'dt.String', (['nullable'], {}), '(nullable)\n', (7893, 7903), True, 'import torcharrow.dtypes as dt\n'), ((7946, 7987), 'torcharrow.dtypes.Map', 'dt.Map', (['t.item_type', 't.key_type', 'nullable'], {}), '(t.item_type, t.key_type, nullable)\n', (7952, 7987), True, 'import torcharrow.dtypes as dt\n'), ((3686, 3705), 'torcharrow.dtypes.is_struct', 'dt.is_struct', (['dtype'], {}), '(dtype)\n', (3698, 3705), True, 'import torcharrow.dtypes as dt\n'), ((5542, 5561), 'torcharrow.dtypes.is_string', 'dt.is_string', (['dtype'], {}), '(dtype)\n', (5554, 5561), True, 'import torcharrow.dtypes as dt\n'), ((6462, 6474), 'numpy.any', 'np.any', (['mask'], {}), '(mask)\n', (6468, 6474), True, 'import numpy as np\n'), ((3015, 3036), 'pandas.Series', 'pd.Series', (['df[f.name]'], {}), '(df[f.name])\n', (3024, 3036), True, 'import pandas as pd\n'), ((4960, 4990), 'torcharrow.dtypes.typeof_np_dtype', 'dt.typeof_np_dtype', (['data.dtype'], {}), '(data.dtype)\n', (4978, 4990), True, 'import torcharrow.dtypes as dt\n'), ((6419, 6444), 'numpy.vectorize', 'np.vectorize', (['_is_not_str'], {}), '(_is_not_str)\n', (6431, 6444), True, 'import numpy as np\n'), ((3281, 3297), 'pandas.Series', 'pd.Series', (['df[n]'], {}), '(df[n])\n', (3290, 3297), True, 'import pandas as pd\n'), ((5080, 5110), 'torcharrow.dtypes.typeof_np_dtype', 'dt.typeof_np_dtype', (['data.dtype'], {}), '(data.dtype)\n', (5098, 5110), True, 'import torcharrow.dtypes as dt\n'), ((5606, 5631), 'numpy.vectorize', 'np.vectorize', (['_is_not_str'], {}), '(_is_not_str)\n', (5618, 5631), True, 'import numpy as np\n'), ((5433, 5453), 'numpy.ma.array', 'ma.array', (['data', 'mask'], {}), '(data, mask)\n', (5441, 5453), True, 'import numpy.ma as ma\n'), ((5632, 5652), 'numpy.ma.array', 'ma.array', (['data', 'mask'], {}), '(data, mask)\n', (5640, 5652), True, 'import numpy.ma as ma\n')]