# -*- coding: utf-8 -*-

from .context import bamboo
import pandas as pd
from pandas._testing import assert_frame_equal
import unittest


class BasicTestSuite(unittest.TestCase):
    """Basic test cases."""

    def test_columns_to_rows_keep_one_columns(self):
        df = pd.DataFrame(
            {
                "A": {0: "a", 1: "b", 2: "c", 3: "c"},
                "B": {0: 1, 1: 3, 2: 5, 3: 455},
                "C": {0: 2, 1: 4, 2: 6, 3: 100},
                "D": {0: "A", 1: "B", 2: "C", 3: "D"},
            }
        )

        dflong = bamboo.columns_to_rows(df, columns_to_keep=["A"])
        dfref = pd.read_json(
            "tests/df_test_columns_to_rows_keep_one_column.json",
            typ="frame",
        )
        assert_frame_equal(dflong, dfref)

    def test_columns_to_rows_two_columns(self):
        df = pd.DataFrame(
            {
                "A": {0: "a", 1: "b", 2: "c", 3: "c"},
                "B": {0: 1, 1: 3, 2: 5, 3: 455},
                "C": {0: 2, 1: 4, 2: 6, 3: 100},
                "D": {0: "A", 1: "B", 2: "C", 3: "D"},
            }
        )
        dflong = bamboo.columns_to_rows(df, columns_to_keep=["A", "B"])
        dfref = pd.read_json(
            "tests/df_test_columns_to_rows_keep_two_columns.json",
            typ="frame",
        )

        assert_frame_equal(dflong, dfref)


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