from __future__ import annotations

import os
import pickle

import numpy as np
import pandas as pd
import pytest
from pyarrow import fs

import dask
from dask._compatibility import WINDOWS
from dask.dataframe.dask_expr import from_array, from_graph, from_pandas, read_parquet
from dask.dataframe.dask_expr._expr import Filter, Lengths, Literal
from dask.dataframe.dask_expr._reductions import Len
from dask.dataframe.dask_expr.io import FusedParquetIO, ReadParquet
from dask.dataframe.dask_expr.io.parquet import (
    _aggregate_statistics_to_file,
    _combine_stats,
    _extract_stats,
)
from dask.dataframe.utils import assert_eq
from dask.utils import key_split


def _make_file(dir, df=None, filename="myfile.parquet", **kwargs):
    fn = os.path.join(str(dir), filename)
    if df is None:
        df = pd.DataFrame({c: range(10) for c in "abcde"})
    df.to_parquet(fn, **kwargs)
    return fn


@pytest.fixture
def parquet_file(tmpdir):
    return _make_file(tmpdir)


@pytest.fixture(params=["arrow", "fsspec"])
def filesystem(request):
    return request.param


def test_parquet_len(tmpdir, filesystem):
    df = read_parquet(_make_file(tmpdir), filesystem=filesystem)
    pdf = df.compute()

    assert len(df[df.a > 5]) == len(pdf[pdf.a > 5])

    s = (df["b"] + 1).astype("Int32")
    assert len(s) == len(pdf)

    assert isinstance(Len(s.expr).optimize(), Literal)
    assert isinstance(Lengths(s.expr).optimize(), Literal)


def test_parquet_missing_stats(tmpdir, filesystem):
    _make_file(tmpdir)
    _make_file(tmpdir, write_statistics=["a", "b"], filename="bla.parquet")

    result = read_parquet(tmpdir, filesystem=filesystem)
    expected = pd.concat(
        [
            pd.DataFrame({c: range(10) for c in "abcde"}),
            pd.DataFrame({c: range(10) for c in "abcde"}),
        ]
    )
    assert_eq(result, expected, check_index=False)


@pytest.mark.parametrize("val", [np.nan, 1])
def test_parquet_all_na_column(tmpdir, filesystem, val):
    pdf = pd.DataFrame({"a": [np.nan] * 299 + [val], "b": [1, 2, 3] * 100})
    _make_file(tmpdir, df=pdf, filename="bla.parquet", row_group_size=100)
    result = read_parquet(tmpdir, filesystem=filesystem)
    assert_eq(result, pdf)


def test_parquet_len_filter(tmpdir, filesystem):
    df = read_parquet(_make_file(tmpdir), filesystem=filesystem)
    expr = Len(df[df.c > 0].expr)
    result = expr.simplify()
    for rp in result.find_operations(ReadParquet):
        assert rp.operand("columns") == ["c"] or rp.operand("columns") == []


def test_parquet_len_empty_dir():
    assert len(read_parquet(["empty_dir/*.parquet"])) == 0


@pytest.mark.skipif(WINDOWS, reason="directory empty on windows")
@pytest.mark.parametrize("write_metadata_file", [True, False])
def test_to_parquet(tmpdir, write_metadata_file):
    pdf = pd.DataFrame({"x": [1, 4, 3, 2, 0, 5]})
    df = from_pandas(pdf, npartitions=2)

    # Check basic parquet round trip
    df.to_parquet(tmpdir, write_metadata_file=write_metadata_file)
    df2 = read_parquet(tmpdir, calculate_divisions=True)
    assert_eq(df, df2)

    # Check overwrite behavior
    df["new"] = df["x"] + 1
    df.to_parquet(tmpdir, overwrite=True, write_metadata_file=write_metadata_file)
    df2 = read_parquet(tmpdir, calculate_divisions=True)
    assert_eq(df, df2)

    # Check that we cannot overwrite a path we are
    # reading from in the same graph
    with pytest.raises(ValueError, match="Cannot overwrite"):
        df2.to_parquet(tmpdir, overwrite=True)


def test_pyarrow_filesystem(parquet_file):
    filesystem = fs.LocalFileSystem()

    df_pa = read_parquet(parquet_file, filesystem=filesystem)
    df = read_parquet(parquet_file)
    assert assert_eq(df, df_pa)
    assert all(tsk.data_producer for tsk in df.dask.values())
    assert all(tsk.data_producer for tsk in df.optimize().dask.values())


