enviPy-bayer/utilities/misc.py

import base64
import hashlib
import hmac
import json
import logging
import uuid
from collections import defaultdict
from datetime import datetime
from typing import Any, Dict, List, Optional, TYPE_CHECKING

from django.conf import settings as s
from django.db import transaction
from ninja import Schema
from pydantic import HttpUrl, ValidationError

from bayer.models import PESCompound, PESStructure
from epdb.exceptions import PackageImportException
from epdb.models import (
    AdditionalInformation,
    Compound,
    CompoundStructure,
    Edge,
    EnviFormer,
    EPModel,
    ExternalDatabase,
    ExternalIdentifier,
    License,
    MLRelativeReasoning,
    Node,
    ParallelRule,
    Pathway,
    Reaction,
    Rule,
    RuleBasedRelativeReasoning,
    Scenario,
    Setting,
    SimpleAmbitRule,
    SimpleRDKitRule,
    SimpleRule, Group,
)
from utilities.chem import FormatConverter

logger = logging.getLogger(__name__)
Package = s.GET_PACKAGE_MODEL()

if TYPE_CHECKING:
    from epdb.logic import SPathway, GroupManager


class LicenseExportSchema(Schema):
    cc_string: str
    link: HttpUrl
    image_link: HttpUrl


##############
# RefSchemas #
##############
class RefExportSchema(Schema):
    uuid: str
    url: str

    @staticmethod
    def resolve_uuid(obj):
        value = obj.get("uuid") if isinstance(obj, dict) else obj.uuid
        return str(value)


class RefGroupExportSchema(RefExportSchema): ...


class RefCompoundExportSchema(RefExportSchema): ...


class RefCompoundStructureExportSchema(RefExportSchema): ...


class RefReactionExportSchema(RefExportSchema): ...


class RefRuleExportSchema(RefExportSchema): ...


class RefNodeExportSchema(RefExportSchema): ...


class RefEdgeExportSchema(RefExportSchema): ...


class RefPathwayExportSchema(RefExportSchema): ...


class RefScenarioExportSchema(RefExportSchema): ...


class RefEnzymeExportSchema(RefExportSchema): ...


############
# Compound #
############
class CompoundExportSchema(RefCompoundExportSchema):
    name: str
    description: str
    aliases: List[str]
    default_structure: RefCompoundStructureExportSchema
    structures: List["CompoundStructureExportSchema"]
    scenarios: List[RefScenarioExportSchema]


class CompoundStructureExportSchema(RefCompoundStructureExportSchema):
    name: str
    description: str
    aliases: List[str]
    smiles: str
    molfile: Optional[str]
    normalized_structure: bool
    scenarios: List[RefScenarioExportSchema]


class PESCompoundExportSchema(RefCompoundExportSchema):
    name: str
    description: str
    aliases: List[str]
    default_structure: RefCompoundStructureExportSchema
    structures: List["PESCompoundStructureExportSchema"]
    scenarios: List[RefScenarioExportSchema]


class PESCompoundStructureExportSchema(RefCompoundStructureExportSchema):
    name: str
    description: str
    aliases: List[str]
    smiles: str
    molfile: Optional[str]
    normalized_structure: bool
    scenarios: List[RefScenarioExportSchema]
    pes_link: HttpUrl


############
# Reaction #
############
class ReactionExportSchema(RefReactionExportSchema):
    name: str
    description: str
    aliases: List[str]
    educts: List[RefCompoundStructureExportSchema]
    products: List[RefCompoundStructureExportSchema]
    rules: List[RefRuleExportSchema]
    multi_step: bool
    medline_references: List[str] | None
    scenarios: List[RefScenarioExportSchema]


#########
# Rules #
#########
class EnzymeExportSchema(RefEnzymeExportSchema):
    ec_number: str
    classification_level: int
    linking_method: str
    reaction_evidence: List[RefReactionExportSchema]
    edge_evidence: List[RefEdgeExportSchema]


