#TAGS lexer, parser
# trick reverse order, to get count first
# TODO rewrite lexer to lex in reverse order
# see 394
#
# @lc app=leetcode id=726 lang=python3
#
# [726] Number of Atoms
#
# https://leetcode.com/problems/number-of-atoms/description/
#
# algorithms
# Hard (45.63%)
# Likes:    248
# Dislikes: 88
# Total Accepted:    12.7K
# Total Submissions: 27.9K
# Testcase Example:  '"H2O"'
#
# Given a chemical formula (given as a string), return the count of each atom.
#
# An atomic element always starts with an uppercase character, then zero or
# more lowercase letters, representing the name.
#
# 1 or more digits representing the count of that element may follow if the
# count is greater than 1.  If the count is 1, no digits will follow.  For
# example, H2O and H2O2 are possible, but H1O2 is impossible.
#
# Two formulas concatenated together produce another formula.  For example,
# H2O2He3Mg4 is also a formula.
#
# A formula placed in parentheses, and a count (optionally added) is also a
# formula.  For example, (H2O2) and (H2O2)3 are formulas.
#
# Given a formula, output the count of all elements as a string in the
# following form: the first name (in sorted order), followed by its count (if
# that count is more than 1), followed by the second name (in sorted order),
# followed by its count (if that count is more than 1), and so on.
#
# Example 1:
#
# Input:
# formula = "H2O"
# Output: "H2O"
# Explanation:
# The count of elements are {'H': 2, 'O': 1}.
#
#
#
# Example 2:
#
# Input:
# formula = "Mg(OH)2"
# Output: "H2MgO2"
# Explanation:
# The count of elements are {'H': 2, 'Mg': 1, 'O': 2}.
#
#
#
# Example 3:
#
# Input:
# formula = "K4(ON(SO3)2)2"
# Output: "K4N2O14S4"
# Explanation:
# The count of elements are {'K': 4, 'N': 2, 'O': 14, 'S': 4}.
#
#
#
# Note:
# All atom names consist of lowercase letters, except for the first character
# which is uppercase.
# The length of formula will be in the range [1, 1000].
# formula will only consist of letters, digits, and round parentheses, and is a
# valid formula as defined in the problem.
#
#

def lexer(formula):
    n = len(formula)
    i = 0
    while i < n:
        c = formula[i]
        if c.isdigit():
            res = 0
            while i < n and formula[i].isdigit():
                res = int(formula[i]) + 10  * res
                i += 1
            yield res
        elif c in {'(', ')'}:
            yield c
            i += 1
        else:
            res = [c]
            i += 1
            while i < n and 'a' <= formula[i] <= 'z':
                res.append(formula[i])
                i += 1
            yield ''.join(res)

from collections import Counter

class Solution:
    def countOfAtoms(self, formula: str) -> str:
        l = list(lexer(formula))
        stack = [1]
        counter = Counter()
        i = len(l) - 1
        while i >= 0:
            token = l[i]
            print(token)
            i -= 1
            if isinstance(token, int):
                num = token
                token = l[i]
                print(token)
                i -= 1
            else:
                num = 1
            num = num * stack[-1]
            if token == ')':
                stack.append(num)
            elif token == '(':
                stack.pop()
            else:
                counter[token] += num
        kv_to_string = lambda p: p[0] + (str(p[1]) if p[1] > 1 else '')
        res = ''.join( kv_to_string((k, counter[k])) for k in sorted(list(counter.keys())))
        return res

