S03: Maximum Common Substructure in Reaction Informatics

import warnings
import rdkit
from rdkit import Chem, RDLogger
import networkx as nx
import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path
from synedu.Utils.vis import draw_molecular_graph  # for visualization

# Keep the lesson output readable: RDKit narrates unmapped atoms on stderr and
# Matplotlib warns about layouts we set deliberately. Neither is a result.
RDLogger.DisableLog("rdApp.*")
warnings.filterwarnings("ignore", category=UserWarning)

print("RDKit version:", rdkit.__version__)
print("NetworkX version:", nx.__version__)

DATA_DIR = Path("data")
CSV_PATH = DATA_DIR / "molecules.csv"
df = pd.read_csv(CSV_PATH)
df['mol'] = df['smiles'].apply(Chem.MolFromSmiles)
display(df)

from rdkit import Chem
from rdkit.Chem import rdFMCS, Draw


def rdkit_mcs(
    mols,
    timeout: int = 10,
    ringMatchesRingOnly: bool = True,
    maximizeBonds: bool = False,
    atomCompare=rdFMCS.AtomCompare.CompareElements,
    bondCompare=rdFMCS.BondCompare.CompareOrder,
):
    """
    Return the MCS SMARTS string and the rdFMCS result object.
    `mols` can be a list of RDKit Mol objects or SMILES strings.
    """
    mol_objs = [Chem.MolFromSmiles(s) if isinstance(s, str) else s for s in mols]
    if not mol_objs or any(mol is None for mol in mol_objs):
        raise ValueError("All MCS inputs must be valid molecules")
    res = rdFMCS.FindMCS(
        mol_objs,
        maximizeBonds=maximizeBonds,
        atomCompare=atomCompare,
        bondCompare=bondCompare,
        ringMatchesRingOnly=ringMatchesRingOnly,
        timeout=timeout,
    )
    return res.smartsString, res


def show_mcs_and_highlight(
    mols, smarts=None, show_all_matches=False, subimg_size=(300, 300)
):
    """
    Compute (or accept) MCS SMARTS, enumerate matches in each molecule,
    print a first-match atom-index mapping, and draw the molecules with
    highlighted MCS atoms.
    """
    # Accept SMILES or Mol
    mol_objs = [Chem.MolFromSmiles(s) if isinstance(s, str) else s for s in mols]
    if smarts is None:
        smarts, _ = rdkit_mcs(mol_objs)
    mcs_query = Chem.MolFromSmarts(smarts)

    # find matches
    matches = [tuple(m.GetSubstructMatches(mcs_query)) for m in mol_objs]
    match_counts = [len(ms) for ms in matches]

    print("MCS SMARTS:", smarts)
    for i, (m, cnt) in enumerate(zip(mol_objs, match_counts)):
        print(f"Mol {i+1}: {m.GetNumAtoms()} atoms, MCS matches: {cnt}")

    # If no matches found (shouldn't happen), return
    if any(cnt == 0 for cnt in match_counts):
        print("Warning: no MCS embedding found in one or more molecules.")
        return

    # Show mapping for the *first* embedding in each molecule
    first_matches = [ms[0] for ms in matches]
    mapping = (
        list(zip(first_matches[0], first_matches[1]))
        if len(first_matches) >= 2
        else None
    )
    if mapping:
        print("\nAtom mapping (mol1_atom -> mol2_atom) for first embeddings:")
        print(mapping)

    # Prepare highlight lists: choose which embedding to highlight for each molecule
    if show_all_matches:
        # create list-of-lists where we highlight all embeddings (flatten duplicates visually)
        hl_lists = [list(set(sum([list(t) for t in ms], []))) for ms in matches]
    else:
        hl_lists = [list(ms[0]) for ms in matches]

    # Legend labels
    legends = [
        mol.GetProp("_Name") if mol.HasProp("_Name") else f"Mol {i+1}"
        for i, mol in enumerate(mol_objs)
    ]

    # Draw side-by-side with highlighted atoms
    img = Draw.MolsToGridImage(
        mol_objs,
        molsPerRow=len(mol_objs),
        subImgSize=subimg_size,
        legends=legends,
        highlightAtomLists=hl_lists,
    )
    display(img)  # in a Jupyter notebook; returns PIL Image otherwise


m1 = Chem.MolFromSmiles("Oc1ccccc1")
m1.SetProp("_Name", "phenol")
m2 = Chem.MolFromSmiles("Cc1ccccc1")
m2.SetProp("_Name", "toluene")

