import unittest
import os
import shutil
import pandas as pd
import numpy as np

from pyedautils.data_io import save_data, load_data

class TestDataIO(unittest.TestCase):
    def setUp(self):
        # Create a temporary directory for test files
        self.test_dir = "test_data"
        os.makedirs(self.test_dir, exist_ok=True)

    def tearDown(self):
        # Remove the temporary directory and its contents after tests
        shutil.rmtree(self.test_dir)

    def test_save_and_load_json(self):
        # Test data
        data = {"key1": "value1", "key2": "value2"}
        file_path = os.path.join(self.test_dir, "test.json")

        # Test saving and loading JSON file
        save_data(data, file_path)
        loaded_data = load_data(file_path)

        self.assertEqual(data, loaded_data)

    def test_save_and_load_csv(self):
        # Test data: time series data with datetime index
        dates = pd.date_range(start="2022-01-01", periods=10, freq="D")
        values = np.round(np.random.randn(10), 3)
        df = pd.DataFrame({"Value": values}, index=dates)
        file_path = os.path.join(self.test_dir, "test.csv")

        # Test saving and loading CSV file
        save_data(df, file_path, index=True)
        loaded_data = load_data(file_path)
        loaded_data.index = pd.to_datetime(loaded_data.iloc[:, 0])
        loaded_data.drop(loaded_data.columns[0], axis=1, inplace=True)

        self.assertTrue(isinstance(loaded_data, pd.DataFrame))
        self.assertTrue(df.equals(loaded_data))
        
    def test_save_and_load_new_directory(self):
        # Test data: time series data with datetime index
        dates = pd.date_range(start="2022-01-01", periods=10, freq="D")
        values = np.round(np.random.randn(10), 3)
        df = pd.DataFrame({"Value": values}, index=dates)
        file_path = os.path.join(self.test_dir, "newdir/test.csv")

        # Test saving and loading CSV file
        save_data(df, file_path, index=True)
        loaded_data = load_data(file_path)
        loaded_data.index = pd.to_datetime(loaded_data.iloc[:, 0])
        loaded_data.drop(loaded_data.columns[0], axis=1, inplace=True)

        self.assertTrue(isinstance(loaded_data, pd.DataFrame))
        self.assertTrue(df.equals(loaded_data))

    def test_save_and_load_pickle(self):
        # Test data: time series data with datetime index
        dates = pd.date_range(start="2022-01-01", periods=10, freq="D")
        values = np.random.randn(10)
        df = pd.DataFrame({"Value": values}, index=dates)
        file_path = os.path.join(self.test_dir, "test.pkl")

        # Test saving and loading pickle file
        save_data(df, file_path)
        loaded_data = load_data(file_path)

        self.assertTrue(isinstance(loaded_data, pd.DataFrame))
        self.assertTrue(df.equals(loaded_data))

    def test_save_and_load_compressed_pickle(self):
        # Test data: time series data with datetime index
        dates = pd.date_range(start="2022-01-01", periods=10, freq="D")
        values = np.random.randn(10)
        df = pd.DataFrame({"Value": values}, index=dates)
        file_path = os.path.join(self.test_dir, "test.pklz")

        # Test saving and loading compressed pickle file
        save_data(df, file_path)
        loaded_data = load_data(file_path)

        self.assertTrue(isinstance(loaded_data, pd.DataFrame))
        self.assertTrue(df.equals(loaded_data))
    
    def test_save_unsupported_format(self):
        # Test data: time series data with datetime index
        dates = pd.date_range(start="2022-01-01", periods=10, freq="D")
        values = np.random.randn(10)
        df = pd.DataFrame({"Value": values}, index=dates)
        file_path = os.path.join(self.test_dir, "test.abc")

        with self.assertLogs("pyedautils.data_io", level="INFO") as cm:
            save_data(df, file_path)

        log_output = "\n".join(cm.output)
        self.assertIn("Saving data to:", log_output)
        self.assertIn("Unsupported file format", log_output)

    def test_load_unsupported_format(self):
        # Test data: time series data with datetime index
        file_path = os.path.join(self.test_dir, "test.def")

        with self.assertLogs("pyedautils.data_io", level="INFO") as cm:
            load_data(file_path)

        log_output = "\n".join(cm.output)
        self.assertIn("Loading data from:", log_output)
        self.assertIn("Unsupported file format", log_output)
        
        
if __name__ == "__main__":
    unittest.main() # pragma: no cover
