from unittest.mock import Mock, patch

from django.test import TestCase

from epdb.logic import SEdge, SNode, SPathway
from epdb.models import Pathway, Setting
from utilities.chem import PredictionResult, ProductSet


class SNodeTest(TestCase):
    def test_snode_eq(self):
        snode1 = SNode("CN1C2C(N(C(N(C)C=2N=C1)=O)C)=O", 0)
        snode2 = SNode("CN1C2C(N(C(N(C)C=2N=C1)=O)C)=O", 0)
        self.assertEqual(snode1, snode2)


class SEdgeTest(TestCase):
    def test_sedge_eq(self):
        sedge1 = SEdge(
            [SNode("CN1C2C(N(C(N(C)C=2N=C1)=O)C)=O", 0)],
            [SNode("CN1C(=O)NC2=C(C1=O)N(C)C=N2", 1), SNode("C=O", 1)],
            rule=None,
        )
        sedge2 = SEdge(
            [SNode("CN1C2C(N(C(N(C)C=2N=C1)=O)C)=O", 0)],
            [SNode("CN1C(=O)NC2=C(C1=O)N(C)C=N2", 1), SNode("C=O", 1)],
            rule=None,
        )
        self.assertEqual(sedge1, sedge2)


class SPathwayTest(TestCase):
    def setUp(self):
        """Set up test data for SPathway tests."""
        self.test_smiles = "CCN(CC)C(=O)C1=CC(=CC=C1)CO"
        self.mock_setting = Mock(spec=Setting)
        self.mock_pathway = Mock(spec=Pathway)

    def test_predict_step_basic(self):
        """Test basic predict_step functionality."""
        spw = SPathway(root_nodes=self.test_smiles, prediction_setting=self.mock_setting)

        # e.g. bt0002
        pr = PredictionResult(
            [
                ProductSet(["CC1=CC=C(C2OC(CO)C(=O)C(O)C2O)C=C1CC1=CC=C(C2=CC=C(F)C=C2)S1"]),
                ProductSet(["CC1=CC=C(C2OC(CO)C(O)C(O)C2=O)C=C1CC1=CC=C(C2=CC=C(F)C=C2)S1"]),
                ProductSet(["CC1=CC=C(C2OC(CO)C(O)C(=O)C2O)C=C1CC1=CC=C(C2=CC=C(F)C=C2)S1"]),
            ],
            0.17,
            None,
        )

        mock_val = {"rule_triggered": True, "transformations": [pr]}

        with patch.object(self.mock_setting, "expand", return_value=mock_val):
            spw.predict_step(from_depth=0)

        self.assertEqual(len(spw.smiles_to_node.keys()), 4)
        self.assertEqual(len(spw.edges), 3)

    def test_to_json(self):
        """Test basic predict_step functionality."""
        spw = SPathway(root_nodes=self.test_smiles, prediction_setting=self.mock_setting)

        # e.g. bt0002
        pr = PredictionResult(
            [
                ProductSet(["CC1=CC=C(C2OC(CO)C(=O)C(O)C2O)C=C1CC1=CC=C(C2=CC=C(F)C=C2)S1"]),
                ProductSet(["CC1=CC=C(C2OC(CO)C(O)C(O)C2=O)C=C1CC1=CC=C(C2=CC=C(F)C=C2)S1"]),
                ProductSet(["CC1=CC=C(C2OC(CO)C(O)C(=O)C2O)C=C1CC1=CC=C(C2=CC=C(F)C=C2)S1"]),
            ],
            0.17,
            None,
        )

        mock_val = {"rule_triggered": True, "transformations": [pr]}

        with patch.object(self.mock_setting, "expand", return_value=mock_val):
            spw.predict_step(from_depth=0)

        self.assertEqual(len(spw.smiles_to_node.keys()), 4)
        self.assertEqual(len(spw.edges), 3)

        json_result = spw.to_json()

        self.assertIsInstance(json_result, dict)
        self.assertIn("nodes", json_result)
        self.assertIn("edges", json_result)
        self.assertEqual(len(json_result["nodes"]), 4)
        self.assertEqual(len(json_result["edges"]), 3)