smarts, res = rdkit_mcs([m1, m2])
show_mcs_and_highlight(
    [m1, m2], smarts=smarts, show_all_matches=False, subimg_size=(400, 400)
)

from __future__ import annotations

from itertools import combinations
from typing import Dict, Set, Tuple

import networkx as nx
from networkx.algorithms import isomorphism as iso
from rdkit import Chem


def nx_mcs(
    G: nx.Graph,
    H: nx.Graph,
    *,
    node_match,
    edge_match,
) -> tuple[nx.Graph, Dict[int, int]]:
    """
    Return one maximum common induced subgraph SG of G, plus a node mapping SG->H.
    Exponential; for teaching only.
    """
    nodes = list(G.nodes())
    for k in range(len(nodes), 0, -1):
        for subset in combinations(nodes, k):
            SG = G.subgraph(subset).copy()

            GM = iso.GraphMatcher(H, SG, node_match=node_match, edge_match=edge_match)
            if GM.subgraph_is_isomorphic():
                # GraphMatcher mapping is H -> SG; invert to get SG -> H
                h_to_sg = next(GM.subgraph_isomorphisms_iter())
                sg_to_h = {sg: h for h, sg in h_to_sg.items()}
                return SG, sg_to_h

    return nx.Graph(), {}


def _map_edges(
    SG: nx.Graph, sg_to_h: Dict[int, int], H: nx.Graph
) -> Set[Tuple[int, int]]:
    """Map SG edges into H using sg_to_h; returns undirected edge tuples."""
    out: Set[Tuple[int, int]] = set()
    for u, v in SG.edges():
        hu, hv = sg_to_h[u], sg_to_h[v]
        if H.has_edge(hu, hv):
            out.add((hu, hv))
    return out

from synedu.Utils import mol_to_graph, node_match, edge_match, draw_molecular_graph
import matplotlib.pyplot as plt


m1 = Chem.MolFromSmiles("Oc1ccccc1")  # phenol
m2 = Chem.MolFromSmiles("Cc1ccccc1")  # toluene

# Convert to labeled NetworkX graphs (nodes must be atom indices)
m1_g = mol_to_graph(m1)
m2_g = mol_to_graph(m2)

# Compute NX MCS + one embedding (SG -> m2_g mapping)
SG, sg_to_h = nx_mcs(m1_g, m2_g, node_match=node_match, edge_match=edge_match)

print("MCS nodes:", SG.number_of_nodes(), "MCS edges:", SG.number_of_edges())
print("One SG->H mapping (first):", sg_to_h)

# Highlight sets
hl_m1_nodes = set(SG.nodes())
hl_m1_edges = set(SG.edges())

hl_m2_nodes = {sg_to_h[n] for n in SG.nodes()}
hl_m2_edges = _map_edges(SG, sg_to_h, m2_g)

fig, axes = plt.subplots(1, 2, figsize=(10, 4))

draw_molecular_graph(
    m1_g,
    ax=axes[0],
    include_mol=False,
    highlight_nodes=hl_m1_nodes,
    highlight_edges=hl_m1_edges,
)
axes[0].set_title("Phenol", fontsize=20, weight='bold')

draw_molecular_graph(
    m2_g,
    ax=axes[1],
    include_mol=False,
    highlight_nodes=hl_m2_nodes,
    highlight_edges=hl_m2_edges,
)
axes[1].set_title("Toluene", fontsize=20, weight='bold')

plt.tight_layout()
plt.show()

from synedu.Utils import draw_rxn_graph

fig = draw_rxn_graph(
    'CCO>>C=C',
    title="Elimination",
    highlight_changes=True,
    label_mode="hetero",
    show_indices=False,
    aromatic_style="circle",
    show_legend=False,
)

from __future__ import annotations

from collections import Counter
from typing import Tuple

from rdkit import Chem


# Reaction decomposition


def parse_reaction_smiles(rsmi: str) -> Tuple[list[Chem.Mol], list[Chem.Mol]]:
    """
    Parse a reaction SMILES into reactant and product molecule lists.

    Parameters
    ----------
    rsmi : str
        Reaction SMILES, e.g. "CCO.O>>CC=O"

    Returns
    -------
    reactants : list[Chem.Mol]
    products  : list[Chem.Mol]
    """
    try:
        lhs, rhs = rsmi.split(">>")
    except ValueError:
        raise ValueError("Reaction SMILES must contain '>>'")

    reactants = [Chem.MolFromSmiles(s) for s in lhs.split(".") if s]
    products = [Chem.MolFromSmiles(s) for s in rhs.split(".") if s]

    if any(m is None for m in reactants + products):
        raise ValueError("Invalid molecule in reaction SMILES")

    return reactants, products


