module3_react.py

import time
from typing import List, Tuple
from rmgpy.species import Species
from rmgpy.molecule import Molecule
from rmgpy.rmg.react import react

# Chunks slower than this threshold (seconds) get their species SMILES logged
# so pathological inputs can be cross-referenced after the run.
SLOW_CHUNK_THRESHOLD_S = 60.0

def smiles_to_species(smiles: str) -> Species:
    """Convert SMILES to RMG Species object with robust resonance generation."""
    try:
        mol = Molecule().from_smiles(smiles)
        try:
            mol.generate_resonance_structures()
        except Exception as e:
            print(f"    Warning: Resonance generation failed for {smiles}: {e}")
        return Species(molecule=[mol], label=smiles)
    except Exception as e:
        print(f"    Error converting SMILES to Species {smiles}: {e}")
        return None

def enumerate_reactions(species_list: List[Species], procnum: int = 1, chunk_size: int = 120) -> List:
    """
    Enumerate unimolecular and bimolecular reactions from species.

    React tuples are processed in chunks so that a single bad species only kills
    its chunk, not the whole batch. If a chunk raises, that chunk falls back to
    per-tuple singleton processing; other chunks keep running in parallel.
    """
    species_list = [s for s in species_list if s is not None]
    if not species_list:
        return []

    all_reactions = []
    react_tuples = []

    # 1. Unimolecular reactions
    for spc in species_list:
        react_tuples.append(((spc,), None))

    # 2. Bimolecular reactions including A + A
    for i, spc_a in enumerate(species_list):
        for j, spc_b in enumerate(species_list):
            if i > j:
                continue
            react_tuples.append(((spc_a, spc_b), None))

    if not react_tuples:
        return []

    num_chunks = (len(react_tuples) + chunk_size - 1) // chunk_size
    print(f"    Reacting {len(react_tuples)} species combinations in {num_chunks} chunk(s) "
          f"of up to {chunk_size} (procnum={procnum})...")

    chunk_timings: List[float] = []
    for chunk_idx in range(num_chunks):
        start = chunk_idx * chunk_size
        chunk = react_tuples[start:start + chunk_size]
        t0 = time.perf_counter()
        try:
            chunk_results = react(chunk, procnum=procnum)
            for sublist in chunk_results:
                if sublist:
                    all_reactions.extend(sublist)
        except Exception as e:
            print(f"    Chunk {chunk_idx + 1}/{num_chunks} failed ({type(e).__name__}: {e}); "
                  f"falling back to singletons for this chunk only.")
            for tup in chunk:
                try:
                    res = react([tup], procnum=1)
                    if res and res[0]:
                        all_reactions.extend(res[0])
                except Exception:
                    pass
        t_chunk = time.perf_counter() - t0
        chunk_timings.append(t_chunk)
        print(f"    Chunk {chunk_idx + 1}/{num_chunks}: {len(chunk)} tuples in {t_chunk:.1f}s")
        if t_chunk > SLOW_CHUNK_THRESHOLD_S:
            slow_smiles = []
            for tup in chunk:
                spcs = tup[0]
                try:
                    smis = [s.molecule[0].to_smiles() for s in spcs]
                    slow_smiles.append("+".join(smis))
                except Exception:
                    pass
            print(f"    [slow-chunk] chunk {chunk_idx + 1} took {t_chunk:.1f}s "
                  f"(threshold {SLOW_CHUNK_THRESHOLD_S:.0f}s); species: {slow_smiles}")

    if chunk_timings:
        print(f"    [chunk-timing] total={sum(chunk_timings):.1f}s "
              f"min={min(chunk_timings):.1f}s max={max(chunk_timings):.1f}s "
              f"mean={sum(chunk_timings)/len(chunk_timings):.1f}s")
            
    # Deduplicate based on Reaction SMILES
    unique_rxns = []
    seen_smiles = set()
    
    for rxn in all_reactions:
        try:
            reactants = sorted([s.molecule[0].to_smiles() for s in rxn.reactants])
            products = sorted([s.molecule[0].to_smiles() for s in rxn.products])
            rxn_smiles = ".".join(reactants) + ">>" + ".".join(products)
            if rxn_smiles not in seen_smiles:
                unique_rxns.append(rxn)
                seen_smiles.add(rxn_smiles)
        except:
            pass
            
    return unique_rxns