from unittest import TestCase
from pathlib import Path

from src.tacq import Relation
from src.tacq.CSP import file_to_examples


class TestTemplateAcquisition(TestCase):
    def test_learn_from_file_csp(self):
        from src.tacq.TemplateAcquisition import TemplateAcquisition
        tacq = TemplateAcquisition()
        # Try to find the file in ../ and then in ./
        data_path = (Path(__file__).parent.parent / "data" / "examtimetabling" /
                     "examtimetabling_4s-3cps-3days-2slots-2r_train.csv")
        self.assertTrue(data_path.exists())
        file_train = str(data_path)
        csp = tacq.learn_from_file(file_train=file_train, max_examples=1000, timeout=None, verbose=True, max_cpu=1)
        self.assertIsNotNone(csp)
        self.assertIsNotNone(tacq.get_network())
        self.assertEqual(len(tacq.get_network().variables), 12)
        expectedRelation1 = Relation(accepted_tuples=[(4, 4), (8, 8), (5, 5), (7, 7), (0, 0), (9, 9), (1, 1), (10, 10),
                                                      (3, 3), (2, 2), (6, 6), (11, 11)])
        expectedScopes1 =  [[0, 1], [1, 0], [0, 2], [2, 0], [0, 3], [3, 0], [0, 4], [4, 0], [0, 5], [5, 0], [0, 6],
                           [6, 0], [0, 7], [7, 0], [0, 8], [8, 0], [0, 9], [9, 0], [0, 10], [10, 0], [0, 11], [11, 0],
                           [1, 2], [2, 1], [1, 3], [3, 1], [1, 4], [4, 1], [1, 5], [5, 1], [1, 6], [6, 1], [1, 7],
                           [7, 1], [1, 8], [8, 1], [1, 9], [9, 1], [1, 10], [10, 1], [1, 11], [11, 1], [2, 3], [3, 2],
                           [2, 4], [4, 2], [2, 5], [5, 2], [2, 6], [6, 2], [2, 7], [7, 2], [2, 8], [8, 2], [2, 9],
                           [9, 2], [2, 10], [10, 2], [2, 11], [11, 2], [3, 4], [4, 3], [3, 5], [5, 3], [3, 6], [6, 3],
                           [3, 7], [7, 3], [3, 8], [8, 3], [3, 9], [9, 3], [3, 10], [10, 3], [3, 11], [11, 3], [4, 5],
                           [5, 4], [4, 6], [6, 4], [4, 7], [7, 4], [4, 8], [8, 4], [4, 9], [9, 4], [4, 10], [10, 4],
                           [4, 11], [11, 4], [5, 6], [6, 5], [5, 7], [7, 5], [5, 8], [8, 5], [5, 9], [9, 5], [5, 10],
                           [10, 5], [5, 11], [11, 5], [6, 7], [7, 6], [6, 8], [8, 6], [6, 9], [9, 6], [6, 10], [10, 6],
                           [6, 11], [11, 6], [7, 8], [8, 7], [7, 9], [9, 7], [7, 10], [10, 7], [7, 11], [11, 7], [8, 9],
                           [9, 8], [8, 10], [10, 8], [8, 11], [11, 8], [9, 10], [10, 9], [9, 11], [11, 9], [10, 11],
                           [11, 10]]
        expectedRelation2 = Relation(accepted_tuples=[(5, 4), (4, 6), (5, 7), (8, 9), (9, 8), (10, 9), (9, 11), (11, 8),
                                                      (4, 5), (5, 6), (8, 11), (9, 10), (11, 10), (10, 8), (7, 6),
                                                      (4, 7), (8, 10), (11, 9), (7, 5)])
        expectedScopes2 = [[0, 1], [0, 2], [1, 0], [1, 2], [2, 0], [2, 1], [3, 4], [3, 5], [4, 3], [4, 5], [5, 3],
                           [5, 4], [6, 7], [6, 8], [7, 6], [7, 8], [8, 6], [8, 7], [9, 10], [9, 11], [10, 9], [10, 11],
                           [11, 9], [11, 10]]
        for scope, relation in csp.constraints:
            if relation == expectedRelation1:
                self.assertIn([scope[0].get_name(), scope[1].get_name()], expectedScopes1)
                expectedScopes1.remove([scope[0].get_name(), scope[1].get_name()])
            elif relation == expectedRelation2:
                self.assertIn([scope[0].get_name(), scope[1].get_name()], expectedScopes2)
                expectedScopes2.remove([scope[0].get_name(), scope[1].get_name()])
            else:
                self.fail(f"Unexpected relation found: {relation}")
        self.assertEqual(len(expectedScopes1), 0, "Some expected scopes were not found in the CSP.")
        self.assertEqual(len(expectedScopes2), 0, "Some expected scopes were not found in the CSP.")

    def test_learn_from_file_accuracy(self):
        from src.tacq.TemplateAcquisition import TemplateAcquisition
        tacq = TemplateAcquisition()
        data_path = (Path(__file__).parent.parent / "data" / "examtimetabling" /
                     "examtimetabling_4s-3cps-3days-2slots-2r_train.csv")
        self.assertTrue(data_path.exists())
        file_train = str(data_path)
        test_path = (Path(__file__).parent.parent / "data" / "examtimetabling" /
                     "examtimetabling_4s-3cps-3days-2slots-2r_test.csv")
        self.assertTrue(test_path.exists())
        file_test = str(test_path)
        tacq.learn_from_file(file_train=file_train, max_examples=300, timeout=None, verbose=True, max_cpu=1)
        self.assertAlmostEqual(0.97, tacq.get_baseline_network().accuracy(file_to_examples(file_test)), 2)
        self.assertEqual(1.0, tacq.get_network().accuracy(file_to_examples(file_test)),
                         "The accuracy of the learned CSP should be 1.0 on the test set.")