# Formula representation


def molecular_formula(mol: Chem.Mol) -> Counter[str]:
    """
    Compute element counts and net formal charge for one molecule.

    Returns
    -------
    Counter[str]
        Mapping element symbol -> atom count
    """
    mol = Chem.AddHs(mol)

    formula = Counter()
    for atom in mol.GetAtoms():
        formula[atom.GetSymbol()] += 1
        formula["__charge__"] += int(atom.GetFormalCharge())
    if formula["__charge__"] == 0:
        del formula["__charge__"]
    return formula


def total_formula(mols: list[Chem.Mol]) -> Counter[str]:
    """
    Compute total formula Φ = sum φ(M) over a list of molecules.
    """
    total = Counter()
    for m in mols:
        for key, value in molecular_formula(m).items():
            total[key] += value
    if total["__charge__"] == 0:
        del total["__charge__"]
    return total


# Imbalance detection

from collections import Counter
from typing import Literal, Tuple


def signed_difference(a: Counter[str], b: Counter[str]) -> Counter[str]:
    """
    Compute signed element-wise difference a - b
    without dropping negative values.
    """
    delta = Counter()
    for elem in set(a) | set(b):
        v = a.get(elem, 0) - b.get(elem, 0)
        if v != 0:
            delta[elem] = v
    return delta


def reaction_imbalance(
    rsmi: str,
) -> Tuple[
    Counter[str],
    Literal[
        "balanced",
        "missing_in_products",
        "missing_in_reactants",
        "missing_in_both_sides",
    ],
]:
    """
    Compute the reaction imbalance vector Δ = Φ_R − Φ_P
    and classify where material is missing.

    Returns
    -------
    delta : Counter[str]
        Element-wise imbalance (positive = excess on reactant side)
    status : str
        One of:
        - 'balanced'
        - 'missing_in_products'
        - 'missing_in_reactants'
        - 'missing_in_both_sides'
    """
    reactants, products = parse_reaction_smiles(rsmi)

    phi_R = total_formula(reactants)
    phi_P = total_formula(products)
    delta = signed_difference(phi_R, phi_P)

    if not delta:
        return delta, "balanced"

    has_pos = any(v > 0 for v in delta.values())
    has_neg = any(v < 0 for v in delta.values())

    if has_pos and not has_neg:
        status = "missing_in_products"
    elif has_neg and not has_pos:
        status = "missing_in_reactants"
    else:
        status = "missing_in_both_sides"

    return delta, status

rsmi = 'CCO>>C=C'
delta, status = reaction_imbalance(rsmi)
print(delta, status)

from collections import Counter
from typing import Dict, Optional

AUX_LIBRARY: Dict[str, Counter[str]] = {
    "O": Counter({"H": 2, "O": 1}),  # H2O
    "[H+]": Counter({"H": 1}),
    "[OH-]": Counter({"H": 1, "O": 1}),
}

def impute_rsmi_rule(
    rsmi: str,
    delta: Counter[str],
    status: str,
    library: Dict[str, Counter[str]] = AUX_LIBRARY,
) -> Optional[str]:
    if status == "balanced":
        return rsmi

    if status == "missing_in_both_sides":
        return None

    lhs, rhs = rsmi.split(">>")
    # decide direction
    if status == "missing_in_products":
        target_delta = delta
        side = "products"
    elif status == "missing_in_reactants":
        target_delta = -delta
        side = "reactants"
    else:
        return None

    # try library
    for aux_smiles, formula in library.items():
        c = can_impute_with_aux(target_delta, formula)
        if c is None:
            continue

        aux_block = ".".join([aux_smiles] * c)

        if side == "products":
            rhs = rhs + "." + aux_block if rhs else aux_block
        else:
            lhs = lhs + "." + aux_block if lhs else aux_block

        return f"{lhs}>>{rhs}"

    return None


