# Licensed to Modin Development Team under one or more contributor license agreements.
# See the NOTICE file distributed with this work for additional information regarding
# copyright ownership.  The Modin Development Team licenses this file to you under the
# Apache License, Version 2.0 (the "License"); you may not use this file except in
# compliance with the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software distributed under
# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific language
# governing permissions and limitations under the License.

import pytest
import numpy as np
import pandas
import ray

import modin.pandas as pd
from modin.config import Engine, NPartitions
from modin.distributed.dataframe.pandas.partitions import from_partitions
from modin.experimental.batch.pipeline import PandasQueryPipeline
from modin.pandas.test.utils import df_equals


@pytest.mark.skipif(
    Engine.get() != "Ray",
    reason="Only Ray supports the Batch Pipeline API",
)
class TestPipelineRayEngine:
    def test_warnings(self):
        """Ensure that creating a Pipeline object raises the correct warnings."""
        arr = np.random.randint(0, 1000, (1000, 1000))
        df = pd.DataFrame(arr)
        # Ensure that building a pipeline warns users that it is an experimental feature
        with pytest.warns(
            UserWarning,
            match="The Batch Pipeline API is an experimental feature and still under development in Modin.",
        ):
            pipeline = PandasQueryPipeline(df)
        with pytest.warns(
            UserWarning,
            match="No outputs to compute. Returning an empty list. Please specify outputs by calling `add_query` with `is_output=True`.",
        ):
            output = pipeline.compute_batch()
        assert output == [], "Empty pipeline did not return an empty list."

    def test_pipeline_simple(self):
        """Create a simple pipeline and ensure that it runs end to end correctly."""
        arr = np.random.randint(0, 1000, (1000, 1000))
        df = pd.DataFrame(arr)

        def add_col(df):
            df["new_col"] = df.sum(axis=1)
            return df

        # Build pipeline
        pipeline = PandasQueryPipeline(df)
        pipeline.add_query(add_col)
        pipeline.add_query(lambda df: df * -30)
        pipeline.add_query(
            lambda df: df.rename(columns={i: f"col {i}" for i in range(1000)})
        )

        def add_row_to_partition(df):
            return pandas.concat([df, df.iloc[[-1]]])

        pipeline.add_query(add_row_to_partition, is_output=True)
        new_df = pipeline.compute_batch()[0]
        # Build df without pipelining to ensure correctness
        correct_df = add_col(pd.DataFrame(arr))
        correct_df *= -30
        correct_df = pd.DataFrame(
            correct_df.rename(columns={i: f"col {i}" for i in range(1000)})._to_pandas()
        )
        correct_modin_frame = correct_df._query_compiler._modin_frame
        partitions = correct_modin_frame._partition_mgr_cls.row_partitions(
            correct_modin_frame._partitions
        )
        partitions = [
            partition.add_to_apply_calls(add_row_to_partition)
            for partition in partitions
        ]
        [partition.drain_call_queue() for partition in partitions]
        partitions = [partition.list_of_blocks for partition in partitions]
        correct_df = from_partitions(partitions, axis=None)
        # Compare pipelined and non-pipelined df
        df_equals(correct_df, new_df)
        # Ensure that setting `num_partitions` when creating a pipeline does not change `NPartitions`
        num_partitions = NPartitions.get()
        PandasQueryPipeline(df, num_partitions=(num_partitions - 1))
        assert (
            NPartitions.get() == num_partitions
        ), "Pipeline did not change NPartitions.get()"

    def test_update_df(self):
        """Ensure that `update_df` updates the df that the pipeline runs on."""
        df = pd.DataFrame([[1, 2, 3], [4, 5, 6]])
        pipeline = PandasQueryPipeline(df)
        pipeline.add_query(lambda df: df + 3, is_output=True)
        new_df = df * -1
        pipeline.update_df(new_df)
        output_df = pipeline.compute_batch()[0]
        df_equals((df * -1) + 3, output_df)

    def test_multiple_outputs(self):
        """Create a pipeline with multiple outputs, and check that all are computed correctly."""
        arr = np.random.randint(0, 1000, (1000, 1000))
        df = pd.DataFrame(arr)
        pipeline = PandasQueryPipeline(df)
        pipeline.add_query(lambda df: df * -30, is_output=True)
        pipeline.add_query(
            lambda df: df.rename(columns={i: f"col {i}" for i in range(1000)}),
            is_output=True,
        )
        pipeline.add_query(lambda df: df + 30, is_output=True)
        new_dfs = pipeline.compute_batch()
        assert len(new_dfs) == 3, "Pipeline did not return all outputs"
        correct_df = pd.DataFrame(arr) * -30
        df_equals(correct_df, new_dfs[0])  # First output computed correctly
        correct_df = correct_df.rename(columns={i: f"col {i}" for i in range(1000)})
        df_equals(correct_df, new_dfs[1])  # Second output computed correctly
        correct_df += 30
        df_equals(correct_df, new_dfs[2])  # Third output computed correctly

    def test_output_id(self):
        """Ensure `output_id` is handled correctly when passed."""
        arr = np.random.randint(0, 1000, (1000, 1000))
        df = pd.DataFrame(arr)
        pipeline = PandasQueryPipeline(df, 0)
        pipeline.add_query(lambda df: df * -30, is_output=True, output_id=20)
        with pytest.raises(
            ValueError, match="Output ID must be specified for all nodes."
        ):
            pipeline.add_query(
                lambda df: df.rename(columns={i: f"col {i}" for i in range(1000)}),
                is_output=True,
            )
        assert (
            len(pipeline.query_list) == 0 and len(pipeline.outputs) == 1
        ), "Invalid `add_query` incorrectly added a node to the pipeline."
        pipeline = PandasQueryPipeline(df)
        pipeline.add_query(lambda df: df * -30, is_output=True)
        with pytest.raises(
            ValueError, match="Output ID must be specified for all nodes."
        ):
            pipeline.add_query(
                lambda df: df.rename(columns={i: f"col {i}" for i in range(1000)}),
                is_output=True,
                output_id=20,
            )
        assert (
            len(pipeline.query_list) == 0 and len(pipeline.outputs) == 1
        ), "Invalid `add_query` incorrectly added a node to the pipeline."
        pipeline = PandasQueryPipeline(df)
        pipeline.add_query(lambda df: df, is_output=True)
        with pytest.raises(
            ValueError,
            match=(
                "`pass_output_id` is set to True, but output ids have not been specified. "
                + "To pass output ids, please specify them using the `output_id` kwarg with pipeline.add_query"
            ),
        ):
            pipeline.compute_batch(postprocessor=lambda df: df, pass_output_id=True)
        with pytest.raises(
            ValueError,
            match="Output ID cannot be specified for non-output node.",
        ):
            pipeline.add_query(lambda df: df, output_id=22)
        assert (
            len(pipeline.query_list) == 0 and len(pipeline.outputs) == 1
        ), "Invalid `add_query` incorrectly added a node to the pipeline."

    def test_output_id_multiple_outputs(self):
        """Ensure `output_id` is handled correctly when multiple outputs are computed."""
        arr = np.random.randint(0, 1000, (1000, 1000))
        df = pd.DataFrame(arr)
        pipeline = PandasQueryPipeline(df)
        pipeline.add_query(lambda df: df * -30, is_output=True, output_id=20)
        pipeline.add_query(
            lambda df: df.rename(columns={i: f"col {i}" for i in range(1000)}),
            is_output=True,
            output_id=21,
        )
        pipeline.add_query(lambda df: df + 30, is_output=True, output_id=22)
        new_dfs = pipeline.compute_batch()
        assert isinstance(
            new_dfs, dict
        ), "Pipeline did not return a dictionary mapping output_ids to dfs"
        assert 20 in new_dfs, "Output ID 1 not cached correctly"
        assert 21 in new_dfs, "Output ID 2 not cached correctly"
        assert 22 in new_dfs, "Output ID 3 not cached correctly"
        assert len(new_dfs) == 3, "Pipeline did not return all outputs"
        correct_df = pd.DataFrame(arr) * -30
        df_equals(correct_df, new_dfs[20])  # First output computed correctly
        correct_df = correct_df.rename(columns={i: f"col {i}" for i in range(1000)})
        df_equals(correct_df, new_dfs[21])  # Second output computed correctly
        correct_df += 30
        df_equals(correct_df, new_dfs[22])  # Third output computed correctly

    def test_postprocessing(self):
        """Check that the `postprocessor` argument to `_compute_batch` is handled correctly."""
        arr = np.random.randint(0, 1000, (1000, 1000))
        df = pd.DataFrame(arr)
        pipeline = PandasQueryPipeline(df)
        pipeline.add_query(lambda df: df * -30, is_output=True)
        pipeline.add_query(
            lambda df: df.rename(columns={i: f"col {i}" for i in range(1000)}),
            is_output=True,
        )
        pipeline.add_query(lambda df: df + 30, is_output=True)

        def new_col_adder(df):
            df["new_col"] = df.iloc[:, -1]
            return df

        new_dfs = pipeline.compute_batch(postprocessor=new_col_adder)
        assert len(new_dfs) == 3, "Pipeline did not return all outputs"
        correct_df = pd.DataFrame(arr) * -30
        correct_df["new_col"] = correct_df.iloc[:, -1]
        df_equals(correct_df, new_dfs[0])
        correct_df = correct_df.drop(columns=["new_col"])
        correct_df = correct_df.rename(columns={i: f"col {i}" for i in range(1000)})
        correct_df["new_col"] = correct_df.iloc[:, -1]
        df_equals(correct_df, new_dfs[1])
        correct_df = correct_df.drop(columns=["new_col"])
        correct_df += 30
        correct_df["new_col"] = correct_df.iloc[:, -1]
        df_equals(correct_df, new_dfs[2])

    def test_postprocessing_with_output_id(self):
        """Check that the `postprocessor` argument is correctly handled when `output_id` is specified."""

        def new_col_adder(df):
            df["new_col"] = df.iloc[:, -1]
            return df

        arr = np.random.randint(0, 1000, (1000, 1000))
        df = pd.DataFrame(arr)
        pipeline = PandasQueryPipeline(df)
        pipeline.add_query(lambda df: df * -30, is_output=True, output_id=20)
        pipeline.add_query(
            lambda df: df.rename(columns={i: f"col {i}" for i in range(1000)}),
            is_output=True,
            output_id=21,
        )
        pipeline.add_query(lambda df: df + 30, is_output=True, output_id=22)
        new_dfs = pipeline.compute_batch(postprocessor=new_col_adder)
        assert len(new_dfs) == 3, "Pipeline did not return all outputs"

    def test_postprocessing_with_output_id_passed(self):
        """Check that the `postprocessor` argument is correctly passed `output_id` when `pass_output_id` is `True`."""
        arr = np.random.randint(0, 1000, (1000, 1000))

        def new_col_adder(df, o_id):
            df["new_col"] = o_id
            return df

        df = pd.DataFrame(arr)
        pipeline = PandasQueryPipeline(df)
        pipeline.add_query(lambda df: df * -30, is_output=True, output_id=20)
        pipeline.add_query(
            lambda df: df.rename(columns={i: f"col {i}" for i in range(1000)}),
            is_output=True,
            output_id=21,
        )
        pipeline.add_query(lambda df: df + 30, is_output=True, output_id=22)
        new_dfs = pipeline.compute_batch(
            postprocessor=new_col_adder, pass_output_id=True
        )
        correct_df = pd.DataFrame(arr) * -30
        correct_df["new_col"] = 20
        df_equals(correct_df, new_dfs[20])
        correct_df = correct_df.drop(columns=["new_col"])
        correct_df = correct_df.rename(columns={i: f"col {i}" for i in range(1000)})
        correct_df["new_col"] = 21
        df_equals(correct_df, new_dfs[21])
        correct_df = correct_df.drop(columns=["new_col"])
        correct_df += 30
        correct_df["new_col"] = 22
        df_equals(correct_df, new_dfs[22])

    def test_postprocessing_with_partition_id(self):
        """Check that the postprocessing is correctly handled when `partition_id` is passed."""
        arr = np.random.randint(0, 1000, (1000, 1000))

        def new_col_adder(df, partition_id):
            df["new_col"] = partition_id
            return df

        df = pd.DataFrame(arr)
        pipeline = PandasQueryPipeline(df)
        pipeline.add_query(lambda df: df * -30, is_output=True, output_id=20)
        pipeline.add_query(
            lambda df: df.rename(columns={i: f"col {i}" for i in range(1000)}),
            is_output=True,
            output_id=21,
        )
        new_dfs = pipeline.compute_batch(
            postprocessor=new_col_adder, pass_partition_id=True
        )
        correct_df = pd.DataFrame(arr) * -30
        correct_modin_frame = correct_df._query_compiler._modin_frame
        partitions = correct_modin_frame._partition_mgr_cls.row_partitions(
            correct_modin_frame._partitions
        )
        partitions = [
            partition.add_to_apply_calls(new_col_adder, i)
            for i, partition in enumerate(partitions)
        ]
        [partition.drain_call_queue() for partition in partitions]
        partitions = [partition.list_of_blocks for partition in partitions]
        correct_df = from_partitions(partitions, axis=None)
        df_equals(correct_df, new_dfs[20])
        correct_df = correct_df.drop(columns=["new_col"])
        correct_df = pd.DataFrame(
            correct_df.rename(columns={i: f"col {i}" for i in range(1000)})._to_pandas()
        )
        correct_modin_frame = correct_df._query_compiler._modin_frame
        partitions = correct_modin_frame._partition_mgr_cls.row_partitions(
            correct_modin_frame._partitions
        )
        partitions = [
            partition.add_to_apply_calls(new_col_adder, i)
            for i, partition in enumerate(partitions)
        ]
        [partition.drain_call_queue() for partition in partitions]
        partitions = [partition.list_of_blocks for partition in partitions]
        correct_df = from_partitions(partitions, axis=None)
        df_equals(correct_df, new_dfs[21])

    def test_postprocessing_with_all_metadata(self):
        """Check that postprocessing is correctly handled when `partition_id` and `output_id` are passed."""
        arr = np.random.randint(0, 1000, (1000, 1000))

        def new_col_adder(df, o_id, partition_id):
            df["new_col"] = f"{o_id} {partition_id}"
            return df

        df = pd.DataFrame(arr)
        pipeline = PandasQueryPipeline(df)
        pipeline.add_query(lambda df: df * -30, is_output=True, output_id=20)
        pipeline.add_query(
            lambda df: df.rename(columns={i: f"col {i}" for i in range(1000)}),
            is_output=True,
            output_id=21,
        )
        new_dfs = pipeline.compute_batch(
            postprocessor=new_col_adder, pass_partition_id=True, pass_output_id=True
        )
        correct_df = pd.DataFrame(arr) * -30
        correct_modin_frame = correct_df._query_compiler._modin_frame
        partitions = correct_modin_frame._partition_mgr_cls.row_partitions(
            correct_modin_frame._partitions
        )
        partitions = [
            partition.add_to_apply_calls(new_col_adder, 20, i)
            for i, partition in enumerate(partitions)
        ]
        [partition.drain_call_queue() for partition in partitions]
        partitions = [partition.list_of_blocks for partition in partitions]
        correct_df = from_partitions(partitions, axis=None)
        df_equals(correct_df, new_dfs[20])
        correct_df = correct_df.drop(columns=["new_col"])
        correct_df = pd.DataFrame(
            correct_df.rename(columns={i: f"col {i}" for i in range(1000)})._to_pandas()
        )
        correct_modin_frame = correct_df._query_compiler._modin_frame
        partitions = correct_modin_frame._partition_mgr_cls.row_partitions(
            correct_modin_frame._partitions
        )
        partitions = [
            partition.add_to_apply_calls(new_col_adder, 21, i)
            for i, partition in enumerate(partitions)
        ]
        [partition.drain_call_queue() for partition in partitions]
        partitions = [partition.list_of_blocks for partition in partitions]
        correct_df = from_partitions(partitions, axis=None)
        df_equals(correct_df, new_dfs[21])

    def test_repartition_after(self):
        """Check that the `repartition_after` argument is appropriately handled."""
        df = pd.DataFrame([list(range(1000))])
        pipeline = PandasQueryPipeline(df)
        pipeline.add_query(
            lambda df: pandas.concat([df] * 1000), repartition_after=True
        )

        def new_col_adder(df, partition_id):
            df["new_col"] = partition_id
            return df

        pipeline.add_query(new_col_adder, is_output=True, pass_partition_id=True)
        new_dfs = pipeline.compute_batch()
        # new_col_adder should set `new_col` to the partition ID
        # throughout the dataframe. We expect there to be
        # NPartitions.get() partitions by the time new_col_adder runs,
        # because the previous step has repartitioned.
        assert len(new_dfs[0]["new_col"].unique()) == NPartitions.get()
        # Test that `repartition_after=True` raises an error when the result has more than
        # one partition.
        partition1 = ray.put(pandas.DataFrame([[0, 1, 2]]))
        partition2 = ray.put(pandas.DataFrame([[3, 4, 5]]))
        df = from_partitions([partition1, partition2], 0)
        pipeline = PandasQueryPipeline(df, 0)
        pipeline.add_query(lambda df: df, repartition_after=True, is_output=True)

        with pytest.raises(
            NotImplementedError,
            match="Dynamic repartitioning is currently only supported for DataFrames with 1 partition.",
        ):
            new_dfs = pipeline.compute_batch()

    def test_fan_out(self):
        """Check that the fan_out argument is appropriately handled."""
        df = pd.DataFrame([[0, 1, 2]])

        def new_col_adder(df, partition_id):
            df["new_col"] = partition_id
            return df

        def reducer(dfs):
            new_cols = "".join([str(df["new_col"].values[0]) for df in dfs])
            dfs[0]["new_col1"] = new_cols
            return dfs[0]

        pipeline = PandasQueryPipeline(df)
        pipeline.add_query(
            new_col_adder,
            fan_out=True,
            reduce_fn=reducer,
            pass_partition_id=True,
            is_output=True,
        )
        new_df = pipeline.compute_batch()[0]
        correct_df = pd.DataFrame([[0, 1, 2]])
        correct_df["new_col"] = 0
        correct_df["new_col1"] = "".join([str(i) for i in range(NPartitions.get())])
        df_equals(correct_df, new_df)
        # Test that `fan_out=True` raises an error when the input has more than
        # one partition.
        partition1 = ray.put(pandas.DataFrame([[0, 1, 2]]))
        partition2 = ray.put(pandas.DataFrame([[3, 4, 5]]))
        df = from_partitions([partition1, partition2], 0)
        pipeline = PandasQueryPipeline(df)
        pipeline.add_query(
            new_col_adder,
            fan_out=True,
            reduce_fn=reducer,
            pass_partition_id=True,
            is_output=True,
        )
        with pytest.raises(
            NotImplementedError,
            match="Fan out is only supported with DataFrames with 1 partition.",
        ):
            new_df = pipeline.compute_batch()[0]

    def test_pipeline_complex(self):
        """Create a complex pipeline with both `fan_out`, `repartition_after` and postprocessing and ensure that it runs end to end correctly."""
        from os.path import exists
        from os import remove
        from time import sleep

        df = pd.DataFrame([[0, 1, 2]])

        def new_col_adder(df, partition_id):
            sleep(60)
            df["new_col"] = partition_id
            return df

        def reducer(dfs):
            new_cols = "".join([str(df["new_col"].values[0]) for df in dfs])
            dfs[0]["new_col1"] = new_cols
            return dfs[0]

        desired_num_partitions = 24
        pipeline = PandasQueryPipeline(df, num_partitions=desired_num_partitions)
        pipeline.add_query(
            new_col_adder,
            fan_out=True,
            reduce_fn=reducer,
            pass_partition_id=True,
            is_output=True,
            output_id=20,
        )
        pipeline.add_query(
            lambda df: pandas.concat([df] * 1000),
            repartition_after=True,
        )

        def to_csv(df, partition_id):
            df = df.drop(columns=["new_col"])
            df.to_csv(f"{partition_id}.csv")
            return df

        pipeline.add_query(to_csv, is_output=True, output_id=21, pass_partition_id=True)

        def post_proc(df, o_id, partition_id):
            df["new_col_proc"] = f"{o_id} {partition_id}"
            return df

        new_dfs = pipeline.compute_batch(
            postprocessor=post_proc,
            pass_partition_id=True,
            pass_output_id=True,
        )
        correct_df = pd.DataFrame([[0, 1, 2]])
        correct_df["new_col"] = 0
        correct_df["new_col1"] = "".join(
            [str(i) for i in range(desired_num_partitions)]
        )
        correct_df["new_col_proc"] = "20 0"
        df_equals(correct_df, new_dfs[20])
        correct_df = pd.concat([correct_df] * 1000)
        correct_df = correct_df.drop(columns=["new_col"])
        correct_df["new_col_proc"] = "21 0"
        new_length = len(correct_df.index) // desired_num_partitions
        for i in range(desired_num_partitions):
            if i == desired_num_partitions - 1:
                correct_df.iloc[i * new_length :, -1] = f"21 {i}"
            else:
                correct_df.iloc[i * new_length : (i + 1) * new_length, -1] = f"21 {i}"
        df_equals(correct_df, new_dfs[21])
        correct_df = correct_df.drop(columns=["new_col_proc"])
        for i in range(desired_num_partitions):
            if i == desired_num_partitions - 1:
                correct_partition = correct_df.iloc[i * new_length :]
            else:
                correct_partition = correct_df.iloc[
                    i * new_length : (i + 1) * new_length
                ]
            assert exists(
                f"{i}.csv"
            ), "CSV File for Partition {i} does not exist, even though dataframe should have been repartitioned."
            df_equals(
                correct_partition,
                pd.read_csv(f"{i}.csv", index_col="Unnamed: 0").rename(
                    columns={"0": 0, "1": 1, "2": 2}
                ),
            )
            remove(f"{i}.csv")