class EnzymeRuleExportSchema(RefRuleExportSchema):
    enzymes: List[EnzymeExportSchema] | None = None

    @staticmethod
    def resolve_enzymes(obj):
        if isinstance(obj, dict):
            res = []
            for e in obj.get("enzymes", []):
                res.append(EnzymeExportSchema.model_validate(e))
            return res
        return obj.enzymelink_set.all()


class RuleExportSchema(EnzymeRuleExportSchema):
    name: str
    description: str
    aliases: List[str]
    smirks: str
    reactant_filter_smarts: Optional[str]
    product_filter_smarts: Optional[str]
    scenarios: List[RefScenarioExportSchema]


class ParallelRuleExportSchema(EnzymeRuleExportSchema):
    name: str
    description: str
    aliases: List[str]
    simple_rules: List[RefRuleExportSchema]
    scenarios: List[RefScenarioExportSchema]


###########################
# Pathway / Nodes / Edges #
###########################
class NodeExportSchema(RefNodeExportSchema):
    name: str
    description: str
    aliases: List[str]
    default_node_label: RefCompoundStructureExportSchema
    node_labels: List[RefCompoundStructureExportSchema]
    depth: int
    stereo_removed: bool
    scenarios: List[RefScenarioExportSchema]


class EdgeExportSchema(RefEdgeExportSchema):
    name: str
    description: str
    aliases: List[str]
    edge_label: RefReactionExportSchema
    start_nodes: List[RefNodeExportSchema]
    end_nodes: List[RefNodeExportSchema]
    scenarios: List[RefScenarioExportSchema]


class PathwayExportSchema(RefPathwayExportSchema):
    name: str
    description: str
    aliases: List[str]
    predicted: bool
    nodes: List[NodeExportSchema]
    edges: List[EdgeExportSchema]
    scenarios: List[RefScenarioExportSchema]


class ScenarioExportSchema(RefScenarioExportSchema):
    name: str
    description: str
    scenario_date: str
    scenario_type: str


class AdditionalInformationExportSchema(RefExportSchema):
    type: str
    data: dict
    scenario: RefScenarioExportSchema | None = None
    attach_object: RefExportSchema | None = None

    @staticmethod
    def resolve_attach_object(obj):
        if isinstance(obj, dict):
            if obj.get("attach_object") is None:
                return None
            return RefExportSchema.model_validate(obj["attach_object"])
        return obj.content_object


###########
# Package #
###########
class PackageExportSchema(Schema):
    name: str
    description: str
    uuid: str
    reviewed: bool
    license: LicenseExportSchema | None = None
    classification_level: str = "Internal"
    data_pool: RefGroupExportSchema | None = None
    compounds: List[CompoundExportSchema | PESCompoundExportSchema]
    reactions: List[ReactionExportSchema]
    simple_rules: List[RuleExportSchema]
    composite_rules: List[ParallelRuleExportSchema]
    pathways: List[PathwayExportSchema]
    scenarios: List[ScenarioExportSchema]
    additional_information: List[AdditionalInformationExportSchema]

    @staticmethod
    def resolve_uuid(obj):
        value = obj.get("uuid") if isinstance(obj, dict) else obj.uuid
        return str(value)

    @staticmethod
    def resolve_classification_level(obj):
        if isinstance(obj, dict):
            return obj["classification_level"]
        return obj.Classification(obj.classification_level).name


    @staticmethod
    def resolve_compounds(obj):
        res = []
        if isinstance(obj, dict):
            for c in obj.get("compounds", []):
                is_pes = any([cs.get("pes_link", None) is not None for cs in c.get("structures", [])])
                if is_pes:
                    res.append(PESCompoundExportSchema.model_validate(c))
                else:
                    res.append(CompoundExportSchema.model_validate(c))
        else:
            for c in obj.compounds.all():
                if isinstance(c, PESCompound):
                    res.append(PESCompoundExportSchema.from_orm(c))
                else:
                    res.append(CompoundExportSchema.from_orm(c))
        return res

    @staticmethod
    def resolve_simple_rules(obj):
        if isinstance(obj, dict):
            result = []
            for r in obj.get("simple_rules", []):
                result.append(RuleExportSchema.model_validate(r))
            return result
        return SimpleAmbitRule.objects.filter(package=obj)

    @staticmethod
    def resolve_composite_rules(obj):
        if isinstance(obj, dict):
            result = []
            for r in obj.get("composite_rules", []):
                result.append(ParallelRuleExportSchema.model_validate(r))
            return result
        return ParallelRule.objects.filter(package=obj)

    @staticmethod
    def resolve_additional_information(obj):
        if isinstance(obj, dict):
            result = []
            for ai in obj.get("additional_information", []):
                result.append(AdditionalInformationExportSchema.model_validate(ai))
            return result
        return AdditionalInformation.objects.filter(package=obj)


