# 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, 'torcharrow.dtypes.is_boolean_or_numerical', 'dt.is_boolean_or_numerical', (['dtype'], {}), True, 'import torcharrow.dtypes as dt\n'), (199, 'torcharrow.dtypes.is_boolean_or_numerical', 'dt.is_boolean_or_numerical', (['dtype'], {}), True, 'import torcharrow.dtypes as dt\n'), (229, 'torcharrow.dtypes.typeof_nptype', 'dt.typeof_nptype', (['t', 'nullable'], {}), True, 'import torcharrow.dtypes as dt\n'), (233, 'pyarrow.types.is_boolean', 'pa.types.is_boolean', (['t'], {}), True, 'import pyarrow as pa\n'), (235, 'pyarrow.types.is_int8', 'pa.types.is_int8', (['t'], {}), True, 'import pyarrow as pa\n'), (237, 'pyarrow.types.is_int16', 'pa.types.is_int16', (['t'], {}), True, 'import pyarrow as pa\n'), (239, 'pyarrow.types.is_int32', 'pa.types.is_int32', (['t'], {}), True, 'import pyarrow as pa\n'), (241, 'pyarrow.types.is_int64', 'pa.types.is_int64', (['t'], {}), True, 'import pyarrow as pa\n'), (243, 'pyarrow.types.is_float32', 'pa.types.is_float32', (['t'], {}), True, 'import pyarrow as pa\n'), (245, 'pyarrow.types.is_float64', 'pa.types.is_float64', (['t'], {}), True, 'import pyarrow as pa\n'), (247, 'pyarrow.types.is_list', 'pa.types.is_list', (['t'], {}), True, 'import pyarrow as pa\n'), (249, 'pyarrow.types.is_struct', 'pa.types.is_struct', (['t'], {}), True, 'import pyarrow as pa\n'), (251, 'pyarrow.types.is_null', 'pa.types.is_null', (['t'], {}), True, 'import pyarrow as pa\n'), (253, 'pyarrow.types.is_string', 'pa.types.is_string', (['t'], {}), True, 'import pyarrow as pa\n'), (255, 'pyarrow.types.is_map', 'pa.types.is_map', (['t'], {}), True, 'import pyarrow as pa\n'), (29, 'torcharrow.dtypes.is_struct', 'dt.is_struct', (['dtype'], {}), True, 'import torcharrow.dtypes as dt\n'), (30, 'typing.cast', 'cast', (['dt.Struct', 'dtype'], {}), False, 'from typing import List, Optional, cast\n'), (98, 'torcharrow.dtypes.is_struct', 'dt.is_struct', (['dtype'], {}), True, 'import torcharrow.dtypes as dt\n'), (99, 'typing.cast', 'cast', (['dt.Struct', 'dtype'], {}), False, 'from typing import List, Optional, cast\n'), (168, 'torcharrow.dtypes.is_primitive_type', 'dt.is_primitive_type', (['dtype'], {}), True, 'import torcharrow.dtypes as dt\n'), (190, 'torcharrow.dtypes.typeof_np_dtype', 'dt.typeof_np_dtype', (['data.dtype'], {}), True, 'import torcharrow.dtypes as dt\n'), (194, 'torcharrow.dtypes.is_primitive', 'dt.is_primitive', (['dtype'], {}), True, 'import torcharrow.dtypes as dt\n'), (200, 'numpy.isnan', 'np.isnan', (['data'], {}), True, 'import numpy as np\n'), (202, 'torcharrow.dtypes.is_string', 'dt.is_string', (['dtype'], {}), True, 'import torcharrow.dtypes as dt\n'), (234, 'torcharrow.dtypes.Boolean', 'dt.Boolean', (['nullable'], {}), True, 'import torcharrow.dtypes as dt\n'), (236, 'torcharrow.dtypes.Int8', 'dt.Int8', (['nullable'], {}), True, 'import torcharrow.dtypes as dt\n'), (238, 'torcharrow.dtypes.Int16', 'dt.Int16', (['nullable'], {}), True, 'import torcharrow.dtypes as dt\n'), (240, 'torcharrow.dtypes.Int32', 'dt.Int32', (['nullable'], {}), True, 'import torcharrow.dtypes as dt\n'), (242, 'torcharrow.dtypes.Int64', 'dt.Int64', (['nullable'], {}), True, 'import torcharrow.dtypes as dt\n'), (244, 'torcharrow.dtypes.Float32', 'dt.Float32', (['nullable'], {}), True, 'import torcharrow.dtypes as dt\n'), (246, 'torcharrow.dtypes.Float64', 'dt.Float64', (['nullable'], {}), True, 'import torcharrow.dtypes as dt\n'), (248, 'typing.List', 'List', (['t.value_type', 'nullable'], {}), False, 'from typing import List, Optional, cast\n'), (252, 'torcharrow.dtypes.Void', 'dt.Void', ([], {}), True, 'import torcharrow.dtypes as dt\n'), (254, 'torcharrow.dtypes.String', 'dt.String', (['nullable'], {}), True, 'import torcharrow.dtypes as dt\n'), (256, 'torcharrow.dtypes.Map', 'dt.Map', (['t.item_type', 't.key_type', 'nullable'], {}), True, 'import torcharrow.dtypes as dt\n'), (124, 'torcharrow.dtypes.is_struct', 'dt.is_struct', (['dtype'], {}), True, 'import torcharrow.dtypes as dt\n'), (180, 'torcharrow.dtypes.is_string', 'dt.is_string', (['dtype'], {}), True, 'import torcharrow.dtypes as dt\n'), (204, 'numpy.any', 'np.any', (['mask'], {}), True, 'import numpy as np\n'), (104, 'pandas.Series', 'pd.Series', (['df[f.name]'], {}), True, 'import pandas as pd\n'), (166, 'torcharrow.dtypes.typeof_np_dtype', 'dt.typeof_np_dtype', (['data.dtype'], {}), True, 'import torcharrow.dtypes as dt\n'), (203, 'numpy.vectorize', 'np.vectorize', (['_is_not_str'], {}), True, 'import numpy as np\n'), (111, 'pandas.Series', 'pd.Series', (['df[n]'], {}), True, 'import pandas as pd\n'), (169, 'torcharrow.dtypes.typeof_np_dtype', 'dt.typeof_np_dtype', (['data.dtype'], {}), True, 'import torcharrow.dtypes as dt\n'), (181, 'numpy.vectorize', 'np.vectorize', (['_is_not_str'], {}), True, 'import numpy as np\n'), (177, 'numpy.ma.array', 'ma.array', (['data', 'mask'], {}), True, 'import numpy.ma as ma\n'), (181, 'numpy.ma.array', 'ma.array', (['data', 'mask'], {}), True, 'import numpy.ma as ma\n')]