from django.test import TestCase, override_settings
from django.urls import reverse
from django.conf import settings as s

from epdb.logic import UserManager, PackageManager
from epdb.models import Pathway, Edge

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

    @classmethod
    def setUpClass(cls):
        super(PathwayViewTest, cls).setUpClass()
        cls.user1 = UserManager.create_user("user1", "user1@envipath.com", "SuperSafe",
                                            set_setting=True, add_to_group=True, is_active=True)
        cls.user1_default_package = cls.user1.default_package
        cls.package = PackageManager.create_package(cls.user1, 'Test', 'Test Pack')

    def setUp(self):
        self.client.force_login(self.user1)

    def test_predict_pathway(self):
        response = self.client.post(reverse("pathways"), {
            'name': 'Test Pathway',
            'description': 'Just a Description',
            'predict': 'predict',
            'smiles': 'CCN(CC)C(=O)C1=CC(=CC=C1)CO',
        })
        self.assertEqual(response.status_code, 302)
        pathway_url = response.url

        pw = Pathway.objects.get(url=pathway_url)
        self.assertEqual(self.user1_default_package, pw.package)

        self.assertEqual(pw.name, 'Test Pathway')
        self.assertEqual(pw.description, 'Just a Description')
        self.assertEqual(len(pw.root_nodes), 1)
        self.assertEqual(pw.root_nodes.first().default_node_label.smiles, 'CCN(CC)C(=O)C1=CC(CO)=CC=C1')

        first_level_nodes = {
            # Edge 1
            'CCN(CC)C(=O)C1=CC(C=O)=CC=C1',
            # Edge 2
            'CCNC(=O)C1=CC(CO)=CC=C1',
            'CC=O',
            # Edge 3
            'CCNCC',
            'O=C(O)C1=CC(CO)=CC=C1',
        }

        predicted_nodes = set()
        edges = Edge.objects.filter(start_nodes__in=[pw.root_nodes.first()])

        for edge in edges:
            for n in edge.end_nodes.all():
                predicted_nodes.add(n.default_node_label.smiles)

        self.assertEqual(first_level_nodes, predicted_nodes)

    def test_predict_package_pathway(self):
        response = self.client.post(
            reverse("package pathway list", kwargs={'package_uuid': str(self.package.uuid)}), {
                'name': 'Test Pathway',
                'description': 'Just a Description',
                'predict': 'predict',
                'smiles': 'CCN(CC)C(=O)C1=CC(=CC=C1)CO',
            })
        self.assertEqual(response.status_code, 302)
        pathway_url = response.url

        pw = Pathway.objects.get(url=pathway_url)
        self.assertEqual(self.package, pw.package)

        self.assertEqual(pw.name, 'Test Pathway')
        self.assertEqual(pw.description, 'Just a Description')
        self.assertEqual(len(pw.root_nodes), 1)
        self.assertEqual(pw.root_nodes.first().default_node_label.smiles, 'CCN(CC)C(=O)C1=CC(CO)=CC=C1')

        first_level_nodes = {
            # Edge 1
            'CCN(CC)C(=O)C1=CC(C=O)=CC=C1',
            # Edge 2
            'CCNC(=O)C1=CC(CO)=CC=C1',
            'CC=O',
            # Edge 3
            'CCNCC',
            'O=C(O)C1=CC(CO)=CC=C1',
        }

        predicted_nodes = set()
        edges = Edge.objects.filter(start_nodes__in=[pw.root_nodes.first()])

        for edge in edges:
            for n in edge.end_nodes.all():
                predicted_nodes.add(n.default_node_label.smiles)

        self.assertEqual(first_level_nodes, predicted_nodes)