@pytest.mark.parametrize("dtype_backend", ["pyarrow", "numpy_nullable", None])
def test_pyarrow_filesystem_dtype_backend(parquet_file, dtype_backend):
    filesystem = fs.LocalFileSystem()

    df_pa = read_parquet(
        parquet_file, filesystem=filesystem, dtype_backend=dtype_backend
    )
    df = read_parquet(parquet_file, dtype_backend=dtype_backend)
    assert assert_eq(df, df_pa)


@pytest.mark.parametrize("types_mapper", [None, lambda x: None])
def test_pyarrow_filesystem_types_mapper(parquet_file, types_mapper):
    # This test isn't doing much other than ensuring the stuff is not raising
    # anywhere
    filesystem = fs.LocalFileSystem()

    df_pa = read_parquet(
        parquet_file,
        filesystem=filesystem,
        arrow_to_pandas={"types_mapper": types_mapper},
    )
    df = read_parquet(parquet_file, arrow_to_pandas={"types_mapper": types_mapper})
    assert assert_eq(df, df_pa)


def test_pyarrow_filesystem_serialize(parquet_file):
    filesystem = fs.LocalFileSystem()

    df_pa = read_parquet(parquet_file, filesystem=filesystem)

    roundtripped = pickle.loads(pickle.dumps(df_pa.optimize().dask))
    roundtripped_df = from_graph(
        roundtripped,
        df_pa._meta,
        df_pa.divisions,
        df_pa.__dask_keys__(),
        key_split(df_pa._name),
    )
    assert assert_eq(df_pa, roundtripped_df)


def test_pyarrow_filesystem_filters(parquet_file):
    filesystem = fs.LocalFileSystem()

    df_pa = read_parquet(parquet_file, filesystem=filesystem)
    df_pa = df_pa[df_pa.c == 1]
    expected = read_parquet(
        parquet_file, filesystem=filesystem, filters=[[("c", "==", 1)]]
    )
    assert df_pa.optimize()._name == expected.optimize()._name
    assert len(df_pa.compute()) == 1


second_parquet_file = parquet_file


def test_pyarrow_filesystem_list_of_files(parquet_file, second_parquet_file):
    filesystem = fs.LocalFileSystem()

    result = read_parquet([parquet_file, second_parquet_file], filesystem=filesystem)
    expected = pd.read_parquet([parquet_file, second_parquet_file])
    assert_eq(result, expected, check_index=False)


def test_partition_pruning(tmpdir):
    with dask.config.set({"dataframe.parquet.minimum-partition-size": 1}):
        filesystem = fs.LocalFileSystem()
        df = from_pandas(
            pd.DataFrame(
                {
                    "a": [1, 2, 3, 4, 5] * 10,
                    "b": range(50),
                }
            ),
            npartitions=2,
        )
        df.to_parquet(tmpdir, partition_on=["a"])
        ddf = read_parquet(tmpdir, filesystem=filesystem)
        ddf_filtered = read_parquet(
            tmpdir, filters=[[("a", "==", 1)]], filesystem=filesystem
        )
        assert ddf_filtered.npartitions == ddf.npartitions // 5

        ddf_optimize = read_parquet(tmpdir, filesystem=filesystem)
        ddf_optimize = ddf_optimize[ddf_optimize.a == 1].optimize()
        assert ddf_optimize.npartitions == ddf.npartitions // 5
        assert_eq(
            ddf_filtered,
            ddf_optimize,
            # FIXME ?
            check_names=False,
        )


