import pandas
import pytest
import raydp

import ray
from ray.data.tests.conftest import *  # noqa
from ray.data.tests.test_util import _check_usage_record


# RayDP tests require Ray Java. Make sure ray jar is built before running this test.
@pytest.fixture(scope="function")
def spark(request):
    ray.init(num_cpus=2, include_dashboard=False)
    spark_session = raydp.init_spark("test", 1, 1, "500M")

    def stop_all():
        raydp.stop_spark()
        ray.shutdown()

    request.addfinalizer(stop_all)
    return spark_session


def test_raydp_roundtrip(spark):
    spark_df = spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")], ["one", "two"])
    rows = [(r.one, r.two) for r in spark_df.take(3)]
    ds = ray.data.from_spark(spark_df)
    values = [(r["one"], r["two"]) for r in ds.take(6)]
    assert values == rows
    df = ds.to_spark(spark)
    rows_2 = [(r.one, r.two) for r in df.take(3)]
    assert values == rows_2


def test_raydp_to_spark(spark):
    n = 5
    ds = ray.data.range(n)
    values = [r["id"] for r in ds.take(5)]
    df = ds.to_spark(spark)
    rows = [r.id for r in df.take(5)]
    assert values == rows


def test_from_spark_e2e(spark):
    spark_df = spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")], ["one", "two"])

    rows = [(r.one, r.two) for r in spark_df.take(3)]
    ds = ray.data.from_spark(spark_df)
    assert len(ds.take_all()) == len(rows)
    values = [(r["one"], r["two"]) for r in ds.take(6)]
    assert values == rows

    # Check that metadata fetch is included in stats.
    assert "FromArrow" in ds.stats()
    # Underlying implementation uses `FromArrow` operator
    assert ds._logical_plan.dag.name == "FromArrow"
    _check_usage_record(["FromArrow"])


def test_to_pandas(spark):
    df = spark.range(100)
    ds = ray.data.from_spark(df)
    pdf = ds.to_pandas()
    pdf2 = df.toPandas().astype(pdf.dtypes.to_dict())
    pandas.testing.assert_frame_equal(pdf, pdf2)


if __name__ == "__main__":
    import sys

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