from tempfile import TemporaryDirectory

import numpy as np
from django.conf import settings as s
from django.test import TestCase, override_settings

from epdb.logic import PackageManager
from epdb.models import MLRelativeReasoning, RuleBasedRelativeReasoning, User

Package = s.GET_PACKAGE_MODEL()


@override_settings(MODEL_DIR=s.FIXTURE_DIRS[0] / "models", CELERY_TASK_ALWAYS_EAGER=True)
class ModelTest(TestCase):
    fixtures = ["test_fixtures_incl_model.jsonl.gz"]

    @classmethod
    def setUpClass(cls):
        super(ModelTest, cls).setUpClass()
        cls.user = User.objects.get(username="anonymous")
        cls.package = PackageManager.create_package(cls.user, "Anon Test Package", "No Desc")
        cls.BBD_SUBSET = Package.objects.get(name="Fixtures")

    def test_mlrr(self):
        with TemporaryDirectory() as tmpdir:
            with self.settings(MODEL_DIR=tmpdir):
                threshold = float(0.5)

                rule_package_objs = [self.BBD_SUBSET]
                data_package_objs = [self.BBD_SUBSET]
                eval_packages_objs = [self.BBD_SUBSET]

                mod = MLRelativeReasoning.create(
                    self.package,
                    rule_package_objs,
                    data_package_objs,
                    threshold=threshold,
                    name="ECC - BBD - 0.5",
                    description="Created MLRelativeReasoning in Testcase",
                )

                mod.build_dataset()
                mod.build_model()
                mod.evaluate_model(True, eval_packages_objs, n_splits=2)

                results = mod.predict("CCN(CC)C(=O)C1=CC(=CC=C1)C")

                products = dict()
                for r in results:
                    for ps in r.product_sets:
                        products[tuple(sorted(ps.product_set))] = (r.rule.name, r.probability)

                expected = {
                    ("CC=O", "CCNC(=O)C1=CC(C)=CC=C1"): (
                        "bt0243-4301",
                        np.float64(0.5),
                    ),
                    ("CC1=CC=CC(C(=O)O)=C1", "CCNCC"): ("bt0430-4011", np.float64(0.25)),
                }

                self.assertEqual(products, expected)

                # from pprint import pprint
                # pprint(mod.eval_results)

    def test_applicability(self):
        with TemporaryDirectory() as tmpdir:
            with self.settings(MODEL_DIR=tmpdir):
                threshold = float(0.5)

                rule_package_objs = [self.BBD_SUBSET]
                data_package_objs = [self.BBD_SUBSET]
                eval_packages_objs = [self.BBD_SUBSET]

                mod = MLRelativeReasoning.create(
                    self.package,
                    rule_package_objs,
                    data_package_objs,
                    threshold=threshold,
                    name="ECC - BBD - 0.5",
                    description="Created MLRelativeReasoning in Testcase",
                    build_app_domain=True,  # To test the applicability domain this must be True
                    app_domain_num_neighbours=5,
                    app_domain_local_compatibility_threshold=0.5,
                    app_domain_reliability_threshold=0.5,
                )

                mod.build_dataset()
                mod.build_model()
                mod.evaluate_model(True, eval_packages_objs, n_splits=2)

                _ = mod.predict("CCN(CC)C(=O)C1=CC(=CC=C1)C")

    def test_rbrr(self):
        with TemporaryDirectory() as tmpdir:
            with self.settings(MODEL_DIR=tmpdir):
                threshold = float(0.5)

                rule_package_objs = [self.BBD_SUBSET]
                data_package_objs = [self.BBD_SUBSET]
                eval_packages_objs = [self.BBD_SUBSET]

                mod = RuleBasedRelativeReasoning.create(
                    self.package,
                    rule_package_objs,
                    data_package_objs,
                    threshold=threshold,
                    min_count=5,
                    max_count=0,
                    name="ECC - BBD - 0.5",
                    description="Created MLRelativeReasoning in Testcase",
                )

                mod.build_dataset()
                mod.build_model()
                mod.evaluate_model(True, eval_packages_objs, n_splits=2)

                _ = mod.predict("CCN(CC)C(=O)C1=CC(=CC=C1)C")
