import unittest
import riptable as rt


class AccumTable_Test(unittest.TestCase):
    def test_accum_cols(self):

        num_rows = 10
        data = rt.Dataset(
            {
                'Symb': rt.Cat(['A', 'B'] * int(num_rows / 2)),
                'Count': rt.full(num_rows, 1.0),
                'PlusMinus': [1.0, -1.0]
                * int(num_rows / 2),  # Added to handle edge case of zero footer
            }
        )

        accum = rt.accum_cols(
            data.Symb, [data.Count, data.PlusMinus], ['Count', 'PlusMinus']
        )
        accum_expected = rt.Dataset(
            {'Symb': ['A', 'B'], 'Count': [5.0, 5.0], 'PlusMinus': [5.0, -5.0]}
        )
        accum_expected.footer_set_values(
            'Total', {'Symb': 'Total', 'Count': 10.0, 'PlusMinus': 0.0}
        )
        self.assertTrue((accum == accum_expected).all(axis=None))

    def test_accum_cols_multikey(self):
        num_rows = 12
        data = rt.Dataset(
            {
                'Symb': rt.Cat(['A', 'B'] * int(num_rows / 2)),
                'Exch': rt.Cat(['X', 'Y', 'Y', 'X'] * int(num_rows / 4)),
                'Count': rt.full(num_rows, 1.0),
                'PlusMinus': [1.0, -1.0] * int(num_rows / 2),
            }
        )
        data.MultiKeyCat = rt.Cat([data.Symb, data.Exch])

        accum = rt.accum_cols(
            data.MultiKeyCat, [data.Count, data.PlusMinus], ['Count', 'PlusMinus']
        )
        accum_expected = rt.Dataset(
            {
                'Symb': ['A', 'B', 'A', 'B'],
                'Exch': ['X', 'Y', 'Y', 'X'],
                'Count': [3.0, 3.0, 3.0, 3.0],
                'PlusMinus': [3.0, -3.0, 3.0, -3.0],
            }
        )
        accum_expected.footer_set_values(
            'Total', {'Exch': 'Total', 'Count': 12.0, 'PlusMinus': 0.0}
        )

        self.assertTrue((accum == accum_expected).all(axis=None))

    # When a raw FA is passed as the pointers instead of Categorical
    def test_accum_cols_noncat(self):
        num_rows = 10
        pointer = rt.FA([0, 1] * int(num_rows / 2))
        count = rt.full(num_rows, 1.0)

        accum = rt.accum_cols(pointer, count)
        accum_expected = rt.Dataset({'YLabel': [0, 1], 'col0': [5.0, 5.0]})
        accum_expected.footer_set_values('Total', {'YLabel': 'Total', 'col0': 10.0})

        self.assertTrue((accum == accum_expected).all(axis=None))

    # Test basic accum_ratiop
    def test_accum_ratiop(self):

        num_rows = 12
        data = rt.Dataset(
            {
                'Symb': rt.Cat(['A', 'A', 'A', 'B'] * int(num_rows / 4)),
                'Exch': rt.Cat(['Z', 'Z', 'X', 'X'] * int(num_rows / 4)),
                'Count': rt.full(num_rows, 1.0),
            }
        )

        # Invalid input
        with self.assertRaises(
            ValueError, msg=f'Failed to raise an error when passing invalid norm_by arg'
        ):
            rt.accum_ratiop(data.Symb, data.Exch, data.Count, norm_by='z')

        # Ratio within total
        accum = rt.accum_ratiop(data.Symb, data.Exch, data.Count, norm_by='T')
        accum_expected = rt.Dataset(
            {
                'Symb': ['A', 'B'],
                'X': [25.0, 25.0],
                'Z': [50.0, 0.0],
                'TotalRatio': [75.0, 25.0],
                'Total': [9.0, 3.0],
            }
        )
        accum_expected.footer_set_values(
            'TotalRatio',
            {'Symb': 'TotalRatio', 'X': 50.0, 'Z': 50.0, 'TotalRatio': 100.0},
        )
        accum_expected.footer_set_values(
            'Total', {'Symb': 'Total', 'X': 6.0, 'Z': 6.0, 'Total': 12.0}
        )
        self.assertTrue((accum == accum_expected).all(axis=None))

        # Ratio within columns
        accum = rt.accum_ratiop(data.Symb, data.Exch, data.Count, norm_by='c')
        accum_expected = rt.Dataset(
            {
                'Symb': ['A', 'B'],
                'X': [50.0, 50.0],
                'Z': [100.0, 0.0],
                'TotalRatio': [75.0, 25.0],
                'Total': [9.0, 3.0],
            }
        )
        accum_expected.footer_set_values(
            'TotalRatio',
            {'Symb': 'TotalRatio', 'X': 100.0, 'Z': 100.0, 'TotalRatio': 100.0},
        )
        accum_expected.footer_set_values(
            'Total', {'Symb': 'Total', 'X': 6.0, 'Z': 6.0, 'Total': 12.0}
        )
        self.assertTrue((accum == accum_expected).all(axis=None))


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