enviPy-bayer/bayer/models.py

from typing import List
import urllib.parse
import nh3
from django.conf import settings as s
from django.db import models, transaction
from django.db.models import QuerySet
from django.urls import reverse

from epdb.models import (
    EnviPathModel,
    Compound,
    CompoundStructure,
    ParallelRule,
    SequentialRule,
    SimpleAmbitRule,
    SimpleRDKitRule,
)
from utilities.chem import FormatConverter


class Package(EnviPathModel):
    reviewed = models.BooleanField(verbose_name="Reviewstatus", default=False)
    license = models.ForeignKey(
        "epdb.License", on_delete=models.SET_NULL, blank=True, null=True, verbose_name="License"
    )

    class Classification(models.IntegerChoices):
        INTERNAL = 0, "Internal"
        RESTRICTED = 10 , "Restricted"
        SECRET = 20, "Secret"

    classification_level = models.IntegerField(
        choices=Classification,
        default=Classification.RESTRICTED,
    )

    data_pool = models.ForeignKey("epdb.Group", on_delete=models.SET_NULL, blank=True, null=True,
                                  verbose_name="Data pool", default=None)

    def delete(self, *args, **kwargs):
        # explicitly handle related Rules
        for r in self.rules.all():
            r.delete()
        super().delete(*args, **kwargs)

    def __str__(self):
        return f"{self.name} (pk={self.pk})"

    @property
    def compounds(self) -> QuerySet:
        return self.compound_set.all()

    @property
    def rules(self) -> QuerySet:
        return self.rule_set.all()

    @property
    def reactions(self) -> QuerySet:
        return self.reaction_set.all()

    @property
    def pathways(self) -> QuerySet:
        return self.pathway_set.all()

    @property
    def scenarios(self) -> QuerySet:
        return self.scenario_set.all()

    @property
    def models(self) -> QuerySet:
        return self.epmodel_set.all()

    def _url(self):
        return "{}/package/{}".format(s.SERVER_URL, self.uuid)

    def get_applicable_rules(self) -> List["Rule"]:
        """
        Returns a ordered set of rules where the following applies:
        1. All Composite will be added to result
        2. All SimpleRules will be added if theres no CompositeRule present using the SimpleRule
        Ordering is based on "url" field.
        """
        rules = []
        rule_qs = self.rules

        reflected_simple_rules = set()

        for r in rule_qs:
            if isinstance(r, ParallelRule) or isinstance(r, SequentialRule):
                rules.append(r)
                for sr in r.simple_rules.all():
                    reflected_simple_rules.add(sr)

        for r in rule_qs:
            if isinstance(r, SimpleAmbitRule) or isinstance(r, SimpleRDKitRule):
                if r not in reflected_simple_rules:
                    rules.append(r)

        rules = sorted(rules, key=lambda x: x.url)
        return rules

    class Meta:
        db_table = "epdb_package"


class PESCompound(Compound):

    @staticmethod
    @transaction.atomic
    def create(
        package: "Package", pes_data: dict, name: str = None, description: str = None, *args, **kwargs
    ) -> "Compound":

        pes_url = pes_data["pes_url"]

        # Check if we find a direct match for a given pes_link
        if PESStructure.objects.filter(pes_link=pes_url, compound__package=package).exists():
            # Due to normalization we might end up in having multiple structures
            # All of them point to the same compound -> pick any
            return PESStructure.objects.filter(pes_link=pes_url, compound__package=package).first().compound

        # Generate Compound
        c = PESCompound()
        c.package = package

        if name is not None:
            # Clean for potential XSS
            name = nh3.clean(name, tags=s.ALLOWED_HTML_TAGS).strip()

        if name is None or name == "":
            name = f"Compound {Compound.objects.filter(package=package).count() + 1}"

        c.name = name

        # We have a default here only set the value if it carries some payload
        if description is not None and description.strip() != "":
            c.description = nh3.clean(description, tags=s.ALLOWED_HTML_TAGS).strip()

        c.save()

        molfile = pes_data.get("representativeStructures", [{}])[0].get("ctab")

        if molfile is None:
            raise ValueError("PES data does not contain a valid mol file!")

        smiles = FormatConverter.to_smiles(FormatConverter.from_molfile(molfile))

        standardized_smiles = FormatConverter.standardize(smiles, remove_stereo=True)

        is_standardized = standardized_smiles == smiles

        if not is_standardized:
            _ = PESStructure.create(
                c,
                pes_url,
                molfile,
                standardized_smiles,
                name="Normalized structure of {}".format(name),
                description="{} (in its normalized form)".format(description),
                normalized_structure=True,
            )


        cs = PESStructure.create(
            c,
            pes_url,
            molfile,
            smiles,
            name=name,
            description=description,
            normalized_structure=is_standardized
        )

        c.default_structure = cs
        c.save()

        return c


class PESStructure(CompoundStructure):
    pes_link = models.URLField(blank=False, null=False, verbose_name="PES Link")

    @staticmethod
    @transaction.atomic
    def create(
            compound: Compound,
            pes_link: str,
            mol_file: str,
            smiles: str,
            name: str = None,
            description: str = None,
            *args,
            **kwargs
    ):
        if compound.pk is None:
            raise ValueError("Unpersisted Compound! Persist compound first!")

        cs = PESStructure()
        # Clean for potential XSS
        if name is not None:
            cs.name = nh3.clean(name, tags=s.ALLOWED_HTML_TAGS).strip()

        if description is not None:
            cs.description = nh3.clean(description, tags=s.ALLOWED_HTML_TAGS).strip()

        cs.smiles = smiles
        cs.mol_file = mol_file
        cs.pes_link = pes_link
        cs.compound = compound

        if "normalized_structure" in kwargs:
            cs.normalized_structure = kwargs["normalized_structure"]

        cs.save()

        return cs

    @transaction.atomic
    def add_structure(
        self,
        smiles: str,
        name: str = None,
        description: str = None,
        default_structure: bool = False,
        *args,
        **kwargs,
    ) -> "CompoundStructure":
        raise ValueError("Not supported!")

    def d3_json(self):
        return {
            "is_pes": True,
            "pes_link": self.pes_link,
            # Will overwrite image from Node
            "image": f"{reverse("depict_pes")}?pesLink={urllib.parse.quote(self.pes_link)}"
        }