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

from django.conf import settings as s
from envipy_additional_information import EnviPyModel, UIConfig
from ninja import Schema
from pydantic import HttpUrl, ValidationError

from epdb.exceptions import PackageImportException
from epdb.models import (
    AdditionalInformation,
    Compound,
    CompoundStructure,
    Edge,
    License,
    Node,
    ParallelRule,
    Pathway,
    Reaction,
    Rule,
    Scenario,
    Setting,
    SimpleAmbitRule,
)
from utilities.chem import FormatConverter

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

if TYPE_CHECKING:
    from epdb.logic import SPathway


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 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]


############
# 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] = []

    @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
    compounds: List[CompoundExportSchema]
    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_compounds(obj):
        res = []
        if isinstance(obj, dict):
            for c in obj.get("compounds", []):
                res.append(CompoundExportSchema.model_validate(c))
            return res
        return obj.compounds.all()

    @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 Exporter(ABC):
    def __init__(self, package: Package):
        self._raw_package = package

    def do_export(self):
        return self._export()

    @abstractmethod
    def _export(self):
        pass


class PackageExporter(Exporter):
    def _export(self) -> 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(self._raw_package)

        return data.model_dump(mode="json")


class PathwayExporter(Exporter):
    def __init__(self, package: Package, add_infs_to_export: List[str] = []):
        super().__init__(package)
        self._add_infs_to_export = add_infs_to_export

    def _flatten_additional_information(self, ai: AdditionalInformation) -> dict[str, Any]:
        model_cls: Type[EnviPyModel] = type(ai.get())

        def _flatten(d: dict, parent_key: str = "") -> dict[str, Any]:
            items: dict[str, Any] = {}
            for key, value in d.items():
                new_key = f"{parent_key}__{key.lower()}" if parent_key else key.lower()
                if isinstance(value, dict):
                    items.update(_flatten(value, new_key))
                else:
                    items[new_key] = value
            return items

        flat = _flatten(ai.data, ai.type)

        ui_class = getattr(model_cls, "UI", None)
        if ui_class is None:
            return flat

        for f in model_cls.model_fields:
            ui_info = getattr(ui_class, f, None)

            if not isinstance(ui_info, UIConfig) or ui_info.unit is None:
                continue

            flat[f"{model_cls.__name__}__{f}__unit"] = ui_info.unit

        return flat

    def _export(self):
        from io import StringIO
        from csv import DictWriter

        rows = []

        for pw in self._raw_package.pathways.all():
            for n in pw.nodes:
                for scen in pw.scenarios.all():
                    row = {
                        "pathway_name": pw.name,
                        "pathway_id": str(pw.url),
                        "node_depth": n.depth,
                        "compound_id": str(n.default_node_label.compound.url),
                        "pubchem_ID": n.default_node_label.pubchem_compound_id,
                        "compound_name": n.default_node_label.compound.name,
                        "compound_smiles": n.default_node_label.smiles,
                        "scenario_id": str(scen.url),
                        "scenario_name": scen.name,
                        "scenario_type": scen.scenario_type,
                        "scenario_description": scen.description,
                    }

                    if self._add_infs_to_export:
                        ai_qs = AdditionalInformation.objects.filter(
                            scenario=scen, type__in=self._add_infs_to_export
                        )
                    else:
                        ai_qs = AdditionalInformation.objects.filter(scenario=scen)

                    for ai in ai_qs:
                        if ai.type == "ProposedIntermediate" and ai.content_object == n:
                            row.update({"proposed_intermediate": True})
                        elif ai.type == "SpikeCompound":
                            spike = {"SpikeCompound__url": ai.get().url}

                            try:
                                struc = CompoundStructure.objects.get(
                                    compound__package=self._raw_package, url=ai.get().url
                                )
                                spike["SpikeCompound__smiles"] = struc.smiles
                            except Exception:
                                spike["SpikeCompound__smiles"] = None

                            row.update(**spike)
                        else:
                            row.update(self._flatten_additional_information(ai))
                    rows.append(row)

        # Get all header fields
        all_header_fields = set()
        for row in rows:
            all_header_fields.update(row.keys())

        # Per request the CSV should start with these fields
        header = [
            "pathway_name",
            "pathway_id",
            "node_depth",
            "compound_id",
            "pubchem_ID",
            "compound_name",
            "compound_smiles",
            "scenario_id",
            "scenario_name",
            "scenario_type",
            "scenario_description",
        ]

        # User remaining fields and place them after the predefined values in a sorted manner
        remainder = sorted(list(all_header_fields.difference(set(header))))
        header.extend(remainder)

        buffer = StringIO()
        writer = DictWriter(buffer, fieldnames=header, delimiter="\t")
        writer.writeheader()
        writer.writerows(rows)
        buffer.seek(0)
        return buffer.getvalue()


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],
    ):
        for c in compounds:
            new_c = Compound()
            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:
                new_cs = CompoundStructure()
                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

                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 scen in scenarios:
            new_s = Scenario()
            new_s.uuid = str(uuid.uuid4()) if not self.preserve_uuids else scen.uuid
            new_s.package = package
            new_s.name = scen.name
            new_s.description = scen.description
            new_s.scenario_date = scen.scenario_date
            new_s.scenario_type = scen.scenario_type
            new_s.save()

            self._cache[scen.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 scen in elem.scenarios:
                    elem_obj.scenarios.add(self._cache[scen.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 = 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.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
