map.py

import os
import json
import re
from collections import defaultdict, Counter
from rdkit import Chem
from rdkit.Chem import Descriptors, rdFMCS
from rxnmapper import RXNMapper

def has_metal(smi):
    mol = Chem.MolFromSmiles(smi)
    if not mol:
        return False

    # Non-metals: H, C, N, O, F, Si, P, S, Cl, As, Se, Br, I
    non_metals = {1, 6, 7, 8, 9, 14, 15, 16, 17, 33, 34, 35, 53}

    # Check all atoms, ignore * (atomic number 0)
    for atom in mol.GetAtoms():
        atom_num = atom.GetAtomicNum()
        if atom_num != 0 and atom_num not in non_metals:
            return True
    
    return False

def find_radical(smi):
    mol = Chem.MolFromSmiles(smi)
    if not mol:
        return False
    num_radicals = Descriptors.NumRadicalElectrons(mol)
    return num_radicals > 0

def get_total_charge(smiles):
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        return "Invalid SMILES"
    return Chem.GetFormalCharge(mol)

def generate_atom_mapping(reactant, product, rxn_mapper):
    reaction_smiles = f"{reactant}>>{product}"
    try:
        results = rxn_mapper.get_attention_guided_atom_maps([reaction_smiles])
        mapped_rxn = results[0]['mapped_rxn']
        confidence = results[0]['confidence']
        mapped_reactant, mapped_product = mapped_rxn.split('>>')
        return mapped_reactant, mapped_product
    except Exception as e:
        print(f"❌ Mapping failed: {e}")
        return None, None

def check_mapping_consistency(reactant_smi: str, product_smi: str):
    """
    Check if the Atom Mapping between reactants and products is chemically consistent.
    """
    mol_r = Chem.MolFromSmiles(reactant_smi)
    mol_p = Chem.MolFromSmiles(product_smi)
    
    if not mol_r or not mol_p:
        return False

    def get_map_info(mol):
        mapping = {}
        for atom in mol.GetAtoms():
            map_id = atom.GetAtomMapNum()
            if map_id != 0:
                mapping[map_id] = atom.GetSymbol()
        return mapping

    r_map = get_map_info(mol_r)
    p_map = get_map_info(mol_p)

    is_valid = True
    r_ids = set(r_map.keys())
    p_ids = set(p_map.keys())
    
    missing_in_p = r_ids - p_ids
    if missing_in_p:
        is_valid = False

    missing_in_r = p_ids - r_ids
    if missing_in_r:
        is_valid = False

    common_ids = r_ids.intersection(p_ids)
    for mid in common_ids:
        if r_map[mid] != p_map[mid]:
            is_valid = False

    return is_valid

def common_items(list_a, list_b):
    count_a = Counter(list_a)
    count_b = Counter(list_b)

    common_counts = count_a & count_b
    remain_a_counts = count_a - count_b
    remain_b_counts = count_b - count_a

    remain_a = list(remain_a_counts.elements())
    remain_b = list(remain_b_counts.elements())
    removed_items = list(common_counts.elements())

    return remain_a, remain_b, removed_items

def extract_molecule(smi):
    parts = smi.split(">>")
    reactant = parts[0]
    product = parts[1]
    reactants = reactant.split('.')
    products = product.split('.')
    r, p, secondary = common_items(reactants, products)
    return r, p, secondary

def map_H(reactant, product):
    """
    Input: rxnmapper-mapped reaction SMILES (implicit H)
    Output: number of proton transfer, reaction SMILES with explicit H and mapped protons
    """
    reactants = [Chem.AddHs(Chem.MolFromSmiles(smi)) for smi in reactant.split('.')]
    products = [Chem.AddHs(Chem.MolFromSmiles(smi)) for smi in product.split('.')]

    # --- Index heavy atoms by atom map ---
    def index_atoms(mols):
        idx = defaultdict(list)
        for mol in mols:
            for atom in mol.GetAtoms():
                if atom.GetAtomicNum() != 1:  # heavy atom
                    amap = atom.GetAtomMapNum()
                    idx[amap].append(atom)
        return idx
    
    r_atoms = index_atoms(reactants)
    p_atoms = index_atoms(products)
    current_max_num = max(r_atoms.keys()) if r_atoms else 0

    # --- Collect hydrogens per heavy atom ---
    def get_attached_H(atom):
        return [
            nbr for nbr in atom.GetNeighbors()
            if nbr.GetAtomicNum() == 1
        ]

    lost_H = []
    gained_H = []

    for amap in r_atoms:
        if amap not in p_atoms:
            continue

        r_atom = r_atoms[amap][0]
        p_atom = p_atoms[amap][0]

        r_H = get_attached_H(r_atom)
        p_H = get_attached_H(p_atom)

        delta = len(p_H) - len(r_H)

        if delta < 0:
            lost_H.extend(r_H[:abs(delta)])
            r_H = r_H[abs(delta):]
        elif delta > 0:
            gained_H.extend(p_H[:delta])
            p_H = p_H[delta:]

        assert len(p_H) == len(r_H)
        
        for h_r, h_p in zip(r_H, p_H):
            h_map = current_max_num + 1
            current_max_num = h_map
            h_r.SetAtomMapNum(h_map)
            h_p.SetAtomMapNum(h_map)

    print(f'Unpaired H: {len(lost_H)}')
    pt = len(lost_H)
    
    for h_r, h_p in zip(lost_H, gained_H):
        current_max_num += 1
        h_r.SetAtomMapNum(current_max_num)
        h_p.SetAtomMapNum(current_max_num)
        
    # --- Clear unmapped H (optional safety) ---
    for mol in reactants + products:
        for atom in mol.GetAtoms():
            if atom.GetAtomicNum() == 1 and atom.GetAtomMapNum() == 0:
                print("Warning: unmapped atoms")
                return None, None, None, None

    # --- Rebuild reaction ---
    new_reactants = '.'.join([Chem.MolToSmiles(mol, allHsExplicit=True, kekuleSmiles=True) for mol in reactants])
    new_products = '.'.join([Chem.MolToSmiles(mol, allHsExplicit=True, kekuleSmiles=True) for mol in products])

    return pt, new_reactants, new_products, current_max_num

