import re
from collections import defaultdict

INPUT_FILENAME = "example"


def parse_atoms(molecule):
    """Return a list of all Atoms in the molecule string"""
    return re.findall(r"[A-Z][a-z]*", molecule)


ATOMS = set()
REPLACEMENTS = list()

with open(INPUT_FILENAME) as f:
    line = f.readline()
    while line != "\n":
        atom, decays_to = line.strip().split(" => ")
        REPLACEMENTS.append([atom, decays_to])
        ATOMS.update(parse_atoms(decays_to))
        ATOMS.add(atom)
        line = f.readline()

    MOLECULE = f.readline().strip()

ATOMS.update(parse_atoms(MOLECULE))


print("ATOMS", ATOMS)
print()
print("REPLACEMENTS", REPLACEMENTS)
print()
print("MOLECULE", MOLECULE)
print()


def shrink(molecule, from_molecule, to_atom, idx):
    prefix = molecule[:idx]
    suffix = molecule[idx + len(from_molecule) :]
    return prefix + to_atom + suffix


def step(molecule):
    next_molecules = []
    for (atom, decays_to) in REPLACEMENTS:
        possible = [m.start() for m in re.finditer(decays_to, molecule)]
        if not possible:
            continue

        print(decays_to, atom)

        if atom == "e" and len(decays_to) != len(molecule):
            continue

        for idx in possible:
            next_molecule = shrink(molecule, decays_to, atom, idx)
            next_molecules.append(next_molecules)
            print("    ", idx, next_molecule)

    return next_molecules


step(MOLECULE)