class PackageExporter:
    def __init__(self, package: Package):
        self._raw_package = package

    def do_export(self):
        return PackageExporter._export_package_as_json(self._raw_package)

    @staticmethod
    def _export_package_as_json(package: Package) -> Dict[str, Any]:
        """
        Dumps a Package and all its related objects as JSON.

        Args:
            package: The Package instance to dump

        Returns:
            Dict containing the complete package data as JSON-serializable structure
        """

        data = PackageExportSchema.from_orm(package)

        return data.model_dump(mode="json")


class PackageImporter:

    def __init__(self, package: Dict[str, Any], preserve_uuids: bool = False):
        self.preserve_uuids = preserve_uuids
        self._raw_package = package
        self._cache = {}

    def do_import(self) -> Package:
        return self._import_package_from_json()

    def _import_compounds(
            self,
            package: Package,
            compounds: List[CompoundExportSchema | PESCompoundExportSchema],
    ):
        for c in compounds:

            if isinstance(c, PESCompoundExportSchema):
                c_cls = PESCompound
            else:
                c_cls = Compound

            new_c = c_cls()
            new_c.uuid = str(uuid.uuid4()) if not self.preserve_uuids else c.uuid
            new_c.package = package
            new_c.name = c.name
            new_c.description = c.description
            new_c.aliases = c.aliases
            new_c.save()

            self._cache[c.uuid] = new_c

            for cs in c.structures:

                if isinstance(cs, PESCompoundStructureExportSchema):
                    cs_cls = PESStructure
                else:
                    cs_cls = CompoundStructure

                new_cs = cs_cls()
                new_cs.uuid = str(uuid.uuid4()) if not self.preserve_uuids else cs.uuid
                new_cs.compound = new_c
                new_cs.name = cs.name
                new_cs.description = cs.description
                new_cs.aliases = cs.aliases
                new_cs.smiles = cs.smiles
                new_cs.molfile = cs.molfile

                if cs_cls == PESStructure:
                    new_cs.pes_link = cs.pes_link

                new_cs.normalized_structure = cs.normalized_structure
                new_cs.save()

                self._cache[cs.uuid] = new_cs

            # Now we can assigne the default_structure to the compound
            new_c.default_structure = self._cache[c.default_structure.uuid]
            new_c.save()

    def _import_simple_rules(
            self,
            package: Package,
            rules: List[RuleExportSchema]
    ):
        for r in rules:
            new_r = SimpleAmbitRule()
            new_r.uuid = str(uuid.uuid4()) if not self.preserve_uuids else r.uuid
            new_r.package = package
            new_r.name = r.name
            new_r.description = r.description
            new_r.aliases = r.aliases
            new_r.smirks = r.smirks
            new_r.reactant_filter_smarts = r.reactant_filter_smarts
            new_r.product_filter_smarts = r.product_filter_smarts
            new_r.save()

            self._cache[r.uuid] = new_r

    def _import_composite_rules(
        self,
        package: Package,
        rules: List[ParallelRuleExportSchema]
    ):
        for r in rules:
            new_r = ParallelRule()
            new_r.uuid = str(uuid.uuid4()) if not self.preserve_uuids else r.uuid
            new_r.package = package
            new_r.name = r.name
            new_r.description = r.description
            new_r.aliases = r.aliases
            new_r.save()

            for sr in r.simple_rules:
                new_r.simple_rules.add(self._cache[sr.uuid])

            self._cache[r.uuid] = new_r

    def _import_reactions(
            self,
            package: Package,
            reactions: List[ReactionExportSchema],
    ):
        for r in reactions:
            new_r = Reaction()
            new_r.uuid = str(uuid.uuid4()) if not self.preserve_uuids else r.uuid
            new_r.package = package
            new_r.name = r.name
            new_r.description = r.description
            new_r.aliases = r.aliases
            new_r.multi_step = r.multi_step
            new_r.medline_references = r.medline_references
            new_r.save()

            for educt in r.educts:
                new_r.educts.add(self._cache[educt.uuid])

            for product in r.products:
                new_r.products.add(self._cache[product.uuid])

            for rule in r.rules:
                new_r.rules.add(self._cache[rule.uuid])

            self._cache[r.uuid] = new_r

    def _import_pathways(
            self,
            package: Package,
            pathways: List[PathwayExportSchema]
    ):
        for pw in pathways:
            new_pw = Pathway()
            new_pw.uuid = str(uuid.uuid4()) if not self.preserve_uuids else pw.uuid
            new_pw.package = package
            new_pw.name = pw.name
            new_pw.description = pw.description
            new_pw.aliases = pw.aliases
            new_pw.predicted = pw.predicted
            new_pw.save()

            self._cache[pw.uuid] = new_pw

            for n in pw.nodes:
                new_n = Node()
                new_n.uuid = str(uuid.uuid4()) if not self.preserve_uuids else n.uuid
                new_n.pathway = new_pw
                new_n.name = n.name
                new_n.description = n.description
                new_n.aliases = n.aliases
                new_n.default_node_label = self._cache[n.default_node_label.uuid]
                new_n.depth = n.depth
                new_n.stereo_removed = n.stereo_removed
                new_n.save()

                for nl in n.node_labels:
                    new_n.node_labels.add(self._cache[nl.uuid])

                self._cache[n.uuid] = new_n

            for e in pw.edges:
                new_e = Edge()
                new_e.uuid = str(uuid.uuid4()) if not self.preserve_uuids else e.uuid
                new_e.pathway = new_pw
                new_e.name = e.name
                new_e.description = e.description
                new_e.aliases = e.aliases
                new_e.edge_label = self._cache[e.edge_label.uuid]
                new_e.save()

                for sn in e.start_nodes:
                    new_e.start_nodes.add(self._cache[sn.uuid])

                for en in e.end_nodes:
                    new_e.end_nodes.add(self._cache[en.uuid])

                self._cache[e.uuid] = new_e

    def _import_scenarios(
            self,
            package: Package,
            scenarios: List[ScenarioExportSchema]
    ):
        for s in scenarios:
            new_s = Scenario()
            new_s.uuid = str(uuid.uuid4()) if not self.preserve_uuids else s.uuid
            new_s.package = package
            new_s.name = s.name
            new_s.description = s.description
            new_s.scenario_date = s.scenario_date
            new_s.scenario_type = s.scenario_type
            new_s.save()

            self._cache[s.uuid] = new_s

    def _import_additional_information(
            self,
            package: Package,
            additional_information: List[AdditionalInformationExportSchema]
    ):
        for ai in additional_information:
            new_ai = AdditionalInformation()
            new_ai.uuid = str(uuid.uuid4()) if not self.preserve_uuids else ai.uuid
            new_ai.package = package
            new_ai.type = ai.type
            new_ai.data = ai.data

            if ai.scenario:
                new_ai.scenario = self._cache[ai.scenario.uuid]

            if ai.attach_object:
                new_ai.content_object = self._cache[ai.attach_object.uuid]

            new_ai.save()

            self._cache[ai.uuid] = new_ai

    def _link_scenarios_after_import(self, data: PackageExportSchema):
        collections = [
            data.compounds,
            data.simple_rules,
            data.composite_rules,
            data.reactions,
            data.pathways,
            [n for pw in data.pathways for n in pw.nodes],
            [e for pw in data.pathways for e in pw.edges],
        ]

        for coll in collections:
            for elem in coll:
                elem_obj = self._cache[elem.uuid]
                for s in elem.scenarios:
                    elem_obj.scenarios.add(self._cache[s.uuid])


    def _import_package_from_json(
            self,
    ) -> Package:
        try:
            parsed = PackageExportSchema.model_validate(self._raw_package)
        except ValidationError as e:
            logger.error(f"Error validating package data: {e}")
            raise PackageImportException("Deserialization failed")

        package_uuid = str(uuid.uuid4())
        if self.preserve_uuids:
            package_uuid = parsed.uuid

        package_license = License.objects.get(link=parsed.license.link) if parsed.license else None
        package_classification_level = Package.Classification[parsed.classification_level.upper()]

        print(parsed.classification_level, "->", package_classification_level)

        package = Package()
        package.uuid = package_uuid
        package.name = f"{parsed.name} - Imported at {datetime.now()}"
        package.description = parsed.description
        package.reviewed = False
        package.license = package_license
        package.classification_level = package_classification_level

        if package_classification_level == Package.Classification.SECRET:
            package.data_pool = Group.objects.get(uuid=parsed.data_pool.uuid)

        package.save()

        # Import Elements
        self._import_compounds(package, parsed.compounds)
        self._import_simple_rules(package, parsed.simple_rules)
        self._import_composite_rules(package, parsed.composite_rules)
        self._import_reactions(package, parsed.reactions)
        self._import_pathways(package, parsed.pathways)
        self._import_scenarios(package, parsed.scenarios)
        self._import_additional_information(package, parsed.additional_information)
        self._link_scenarios_after_import(parsed)

        return package

    @staticmethod
    def sign(data: Dict[str, Any], key: str) -> Dict[str, Any]:
        json_str = json.dumps(data, sort_keys=True, separators=(",", ":"))
        signature = hmac.new(key.encode(), json_str.encode(), hashlib.sha256).digest()
        data["_signature"] = base64.b64encode(signature).decode()
        return data

    @staticmethod
    def verify(data: Dict[str, Any], key: str) -> bool:
        copied_data = data.copy()
        sig = copied_data.pop("_signature")
        signature = base64.b64decode(sig, validate=True)
        json_str = json.dumps(copied_data, sort_keys=True, separators=(",", ":"))
        expected = hmac.new(key.encode(), json_str.encode(), hashlib.sha256).digest()
        return hmac.compare_digest(signature, expected)


