from django.test import TestCase

from epdb.logic import PackageManager
from epdb.models import Reaction, Compound, User, Rule
from utilities.ml import Dataset


class DatasetTest(TestCase):
    fixtures = ["test_fixture.cleaned.json"]

    def setUp(self):
        self.cs1 = Compound.create(
            self.package,
            name='2,6-Dibromohydroquinone',
            description='http://localhost:8000/package/32de3cf4-e3e6-4168-956e-32fa5ddb0ce1/compound/d6435251-1a54-4327-b4b1-fd6e9a8f4dc9/structure/d8a0225c-dbb5-4e6c-a642-730081c09c5b',
            smiles='C1=C(C(=C(C=C1O)Br)O)Br',
        ).default_structure

        self.cs2 = Compound.create(
            self.package,
            smiles='O=C(O)CC(=O)/C=C(/Br)C(=O)O',
        ).default_structure

        self.rule1 = Rule.create(
            rule_type='SimpleAmbitRule',
            package=self.package,
            smirks='[#8:8]([H])-[c:4]1[c:3]([H])[c:2](-[#1,#17,#35:9])[c:1](-[#8:7]([H]))[c:6](-[#1,#17,#35])[c:5]([H])1>>[#8-]-[#6:6](=O)-[#6:5]-[#6:4](=[O:8])\[#6:3]=[#6:2](\[#1,#17,#35:9])-[#6:1](-[#8-])=[O:7]',
            description='http://localhost:8000/package/32de3cf4-e3e6-4168-956e-32fa5ddb0ce1/simple-ambit-rule/f6a56c0f-a4a0-4ee3-b006-d765b4767cf6'
        )

        self.reaction1 = Reaction.create(
            package=self.package,
            educts=[self.cs1],
            products=[self.cs2],
            rules=[self.rule1],
            multi_step=False
        )

    @classmethod
    def setUpClass(cls):
        super(DatasetGeneratorTest, cls).setUpClass()
        cls.user = User.objects.get(username='anonymous')
        cls.package = PackageManager.create_package(cls.user, 'Anon Test Package', 'No Desc')

    def test_smoke(self):
        reactions = [r for r in Reaction.objects.filter(package=self.package)]
        applicable_rules = [self.rule1]

        ds = Dataset.generate_dataset(reactions, applicable_rules)

        self.assertEqual(len(ds.y()), 1)
        self.assertEqual(sum(ds.y()[0]), 1)