from collections import defaultdict
from typing import Any, Iterator

from ..base import BaseSolver


class Solver(BaseSolver):
    def solve(self, input: str) -> Iterator[Any]:
        replacements_s, molecule = input.split("\n\n")

        REPLACEMENTS: dict[str, list[str]] = defaultdict(list)
        for replacement_s in replacements_s.splitlines():
            p = replacement_s.split(" => ")
            REPLACEMENTS[p[0]].append(p[1])
        molecule = molecule.strip()

        generated = [
            molecule[:i] + replacement + molecule[i + len(symbol) :]
            for symbol, replacements in REPLACEMENTS.items()
            for replacement in replacements
            for i in range(len(molecule))
            if molecule[i:].startswith(symbol)
        ]

        yield len(set(generated))

        inversion: dict[str, str] = {
            replacement: symbol
            for symbol, replacements in REPLACEMENTS.items()
            for replacement in replacements
        }

        # there is actually only one way to create the molecule, and we can greedily replace
        # tokens with their replacements, e.g., if H => OH then we can replace OH by H directly
        # without thinking

        count = 0
        while molecule != "e":
            i = 0
            m2 = ""
            while i < len(molecule):
                found = False
                for replacement in inversion:
                    if molecule[i:].startswith(replacement):
                        m2 += inversion[replacement]
                        i += len(replacement)
                        count += 1
                        found = True
                        break

                if not found:
                    m2 += molecule[i]
                    i += 1

            molecule = m2

        yield count