class PathwayUtils:
    def __init__(self, pathway: "Pathway"):
        self.pathway = pathway

    @staticmethod
    def _get_products(smiles: str, rules: List["Rule"]):
        educt_rule_products: Dict[str, Dict[str, List[str]]] = defaultdict(
            lambda: defaultdict(list)
        )

        for r in rules:
            product_sets = r.apply(smiles)
            for product_set in product_sets:
                for product in product_set:
                    educt_rule_products[smiles][r.url].append(product)

        return educt_rule_products

    def find_missing_rules(self, rules: List["Rule"]):
        print(f"Processing {self.pathway.name}")
        # compute products for each node / rule combination in the pathway
        educt_rule_products = defaultdict(lambda: defaultdict(list))

        for node in self.pathway.nodes:
            educt_rule_products.update(**self._get_products(node.default_node_label.smiles, rules))

        # loop through edges and determine reactions that can't be constructed by
        # any of the rules or a combination of two rules in a chained fashion

        res: Dict[str, List["Rule"]] = dict()

        for edge in self.pathway.edges:
            found = False
            reaction = edge.edge_label

            educts = [cs for cs in reaction.educts.all()]
            products = [cs.smiles for cs in reaction.products.all()]
            rule_chain = []

            for educt in educts:
                educt = educt.smiles
                triggered_rules = list(educt_rule_products.get(educt, {}).keys())
                for triggered_rule in triggered_rules:
                    if rule_products := educt_rule_products[educt][triggered_rule]:
                        # check if this rule covers the reaction
                        if FormatConverter.smiles_covered_by(
                            products, rule_products, standardize=True, canonicalize_tautomers=True
                        ):
                            found = True
                        else:
                            # Check if another prediction step would cover the reaction
                            for product in rule_products:
                                prod_rule_products = self._get_products(product, rules)
                                prod_triggered_rules = list(
                                    prod_rule_products.get(product, {}).keys()
                                )
                                for prod_triggered_rule in prod_triggered_rules:
                                    if second_step_products := prod_rule_products[product][
                                        prod_triggered_rule
                                    ]:
                                        if FormatConverter.smiles_covered_by(
                                            products,
                                            second_step_products,
                                            standardize=True,
                                            canonicalize_tautomers=True,
                                        ):
                                            rule_chain.append(
                                                (
                                                    triggered_rule,
                                                    Rule.objects.get(url=triggered_rule).name,
                                                )
                                            )
                                            rule_chain.append(
                                                (
                                                    prod_triggered_rule,
                                                    Rule.objects.get(url=prod_triggered_rule).name,
                                                )
                                            )
                                            res[edge.url] = rule_chain

            if not found:
                res[edge.url] = rule_chain

        return res

    def engineer(self, setting: "Setting"):
        from epdb.logic import SPathway
        from utilities.chem import FormatConverter
        from utilities.ml import graph_from_pathway, get_shortest_path

        # get a fresh copy
        pw = Pathway.objects.get(id=self.pathway.pk)

        root_nodes = [n.default_node_label.smiles for n in pw.root_nodes]

        if len(root_nodes) != 1:
            logger.warning(f"Pathway {pw.name} has {len(root_nodes)} root nodes")
            # spw, mapping, intermediates
            return None, {}, []

        # Predict the Pathway in memory
        spw = SPathway(root_nodes[0], None, setting)

        level = 0
        while not spw.done:
            spw.predict_step(from_depth=level)
            level += 1

        # Generate SNode -> Node mapping
        node_mapping = {}

        for node in pw.nodes:
            for snode in spw.smiles_to_node.values():
                data_smiles = node.default_node_label.smiles
                pred_smiles = snode.smiles

                # "~" denotes any bond remove and use implicit single bond for comparison
                data_key = FormatConverter.InChIKey(data_smiles.replace("~", ""))
                pred_key = FormatConverter.InChIKey(pred_smiles.replace("~", ""))

                if data_key == pred_key:
                    node_mapping[snode] = node

        reverse_mapping = {v: k for k, v in node_mapping.items()}

        graph = graph_from_pathway(spw)

        intermediate_mapping = []

        # loop through each edge and each reactant <-> product pair
        # and compute the shortest path on the predicted pathway
        for e in pw.edges:
            for start in e.start_nodes.all():
                if start not in reverse_mapping:
                    continue

                start_snode = reverse_mapping[start]

                for end in e.end_nodes.all():
                    if end not in reverse_mapping:
                        continue

                    end_snode = reverse_mapping[end]

                    # If res is non-empty, we've found intermediates
                    intermediate_smiles = get_shortest_path(
                        graph,
                        FormatConverter.standardize(start_snode.smiles, remove_stereo=True),
                        FormatConverter.standardize(end_snode.smiles, remove_stereo=True),
                    )

                    if intermediate_smiles:
                        intermediates = []

                        prev = start_snode.smiles

                        for smi in intermediate_smiles + [end_snode.smiles]:
                            for e in spw.get_edge_for_educt_smiles(prev):
                                if smi in e.product_smiles():
                                    intermediates.append(e)

                            prev = smi

                        intermediate_mapping.append(
                            (start, end, start_snode, end_snode, intermediates)
                        )

        return spw, reverse_mapping, intermediate_mapping

    @staticmethod
    def spathway_to_pathway(
        package: "Package", spw: "SPathway", name: str = None, description: str = None
    ):
        snode_to_node_mapping = dict()

        root_nodes = spw.root_nodes

        pw = Pathway.create(
            package=package,
            smiles=root_nodes[0].smiles,
            name=name,
            description=description,
            predicted=True,
        )

        pw.setting = spw.prediction_setting
        pw.save()

        snode_to_node_mapping[root_nodes[0]] = pw.root_nodes[0]

        if len(root_nodes) > 1:
            for rn in root_nodes[1:]:
                n = Node.create(pw, rn.smiles, depth=0)
                snode_to_node_mapping[rn] = n

        for snode, node in snode_to_node_mapping.items():
            spw.snode_persist_lookup[snode] = node

        spw.persist = pw

        spw._sync_to_pathway()

        return pw