def can_impute_with_aux(
    delta: Counter[str],
    aux_formula: Counter[str],
) -> Optional[int]:
    """
    Check whether Δ can be written as c · φ(A).
    """
    if not delta:
        return 0

    ratios = []

    for elem, count in aux_formula.items():
        if elem not in delta:
            return None
        if delta[elem] % count != 0:
            return None
        ratios.append(delta[elem] // count)

    if not ratios or len(set(ratios)) != 1:
        return None

    c = ratios[0]
    if c <= 0:
        return None

    remainder = delta - Counter({k: c * v for k, v in aux_formula.items()})
    if remainder:
        return None

    return c

delta, status = reaction_imbalance(rsmi)
rsmi_balanced = impute_rsmi_rule(rsmi, delta, status, library=AUX_LIBRARY)

print(rsmi_balanced)
if rsmi_balanced:
    fig = draw_rxn_graph(
        rsmi_balanced,
        title="E1 example",
        highlight_changes=True,
        label_mode="hetero",
        show_indices=False,
        aromatic_style="circle",
        show_legend=False,
    )

rsmi = (
    "COC(=O)[C@H](CCCCNC(=O)OCc1ccccc1)NC(=O)Nc1cc(OC)cc(C(C)(C)C)c1O"
    ">>"
    "COC(=O)[C@H](CCCCN)NC(=O)Nc1cc(OC)cc(C(C)(C)C)c1O"
)


print(reaction_imbalance(rsmi))

fig = draw_rxn_graph(
    rsmi,
)

from synedu.Utils.reaction import (
    rsmi_to_graph,
)  # ► stored in synedu.Utils.reaction.rsmi_to_graph
from synkit.Graph.Matcher.mcs_matcher import MCSMatcher  # fast MCS for molecular graph

r, p = rsmi_to_graph(rsmi, drop_non_aam=False, use_index_as_atom_map=False)

mcs = MCSMatcher(
    node_attrs=['element', 'formal_charge', 'aromatic'],
    edge_attrs=['order'],
    prune_automorphisms=True,  # reduce the redundant maps
)
mcs.find_common_subgraph(r, p, mcs=True)
maps = mcs.get_mappings()
maps

exist_indice = [value for _, value in maps[0].items()]
missing_indice = [value for value in r.nodes if value not in exist_indice]

import networkx as nx
from typing import Iterable


def extract_induced_subgraph(
    G: nx.Graph, nodes: Iterable[int], *, copy: bool = True
) -> nx.Graph:
    nodes = set(nodes)

    # Safety check
    missing = nodes - set(G.nodes)
    if missing:
        raise KeyError(f"Nodes not found in graph: {sorted(missing)}")

    H = G.subgraph(nodes)
    return H.copy() if copy else H

sub = extract_induced_subgraph(r, missing_indice)
draw_molecular_graph(sub)

import networkx as nx
import math
from typing import List


_ALLOWED_VALENCE = {
    "C": {4},
    "N": {3, 4},  # neutral amine (3), ammonium (4)
    "O": {2},
    "S": {2, 4, 6},
    "P": {3, 5},
}


def add_wildcard(
    G: nx.Graph,
    indices: List[int],
) -> nx.Graph:
    """
    Return a copy of the graph with wildcard (*) atoms added
    to balance valence for selected atoms.

    For each atom index in `indices`, compute:
        total_valence = hcount + sum(bond orders)

    If total_valence is smaller than the maximum allowed valence for that
    element, add a wildcard atom (*) with the largest supported integral bond
    order that does not exceed the deficit. This avoids fractional wildcard
    bonds when an aromatic fragment is represented with bond order 1.5.

    Parameters
    ----------
    G : nx.Graph
        Molecular graph (original graph is NOT modified).
    indices : List[int]
        Atom indices to check and repair.

    Returns
    -------
    nx.Graph
        New graph with wildcard atoms added.
    """
    H = G.copy()

    next_idx = max(H.nodes, default=0) + 1

    for atom in indices:
        if atom not in H:
            continue

        data = H.nodes[atom]
        element = data.get("element")
        hcount = data.get("hcount", 0)

        if element not in _ALLOWED_VALENCE:
            continue

        bond_valence = sum(
            H.edges[atom, nbr].get("order", 1.0) for nbr in H.neighbors(atom)
        )

        total_valence = hcount + bond_valence
        allowed = max(_ALLOWED_VALENCE[element])

        deficit = allowed - total_valence
        if deficit <= 0:
            continue
        bond_order = min(3, math.floor(deficit + 1e-9))
        if bond_order < 1:
            continue

        wc = next_idx
        next_idx += 1

        H.add_node(
            wc,
            element="*",
            aromatic=False,
            hcount=0,
            formal_charge=0,
            wildcard=True,
        )
        H.add_edge(atom, wc, order=float(bond_order), aromatic=False)

    return H

from synedu.Utils.graph import (
    print_graph_attributes,
)  # ► stored in synedu.Utils.graph.print_graph_attributes

wc_sub = add_wildcard(sub, indices=list(sub.nodes()))

print_graph_attributes(wc_sub)

# Use a hydroxy group as the fragment to merge.
from synedu.Utils.conversion import (
    smiles_to_graph,
)  # ► stored in synedu.Utils.conversion.smiles_to_graph

smiles = '*[O][H]'
add_sub = smiles_to_graph(smiles)
print_graph_attributes(add_sub)

import networkx as nx


def find_wildcards(G: nx.Graph) -> list[int]:
    """Return indices of wildcard (*) nodes."""
    return [n for n, d in G.nodes(data=True) if d.get("element") == "*"]


def _get_unique_wildcard(G: nx.Graph, label: str) -> int:
    wcs = find_wildcards(G)
    if len(wcs) != 1:
        raise ValueError(f"{label} must contain exactly one wildcard")
    return wcs[0]


def merge_graphs_on_wildcard(
    G1: nx.Graph,
    G2: nx.Graph,
    *,
    wc1: int | None = None,
    wc2: int | None = None,
) -> nx.Graph:
    """
    Merge two molecular graphs by gluing them at wildcard (*) nodes.

    The wildcard nodes are removed, and their unique neighbors are connected.
    """
    # --- identify wildcards ---
    wc1 = wc1 if wc1 is not None else _get_unique_wildcard(G1, "G1")
    wc2 = wc2 if wc2 is not None else _get_unique_wildcard(G2, "G2")

    # --- check 1–1 attachment ---
    nbrs1 = list(G1.neighbors(wc1))
    nbrs2 = list(G2.neighbors(wc2))
    if len(nbrs1) != 1 or len(nbrs2) != 1:
        raise ValueError("Only 1–1 wildcard gluing is supported")

    u1 = nbrs1[0]
    u2 = nbrs2[0]

    # --- initialize merged graph ---
    G = nx.Graph()

    # copy G1 without wildcard
    for n, d in G1.nodes(data=True):
        if n != wc1:
            G.add_node(n, **d)
    for u, v, d in G1.edges(data=True):
        if wc1 not in (u, v):
            G.add_edge(u, v, **d)

    # copy G2 without wildcard (with offset)
    offset = max(G.nodes, default=-1) + 1
    node_map = {}

    for n, d in G2.nodes(data=True):
        if n == wc2:
            continue
        node_map[n] = n + offset
        G.add_node(node_map[n], **d)

    for u, v, d in G2.edges(data=True):
        if wc2 not in (u, v):
            G.add_edge(node_map[u], node_map[v], **d)

    # --- glue the attachment points ---
    bond_order = min(
        G1.edges[wc1, u1]["order"],
        G2.edges[wc2, u2]["order"],
    )
    G.add_edge(u1, node_map[u2], order=bond_order)

    return G

missing_graph = merge_graphs_on_wildcard(wc_sub, add_sub)

fig, axes = plt.subplots(1, 3, figsize=(15, 4))
draw_molecular_graph(
    sub,
    ax=axes[0],
    include_mol=False,
    highlight_nodes=None,
    highlight_edges=None,
)
axes[0].set_title("Raw", fontsize=20, weight='bold')

draw_molecular_graph(
    new_sub,
    ax=axes[1],
    include_mol=False,
    highlight_nodes=None,
    highlight_edges=None,
)
axes[1].set_title("Wildcard", fontsize=20, weight='bold')

draw_molecular_graph(
    missing_graph,
    ax=axes[2],
    include_mol=False,
    highlight_nodes=None,
    highlight_edges=None,
)
axes[2].set_title("Complete", fontsize=20, weight='bold')

plt.tight_layout()
plt.show()

from synedu.Utils.conversion import (
    graph_to_smi,
)  # ► stored in synedu.Utils.conversion.graph_to_smi

missing_smi = graph_to_smi(missing_graph)
print(missing_smi)

# pip install synrbl
from synrbl import Balancer

smiles = (
    "COC(=O)[C@H](CCCCNC(=O)OCc1ccccc1)NC(=O)Nc1cc(OC)cc(C(C)(C)C)c1O>>"
    + "COC(=O)[C@H](CCCCN)NC(=O)Nc1cc(OC)cc(C(C)(C)C)c1O"
)
synrbl = Balancer()

results = synrbl.rebalance(smiles, output_dict=True)
results