import unittest
import numpy as np
import pandas as pd
import random
from cp2025.algorithms.AggregationKAnonymityTimeOptimal import AggregationKAnonymityTimeOptimal
from UnorderedClass import UnorderedClass

class TestGeneralizationKAnonymityTimeOptimal(unittest.TestCase):

    def test_initially_k_anonymus(self):
        df = [
            [1, 1, 1, 1],
            [1, 1, 1, 1],
            [2, 2, 2, 2],
            [2, 2, 2, 2],
        ]
        k = 2
        k_anonymus_df, k_replaced = AggregationKAnonymityTimeOptimal(k, ['real']*4).depersonalize(df)
        self.assertEqual(df, k_anonymus_df)
        self.assertEqual(k_replaced, 0)

    def test_generalize_last_element(self):
        df = [
            [1, 1, 1, 1],
            [1, 1, 1, 1],
            [2, 2, 2, 2],
            [2, 2, 2, 3],
        ]
        k = 2
        k_anonymus_df, k_replaced = AggregationKAnonymityTimeOptimal(k, ['real']*4).depersonalize(df)
        self.assertEqual(k_replaced, 2)

    def test_generalize_everything(self):
        df = [
            [1, 1, 1, 1],
            [2, 2, 2, 2],
            [3, 3, 3, 3],
            [4, 4, 4, 4],
        ]
        k = 4
        k_anonymus_df, k_replaced = AggregationKAnonymityTimeOptimal(k, ['real']*4).depersonalize(df)
        none_df = [[2.5] * 4]*4
        self.assertEqual(k_anonymus_df, none_df)
        self.assertEqual(k_replaced, 16)

    def test_big_k(self):
        df = [
            [1, 1, 1, 1],
            [1, 1, 1, 1],
            [1, 1, 1, 1],
            [1, 1, 1, 1],
        ]
        k = 5
        k_anonymus_df, k_replaced = AggregationKAnonymityTimeOptimal(k, ['real']*4).depersonalize(df)
        self.assertEqual(k_anonymus_df, None)
        self.assertEqual(k_replaced, None)

    def test_k_equals_one(self):
        df = [
            [1, 1, 1, 1],
            [2, 2, 2, 2],
            [3, 3, 3, 3],
            [4, 4, 4, 4],
        ]
        k = 1
        k_anonymus_df, k_replaced = AggregationKAnonymityTimeOptimal(k, ['real']*4).depersonalize(df)
        self.assertEqual(k_anonymus_df, df)
        self.assertEqual(k_replaced, 0)

    def test_numpy(self):
        df = [
            [1, 1, 1, 1],
            [1, 1, 1, 1],
            [2, 2, 2, 2],
            [2, 2, 2, 2],
        ]
        df = np.array(df)
        k = 2
        k_anonymus_df, k_replaced = AggregationKAnonymityTimeOptimal(k, ['real']*4).depersonalize(df)
        self.assertTrue((df == k_anonymus_df).all())
        self.assertEqual(k_replaced, 0)

    def test_pandas(self):
        df = [
            [1, 1, 1, 1],
            [1, 1, 1, 1],
            [2, 2, 2, 2],
            [2, 2, 2, 2],
        ]
        df = pd.DataFrame(data=np.array(df), columns=['a', 'b', 'c', 'd'])
        k = 2
        k_anonymus_df, k_replaced = AggregationKAnonymityTimeOptimal(k, ['real'] * 4).depersonalize(df)
        self.assertTrue((df == k_anonymus_df).all().all())
        self.assertEqual(k_replaced, 0)

    def test_string_data(self):
        df = [
            ["a", "a", "a", "a"],
            ["a", "a", "a", "a"],
            ["b", "b", "b", "b"],
            ["b", "b", "b", "b"],
        ]
        k = 2
        k_anonymus_df, k_replaced = AggregationKAnonymityTimeOptimal(k, ['ordered']*4).depersonalize(df)
        self.assertEqual(df, k_anonymus_df)
        self.assertEqual(k_replaced, 0)

    def test_float_data(self):
        df = [
            [1.0, 1.0, 1.0, 1.0],
            [1.0, 1.0, 1.0, 1.0],
            [2.0, 2.0, 2.0, 2.0],
            [2.0, 2.0, 2.0, 2.0],
        ]
        k = 2
        k_anonymus_df, k_replaced = AggregationKAnonymityTimeOptimal(k, ['real']*4).depersonalize(df)
        self.assertEqual(df, k_anonymus_df)
        self.assertEqual(k_replaced, 0)

    def test_mixed_data(self):
        df = [
            [1.0, 1.0, "a", 1],
            [1.0, 1.0, "a", 1],
            [2.0, 2.0, "b", 2],
            [2.0, 2.0, "b", 2],
        ]
        k = 2
        k_anonymus_df, k_replaced = AggregationKAnonymityTimeOptimal(k, ['real', 'real', 'ordered', 'real']).depersonalize(df)
        self.assertEqual(df, k_anonymus_df)
        self.assertEqual(k_replaced, 0)

    def test_truly_unordered(self):
        df = [
            [UnorderedClass(1), UnorderedClass(1), UnorderedClass(1), UnorderedClass(1)],
            [UnorderedClass(1), UnorderedClass(1), UnorderedClass(1), UnorderedClass(1)],
            [UnorderedClass(2), UnorderedClass(2), UnorderedClass(2), UnorderedClass(2)],
            [UnorderedClass(2), UnorderedClass(2), UnorderedClass(2), UnorderedClass(3)],
            [UnorderedClass(2), UnorderedClass(2), UnorderedClass(2), UnorderedClass(3)],
        ]
        df_expected = [
            [UnorderedClass(1), UnorderedClass(1), UnorderedClass(1), UnorderedClass(1)],
            [UnorderedClass(1), UnorderedClass(1), UnorderedClass(1), UnorderedClass(1)],
            [UnorderedClass(2), UnorderedClass(2), UnorderedClass(2), UnorderedClass(3)],
            [UnorderedClass(2), UnorderedClass(2), UnorderedClass(2), UnorderedClass(3)],
            [UnorderedClass(2), UnorderedClass(2), UnorderedClass(2), UnorderedClass(3)],
        ]
        k = 2
        k_anonymus_df, k_replaced = AggregationKAnonymityTimeOptimal(k, ['unordered'] * 4).depersonalize(df)
        self.assertEqual(df_expected, k_anonymus_df)
        self.assertEqual(k_replaced, 1)

    def test_random_df(self):
        seed = 1234
        random.seed(seed)
        np.random.seed(seed)
        for i in range(1000):
            rows = random.randint(4, 50)
            cols = random.randint(1, 5)
            df = np.random.randint(0, 3, (rows, cols))
            k = random.randint(2, 4)
            quasi_identifiers_types = []
            for j in range(cols):
                quasi_identifiers_types.append(random.choice(['real', 'ordered', 'unordered']))
            k_anonymus_df, k_replaced = AggregationKAnonymityTimeOptimal(k, quasi_identifiers_types).depersonalize(df)



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