def fill_in_common(common, current_max_num):
    new_smiles_list = []
    counter = current_max_num

    for smi in common:
        mol = Chem.MolFromSmiles(smi)
        if mol is None:
            continue

        mol = Chem.AddHs(mol)
        Chem.Kekulize(mol, clearAromaticFlags=True)

        for atom in mol.GetAtoms():
            counter += 1
            atom.SetAtomMapNum(counter)

        new_smi = Chem.MolToSmiles(mol, kekuleSmiles=True, allHsExplicit=True)
        new_smiles_list.append(new_smi)

    return new_smiles_list

def clean(smi):
    r, p = smi.split(">>")
    r_mol = Chem.MolFromSmiles(r)
    p_mol = Chem.MolFromSmiles(p)
    
    for atom in r_mol.GetAtoms():
        atom.SetAtomMapNum(0)
    for atom in p_mol.GetAtoms():
        atom.SetAtomMapNum(0)
    return Chem.MolToSmiles(r_mol, kekuleSmiles=True) + ">>" + Chem.MolToSmiles(p_mol, kekuleSmiles=True)

def main():
    # Read and write files relative to the directory of this script
    script_dir = os.path.dirname(os.path.abspath(__file__))
    
    reactions_path = os.path.join(script_dir, 'MCSA_reactions.json')
    mapped_data_output_path = os.path.join(script_dir, 'mapped_reactions.json')
    subdirs_output_dir = os.path.join(script_dir, 'mapped_reactions')
    
    print(f"Reading reaction data: {reactions_path}")
    if not os.path.exists(reactions_path):
        print(f"❌ Reaction data file not found: {reactions_path}")
        return

    with open(reactions_path, 'r', encoding='utf-8') as f:
        reactions = json.load(f)

    # Instantiate RXNMapper
    print("Initializing RXNMapper...")
    rxn_mapper = RXNMapper()

    MCSA_mapped_data = defaultdict(list)
    count = 0
    for k, v in reactions.items():
        if count >100:
            break
        count+=1
        status = v[0]
        smi = v[1]
        print(f'Starting processing {k}')

        if status == 'failed' or 'charge' in status:
            continue
        
        r, p, common = extract_molecule(smi)
        reactant = '.'.join(r)
        product = '.'.join(p)

        # Atom mapping
        r_mapped, p_mapped = generate_atom_mapping(reactant, product, rxn_mapper)

        if r_mapped is None:
            continue
        
        # Chemistry and metal detection
        if find_radical(r_mapped) or find_radical(p_mapped):
            chemistry = 'Radical'
        else:
            chemistry = 'Polar'

        if has_metal(r_mapped) or has_metal(p_mapped):
            metal = 'metal'
        else:
            metal = 'non-metal'

        # Consistency check and hydrogen atom mapping
        if check_mapping_consistency(r_mapped, p_mapped):
            pt, new_reactant, new_product, current_max_num = map_H(r_mapped, p_mapped)
            if pt is not None:
                map_result = f'{new_reactant}>>{new_product}'
                if common:
                    common_list = fill_in_common(common, current_max_num)
                    common_smi = '.'.join(common_list)
                    map_result = new_reactant + '.' + common_smi + '>>' + new_product + '.' + common_smi
                
                MCSA_mapped_data[k] = [chemistry, metal, pt, map_result]

    # Output complete mapping results
    print(f"Writing complete mapping results to: {mapped_data_output_path}")
    with open(mapped_data_output_path, 'w', encoding='utf-8') as f:
        json.dump(MCSA_mapped_data, f, sort_keys=True, indent=4, ensure_ascii=False)

    # Classify data
    polar_nonmetal = {}
    polar_metal = {}
    radical_nonmetal = {}
    radical_metal = {}

    for k, v in MCSA_mapped_data.items():
        chemistry = v[0]
        metal = v[1]
        pt = v[2]
        smi = v[3]

        if not pt:
            continue

        if chemistry == 'Polar' and metal == 'non-metal' and pt < 2:
            polar_nonmetal[k] = smi

        elif chemistry == 'Polar' and metal == 'metal' and pt < 2:
            polar_metal[k] = smi

        elif chemistry == 'Radical' and metal == 'non-metal' and pt < 2:
            radical_nonmetal[k] = smi

        elif chemistry == 'Radical' and metal == 'metal' and pt < 2:
            radical_metal[k] = smi

    os.makedirs(subdirs_output_dir, exist_ok=True)
    
    outputs = {
        'Polar_NM.json': polar_nonmetal,
        'Polar_M.json': polar_metal,
        'Radical_NM.json': radical_nonmetal,
        'Radical_M.json': radical_metal
    }

    for filename, data in outputs.items():
        out_path = os.path.join(subdirs_output_dir, filename)
        print(f"Writing classified file to: {out_path} (Count: {len(data)})")
        with open(out_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, sort_keys=True, indent=4, ensure_ascii=False)

    print("🎉 Processing completed!")

if __name__ == '__main__':
    main()