def test_predicate_pushdown(tmpdir):
    original = pd.DataFrame(
        {
            "a": [1, 2, 3, 4, 5] * 10,
            "b": [0, 1, 2, 3, 4] * 10,
            "c": range(50),
            "d": [6, 7] * 25,
            "e": [8, 9] * 25,
        }
    )
    fn = _make_file(tmpdir, df=original)
    df = read_parquet(fn, filesystem="arrow")
    assert_eq(df, original)
    x = df[df.a == 5][df.c > 20]["b"]
    y = x.optimize(fuse=False)
    assert isinstance(y.expr.frame, FusedParquetIO)
    assert ("a", "==", 5) in y.expr.frame.operands[0].operand("filters")[0]
    assert ("c", ">", 20) in y.expr.frame.operands[0].operand("filters")[0]
    assert list(y.columns) == ["b"]

    # Check computed result
    y_result = y.compute()
    assert y_result.name == "b"
    assert len(y_result) == 6
    assert (y_result == 4).all()

    # Don't push down if replace is in there
    x = df[df.replace(5, 50).a == 5]["b"]
    y = x.optimize(fuse=False)
    assert isinstance(y.expr, Filter)
    assert len(y.compute()) == 0


def test_predicate_pushdown_compound(tmpdir):
    pdf = pd.DataFrame(
        {
            "a": [1, 2, 3, 4, 5] * 10,
            "b": [0, 1, 2, 3, 4] * 10,
            "c": range(50),
            "d": [6, 7] * 25,
            "e": [8, 9] * 25,
        }
    )
    fn = _make_file(tmpdir, df=pdf)
    df = read_parquet(fn, filesystem="arrow")

    # Test AND
    x = df[(df.a == 5) & (df.c > 20)]["b"]
    y = x.optimize(fuse=False)
    assert isinstance(y.expr.frame, FusedParquetIO)
    assert {("c", ">", 20), ("a", "==", 5)} == set(y.expr.frame.operands[0].filters[0])
    assert_eq(
        y,
        pdf[(pdf.a == 5) & (pdf.c > 20)]["b"],
        check_index=False,
    )

    # Test OR
    x = df[(df.a == 5) | (df.c > 20)]
    x = x[x.b != 0]["b"]
    y = x.optimize(fuse=False)
    assert isinstance(y.expr.frame, FusedParquetIO)
    filters = [
        set(y.expr.frame.operands[0].filters[0]),
        set(y.expr.frame.operands[0].filters[1]),
    ]
    assert {("c", ">", 20), ("b", "!=", 0)} in filters
    assert {("a", "==", 5), ("b", "!=", 0)} in filters
    expect = pdf[(pdf.a == 5) | (pdf.c > 20)]
    expect = expect[expect.b != 0]["b"]
    assert_eq(
        y,
        expect,
        check_index=False,
    )

    # Test OR and AND
    x = df[((df.a == 5) | (df.c > 20)) & (df.b != 0)]["b"]
    z = x.optimize(fuse=False)
    assert isinstance(z.expr.frame, FusedParquetIO)
    filters = [
        set(z.expr.frame.operands[0].filters[0]),
        set(z.expr.frame.operands[0].filters[1]),
    ]
    assert {("c", ">", 20), ("b", "!=", 0)} in filters
    assert {("a", "==", 5), ("b", "!=", 0)} in filters
    assert_eq(y, z)


def test_aggregate_rg_stats_to_file(tmpdir):
    filesystem = fs.LocalFileSystem()
    fn = str(tmpdir)
    ddf = from_pandas(pd.DataFrame({"a": range(10)}), npartitions=1)
    ddf.to_parquet(fn)
    ddf = read_parquet(fn, filesystem=filesystem)
    frag = ddf._expr.fragments[0]
    # Make sure this doesn't raise. We'll test the actual aggregation below
    _aggregate_statistics_to_file([frag.metadata.to_dict()])
    # In reality, we'll strip the metadata down
    assert (
        len(_aggregate_statistics_to_file([_extract_stats(frag.metadata.to_dict())]))
        > 0
    )


