import os

import numpy as np
import pandas as pd
import pyarrow as pa
import pytest

import ray
from ray.data.context import DataContext
from ray.data.dataset import Schema
from ray.data.tests.conftest import *  # noqa
from ray.data.tests.util import extract_values
from ray.tests.conftest import *  # noqa


@pytest.mark.parametrize("from_ref", [False, True])
def test_from_numpy(ray_start_regular_shared, from_ref):
    arr1 = np.expand_dims(np.arange(0, 4), axis=1)
    arr2 = np.expand_dims(np.arange(4, 8), axis=1)
    arrs = [arr1, arr2]
    if from_ref:
        ds = ray.data.from_numpy_refs([ray.put(arr) for arr in arrs])
    else:
        ds = ray.data.from_numpy(arrs)
    values = np.stack(extract_values("data", ds.take(8)))
    np.testing.assert_array_equal(values, np.concatenate((arr1, arr2)))
    # Check that conversion task is included in stats.
    assert "FromNumpy" in ds.stats()

    # Test from single NumPy ndarray.
    if from_ref:
        ds = ray.data.from_numpy_refs(ray.put(arr1))
    else:
        ds = ray.data.from_numpy(arr1)
    values = np.stack(extract_values("data", ds.take(4)))
    np.testing.assert_array_equal(values, arr1)
    # Check that conversion task is included in stats.
    assert "FromNumpy" in ds.stats()


def test_from_numpy_variable_shaped(ray_start_regular_shared):
    arr = np.array([np.ones((2, 2)), np.ones((3, 3))], dtype=object)
    ds = ray.data.from_numpy(arr)
    values = np.array(extract_values("data", ds.take(2)), dtype=object)

    def recursive_to_list(a):
        if not isinstance(a, (list, np.ndarray)):
            return a
        return [recursive_to_list(e) for e in a]

    # Convert to a nested Python list in order to circumvent failed comparisons on
    # ndarray raggedness.
    np.testing.assert_equal(recursive_to_list(values), recursive_to_list(arr))


def test_to_numpy_refs(ray_start_regular_shared):
    # Tensor Dataset
    ds = ray.data.range_tensor(10, override_num_blocks=2)
    arr = np.concatenate(extract_values("data", ray.get(ds.to_numpy_refs())))
    np.testing.assert_equal(arr, np.expand_dims(np.arange(0, 10), 1))

    # Table Dataset
    ds = ray.data.range(10)
    arr = np.concatenate([t["id"] for t in ray.get(ds.to_numpy_refs())])
    np.testing.assert_equal(arr, np.arange(0, 10))

    # Test multi-column Arrow dataset.
    ds = ray.data.from_arrow(pa.table({"a": [1, 2, 3], "b": [4, 5, 6]}))
    arrs = ray.get(ds.to_numpy_refs())
    np.testing.assert_equal(
        arrs, [{"a": np.array([1, 2, 3]), "b": np.array([4, 5, 6])}]
    )

    # Test multi-column Pandas dataset.
    ds = ray.data.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}))
    arrs = ray.get(ds.to_numpy_refs())
    np.testing.assert_equal(
        arrs, [{"a": np.array([1, 2, 3]), "b": np.array([4, 5, 6])}]
    )


def test_numpy_roundtrip(ray_start_regular_shared, tmp_path):
    tensor_type = DataContext.get_current().arrow_fixed_shape_tensor_format.to_type()

    ds = ray.data.range_tensor(10, override_num_blocks=2)
    ds.write_numpy(tmp_path, column="data")
    ds = ray.data.read_numpy(tmp_path)
    assert ds.count() == 10
    assert ds.schema() == Schema(pa.schema([("data", tensor_type((1,), pa.int64()))]))
    assert sorted(ds.take_all(), key=lambda row: row["data"]) == [
        {"data": np.array([i])} for i in range(10)
    ]


def test_numpy_read_x(ray_start_regular_shared, tmp_path):
    tensor_type = DataContext.get_current().arrow_fixed_shape_tensor_format.to_type()

    path = os.path.join(tmp_path, "test_np_dir")
    os.mkdir(path)
    np.save(os.path.join(path, "test.npy"), np.expand_dims(np.arange(0, 10), 1))
    ds = ray.data.read_numpy(path, override_num_blocks=1)
    assert ds.count() == 10
    assert ds.schema() == Schema(pa.schema([("data", tensor_type((1,), pa.int64()))]))
    np.testing.assert_equal(
        extract_values("data", ds.take(2)), [np.array([0]), np.array([1])]
    )


def test_numpy_write(ray_start_regular_shared, tmp_path):
    ds = ray.data.range_tensor(1)

    ds.write_numpy(tmp_path, column="data")

    actual_array = np.concatenate(
        [np.load(os.path.join(tmp_path, filename)) for filename in os.listdir(tmp_path)]
    )
    assert actual_array == np.array((0,))


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))