@pytest.mark.skipif(
    Engine.get() == "Ray",
    reason="Ray supports the Batch Pipeline API",
)
def test_pipeline_unsupported_engine():
    """Ensure that trying to use the Pipeline API with an unsupported Engine raises errors."""
    # Check that pipeline does not allow `Engine` to not be Ray.
    df = pd.DataFrame([[1]])
    with pytest.raises(
        NotImplementedError,
        match="Batch Pipeline API is only implemented for `PandasOnRay` execution.",
    ):
        PandasQueryPipeline(df)

    eng = Engine.get()
    Engine.put("Ray")
    # Check that even if Engine is Ray, if the df is not backed by Ray, the Pipeline does not allow initialization.
    with pytest.raises(
        NotImplementedError,
        match="Batch Pipeline API is only implemented for `PandasOnRay` execution.",
    ):
        PandasQueryPipeline(df, 0)
    df_on_ray_engine = pd.DataFrame([[1]])
    pipeline = PandasQueryPipeline(df_on_ray_engine)
    # Check that even if Engine is Ray, if the new df is not backed by Ray, the Pipeline does not allow an update.
    with pytest.raises(
        NotImplementedError,
        match="Batch Pipeline API is only implemented for `PandasOnRay` execution.",
    ):
        pipeline.update_df(df)
    Engine.put(eng)
    # Check that pipeline does not allow an update when `Engine` is not Ray.
    with pytest.raises(
        NotImplementedError,
        match="Batch Pipeline API is only implemented for `PandasOnRay` execution.",
    ):
        pipeline.update_df(df)