def test_aggregate_statistics_to_file():
    file_in = {
        "top-level-file-stat": "not-interested",
        "row_groups": [
            # RG 1
            {
                "num_rows": 100,
                "total_byte_size": 1000,
                "columns": [
                    {
                        "total_compressed_size": 10,
                        "total_uncompressed_size": 15,
                        "statistics": {
                            "min": 0,
                            "max": 10,
                            "null_count": 5,
                            "distinct_count": None,
                            "new_powerful_yet_unknown_stat": 42,
                        },
                        "path_in_schema": "a",
                    },
                    {
                        "total_compressed_size": 7,
                        "total_uncompressed_size": 23,
                        "statistics": {
                            "min": 11,
                            "max": 20,
                            "null_count": 1,
                            "distinct_count": 12,
                            "new_powerful_yet_unknown_stat": 42,
                        },
                        "path_in_schema": "b",
                    },
                ],
            },
            # RG 2
            {
                "num_rows": 50,
                "total_byte_size": 500,
                "columns": [
                    {
                        "total_compressed_size": 5,
                        "total_uncompressed_size": 7,
                        "statistics": {
                            "min": 40,
                            "max": 50,
                            "null_count": 0,
                            "distinct_count": None,
                            "new_powerful_yet_unknown_stat": 42,
                        },
                        "path_in_schema": "a",
                    },
                    {
                        "total_compressed_size": 7,
                        "total_uncompressed_size": 23,
                        "statistics": {
                            "min": 0,
                            "max": 20,
                            "null_count": 0,
                            "distinct_count": None,
                            "new_powerful_yet_unknown_stat": 42,
                        },
                        "path_in_schema": "b",
                    },
                ],
            },
        ],
    }

    expected = {
        "top-level-file-stat": "not-interested",
        "num_rows": 150,
        "total_byte_size": 1500,
        "columns": [
            {
                "path_in_schema": "a",
                "total_compressed_size": 15,
                "total_uncompressed_size": 22,
                "statistics": {
                    "min": 0,
                    "max": 50,
                },
            },
            {
                "path_in_schema": "b",
                "total_compressed_size": 14,
                "total_uncompressed_size": 46,
                "statistics": {
                    "min": 0,
                    "max": 20,
                },
            },
        ],
    }
    actual = _aggregate_statistics_to_file([file_in])
    assert len(actual) == 1
    assert actual[0] == expected


def test_combine_statistics():
    file_in = [
        # File 1
        {
            "top-level-file-stat": "not-interested",
            "num_row_groups": 2,
            "row_groups": [
                # RG 1
                {
                    "num_rows": 200,
                    "total_byte_size": 1000,
                    "columns": [
                        {
                            "total_compressed_size": 10,
                            "total_uncompressed_size": 15,
                            "statistics": {
                                "min": 0,
                                "max": 10,
                                "null_count": 5,
                                "distinct_count": None,
                                "new_powerful_yet_unknown_stat": 42,
                            },
                            "path_in_schema": "a",
                        },
                        {
                            "total_compressed_size": 7,
                            "total_uncompressed_size": 23,
                            "statistics": {
                                "min": 11,
                                "max": 20,
                                "null_count": 1,
                                "distinct_count": 12,
                                "new_powerful_yet_unknown_stat": 42,
                            },
                            "path_in_schema": "b",
                        },
                    ],
                },
                # RG 2
                {
                    "num_rows": 50,
                    "total_byte_size": 500,
                    "columns": [
                        {
                            "total_compressed_size": 5,
                            "total_uncompressed_size": 7,
                            "statistics": {
                                "min": 40,
                                "max": 50,
                                "null_count": 0,
                                "distinct_count": None,
                                "new_powerful_yet_unknown_stat": 42,
                            },
                            "path_in_schema": "a",
                        },
                        {
                            "total_compressed_size": 7,
                            "total_uncompressed_size": 23,
                            "statistics": {
                                "min": 0,
                                "max": 20,
                                "null_count": 0,
                                "distinct_count": None,
                                "new_powerful_yet_unknown_stat": 42,
                            },
                            "path_in_schema": "b",
                        },
                    ],
                },
            ],
        },
        # File 2
        {
            "top-level-file-stat": "not-interested",
            "num_row_groups": 1,
            "row_groups": [
                # RG 1
                {
                    "num_rows": 100,
                    "total_byte_size": 2000,
                    "columns": [
                        {
                            "total_compressed_size": 10,
                            "total_uncompressed_size": 15,
                            "statistics": {
                                "min": 0,
                                "max": 10,
                                "null_count": 5,
                                "distinct_count": None,
                                "new_powerful_yet_unknown_stat": 42,
                            },
                            "path_in_schema": "a",
                        },
                        {
                            "total_compressed_size": 7,
                            "total_uncompressed_size": 23,
                            "statistics": {
                                "min": 11,
                                "max": 20,
                                "null_count": 1,
                                "distinct_count": 12,
                                "new_powerful_yet_unknown_stat": 42,
                            },
                            "path_in_schema": "b",
                        },
                    ],
                },
            ],
        },
    ]
    actual = _combine_stats(file_in)
    expected = {
        "num_rows": ((200 + 50) + 100) // 2,
        "total_byte_size": ((1000 + 500) + 2000) // 2,
        "num_row_groups": 1.5,
        "columns": [
            {
                "total_compressed_size": ((10 + 5) + 10) / 2,
                "total_uncompressed_size": ((15 + 7) + 15) / 2,
                "path_in_schema": "a",
            },
            {
                "total_compressed_size": ((7 + 7) + 7) / 2,
                "total_uncompressed_size": ((23 + 23) + 23) / 2,
                "path_in_schema": "b",
            },
        ],
    }
    assert actual == expected


