# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import os
import unittest
import warnings
from itertools import chain

import expecttest
from _utils._common_utils_for_test import create_temp_dir, reset_after_n_next_calls
from torchdata.datapipes.iter import DataFrameMaker, FileLister, FileOpener, IterableWrapper, ParquetDataFrameLoader

try:
    import torcharrow
    import torcharrow.dtypes as dt

    HAS_TORCHARROW = True
except ImportError:
    HAS_TORCHARROW = False

try:
    import pyarrow
    import pyarrow.parquet as parquet

    HAS_PYARROW = True
except ImportError:
    HAS_PYARROW = False
skipIfNoPyArrow = unittest.skipIf(not HAS_PYARROW, "no PyArrow.")
skipIfNoTorchArrow = unittest.skipIf(not HAS_TORCHARROW, "no TorchArrow.")


@skipIfNoTorchArrow
class TestDataFrame(expecttest.TestCase):
    def setUp(self) -> None:
        self.temp_dir = create_temp_dir()
        if HAS_PYARROW:
            self._write_parquet_files()

    def tearDown(self) -> None:
        try:
            self.temp_dir.cleanup()
        except Exception as e:
            warnings.warn(f"TestDataFrame was not able to cleanup temp dir due to {e}")

    def _write_parquet_files(self):
        # Create TorchArrow DataFrames
        DTYPE = dt.Struct([dt.Field("Values", dt.int32)])
        df1 = torcharrow.dataframe([(i,) for i in range(10)], dtype=DTYPE)
        df2 = torcharrow.dataframe([(i,) for i in range(100)], dtype=DTYPE)
        # Write them as parquet files
        for i, df in enumerate([df1, df2]):
            fname = f"df{i}.parquet"
            self._write_df_as_parquet(df, fname)
        self._write_multiple_dfs_as_parquest([df1, df2], fname="merged.parquet")

    def _custom_files_set_up(self, files):
        for fname, content in files.items():
            temp_file_path = os.path.join(self.temp_dir.name, fname)
            with open(temp_file_path, "w") as f:
                f.write(content)

    def _compare_dataframes(self, expected_df, actual_df):
        self.assertEqual(len(expected_df), len(actual_df))
        for exp, act in zip(expected_df, actual_df):
            self.assertEqual(exp, act)

    def _write_df_as_parquet(self, df, fname: str) -> None:
        table = df.to_arrow()
        parquet.write_table(table, os.path.join(self.temp_dir.name, fname))

    def _write_multiple_dfs_as_parquest(self, dfs, fname: str) -> None:
        tables = [df.to_arrow() for df in dfs]
        merged_table = pyarrow.concat_tables(tables)
        parquet.write_table(merged_table, os.path.join(self.temp_dir.name, fname))

    def test_dataframe_maker_iterdatapipe(self):
        source_data = [(i,) for i in range(10)]
        source_dp = IterableWrapper(source_data)
        DTYPE = dt.Struct([dt.Field("Values", dt.int32)])

        # Functional Test: DataPipe correctly converts into a single TorchArrow DataFrame
        df_dp = source_dp.dataframe(dtype=DTYPE)
        df = list(df_dp)[0]
        expected_df = torcharrow.dataframe([(i,) for i in range(10)], dtype=DTYPE)
        self._compare_dataframes(expected_df, df)

        # Functional Test: DataPipe correctly converts into multiple TorchArrow DataFrames, based on size argument
        df_dp = DataFrameMaker(source_dp, dataframe_size=5, dtype=DTYPE)
        dfs = list(df_dp)
        expected_dfs = [
            torcharrow.dataframe([(i,) for i in range(5)], dtype=DTYPE),
            torcharrow.dataframe([(i,) for i in range(5, 10)], dtype=DTYPE),
        ]
        for exp_df, act_df in zip(expected_dfs, dfs):
            self._compare_dataframes(exp_df, act_df)

        # __len__ Test:
        df_dp = source_dp.dataframe(dtype=DTYPE)
        self.assertEqual(1, len(df_dp))
        self.assertEqual(10, len(list(df_dp)[0]))
        df_dp = source_dp.dataframe(dataframe_size=5, dtype=DTYPE)
        self.assertEqual(2, len(df_dp))
        self.assertEqual(5, len(list(df_dp)[0]))

        # Reset Test:
        n_elements_before_reset = 1
        res_before_reset, res_after_reset = reset_after_n_next_calls(df_dp, n_elements_before_reset)
        for exp_df, act_df in zip(expected_dfs[:1], res_before_reset):
            self._compare_dataframes(exp_df, act_df)
        for exp_df, act_df in zip(expected_dfs, res_after_reset):
            self._compare_dataframes(exp_df, act_df)

    def test_dataframe_maker_with_csv(self):
        def get_name(path_and_stream):
            return os.path.basename(path_and_stream[0]), path_and_stream[1]

        csv_files = {"1.csv": "key,item\na,1\nb,2"}
        self._custom_files_set_up(csv_files)
        datapipe1 = FileLister(self.temp_dir.name, "*.csv")
        datapipe2 = FileOpener(datapipe1, mode="b")
        datapipe3 = datapipe2.map(get_name)
        csv_dict_parser_dp = datapipe3.parse_csv_as_dict()

        # Functional Test: Correctly generate TorchArrow DataFrame from CSV
        DTYPE = dt.Struct([dt.Field("key", dt.string), dt.Field("item", dt.string)])
        df_dp = csv_dict_parser_dp.dataframe(dtype=DTYPE, columns=["key", "item"])
        expected_dfs = [torcharrow.dataframe([{"key": "a", "item": "1"}, {"key": "b", "item": "2"}], dtype=DTYPE)]
        for exp_df, act_df in zip(expected_dfs, list(df_dp)):
            self._compare_dataframes(exp_df, act_df)

        # Functional: making sure DataPipe works even without `columns` input
        df_dp = csv_dict_parser_dp.dataframe(dtype=DTYPE)
        for exp_df, act_df in zip(expected_dfs, list(df_dp)):
            self._compare_dataframes(exp_df, act_df)

    @skipIfNoPyArrow
    def test_parquet_dataframe_reader_iterdatapipe(self):
        DTYPE = dt.Struct([dt.Field("Values", dt.int32)])

        # Functional Test: read from Parquet files and output TorchArrow DataFrames
        source_dp = FileLister(self.temp_dir.name, masks="df*.parquet")
        parquet_df_dp = ParquetDataFrameLoader(source_dp, dtype=DTYPE)
        expected_dfs = [
            torcharrow.dataframe([(i,) for i in range(10)], dtype=DTYPE),
            torcharrow.dataframe([(i,) for i in range(100)], dtype=DTYPE),
        ]
        for exp_df, act_df in zip(expected_dfs, list(parquet_df_dp)):
            self._compare_dataframes(exp_df, act_df)

        # Functional Test: correctly read from a Parquet file that was a merged DataFrame
        merged_source_dp = FileLister(self.temp_dir.name, masks="merged.parquet")
        merged_parquet_df_dp = ParquetDataFrameLoader(merged_source_dp, dtype=DTYPE)
        expected_merged_dfs = [torcharrow.dataframe([(i,) for i in chain(range(10), range(100))], dtype=DTYPE)]
        for exp_df, act_df in zip(expected_merged_dfs, list(merged_parquet_df_dp)):
            self._compare_dataframes(exp_df, act_df)

        # __len__ Test: no valid length because we do not know the number of row groups in advance
        with self.assertRaisesRegex(TypeError, "has no len"):
            len(parquet_df_dp)

        # Reset Test:
        n_elements_before_reset = 1
        res_before_reset, res_after_reset = reset_after_n_next_calls(parquet_df_dp, n_elements_before_reset)
        for exp_df, act_df in zip(expected_dfs[:1], res_before_reset):
            self._compare_dataframes(exp_df, act_df)
        for exp_df, act_df in zip(expected_dfs, res_after_reset):
            self._compare_dataframes(exp_df, act_df)


if __name__ == "__main__":
    unittest.main()