def test_index_only_from_parquet(tmpdir):
    pdf = pd.DataFrame({"foo": range(5)}, index=range(50, 55))
    pdf.to_parquet(f"{tmpdir}/test.parquet")
    result = read_parquet(f"{tmpdir}/test.parquet").index
    assert_eq(result, pdf.index)


def test_timestamp_divisions(tmpdir):
    pdf = pd.DataFrame.from_dict(
        {
            "Date": ["11/26/2017", "11/26/2017"],
            "Time": ["17:00:00.067", "17:00:00.102"],
            "Volume": [403, 3],
        }
    )
    pdf["Timestamp"] = pd.to_datetime(pdf.Date) + pd.to_timedelta(pdf.Time)
    pdf.to_parquet(f"{tmpdir}/test.parquet")
    df = read_parquet(
        f"{tmpdir}/test.parquet", index="Timestamp", calculate_divisions=True
    )["Volume"]
    assert df.optimize().divisions == (
        pd.Timestamp("2017-11-26 17:00:00.067000"),
        pd.Timestamp("2017-11-26 17:00:00.102000"),
    )


def test_read_parquet_index_projection(tmpdir):
    df = from_array(
        np.zeros((201 * 10, 8), dtype=np.int64),
        columns=["val1", "val2", "val3", "val4", "val5", "val6", "val7", "tsprv"],
    )
    df = df.repartition(npartitions=201)
    df.to_parquet(f"{tmpdir}/index", write_index=False)

    df = read_parquet(f"{tmpdir}/index", calculate_divisions=True)
    result = df.assign(dts=df.index - df.tsprv)
    expected = df.compute()
    expected = expected.assign(dts=expected.index - expected.tsprv)

    assert_eq(result.dts.min(), expected.dts.min())


@pytest.mark.filterwarnings(
    # The stats collection encounters the async client and falls back to a
    # sync scheduler. Irrelevant for this test.
    "ignore:.*Client.*asynchronous:UserWarning",
)
def test_ensure_plan_computed_during_optimization(tmpdir, filesystem):

    pytest.importorskip("distributed")
    from distributed.utils_test import gen_cluster

    @gen_cluster(
        client=True,
        clean_kwargs={"threads": False},
        config={"admin.async-client-fallback": "sync"},
    )
    async def _test(c, s, *workers):
        # https://github.com/dask/dask/issues/11932

        df = read_parquet(_make_file(tmpdir), filesystem=filesystem)
        df = df[df.a == 1]
        res = await c.compute(df.size)
        assert res == 5

    _test()


def test_tune_optimization_disabled(tmpdir):
    """Test that optimization.tune.active=False preserves the original partition structure."""
    n_partitions = 10
    pdf = pd.DataFrame(
        {
            "x": np.random.randn(1000),
            "y": np.random.randn(1000),
        }
    )
    ddf = from_pandas(pdf, npartitions=n_partitions)

    # Write parquet with multiple partitions
    path = str(tmpdir.join("test.parquet"))
    ddf.to_parquet(path)

    # Read with tune optimization disabled via config
    with dask.config.set({"optimization.tune.active": False}):
        result = read_parquet(path)

        # Verify that the number of partitions is preserved
        assert result.npartitions == n_partitions

        # Verify that map_partitions(len) gives consistent results
        # for both the full DataFrame and a column selection
        full_lens = result.map_partitions(len).compute()
        col_lens = result["x"].map_partitions(len).compute()

        assert len(full_lens) == n_partitions
        assert len(col_lens) == n_partitions
        assert_eq(full_lens, col_lens)
