package main

import (
	"sort"
	"strconv"
	"unicode"
)

type atom struct {
	name string
	num  int
}

func countOfAtoms(formula string) string {
	// 1.建立一个类型为map[string]int  key[原子名]value[原子数量] 的stack
	stack := make([]map[string]int, 0, 0)
	// 2.Push空的map作为最外层
	stack = append(stack, map[string]int{})
	// 3.遍历formula，进行压栈、弹栈
	for i := 0; i < len(formula); i++ {
		if formula[i] == '(' {
			// 如果这是一个左括号，压入空栈
			stack = append(stack, map[string]int{})
		} else if unicode.IsUpper(rune(formula[i])) {
			// 如果是大写字母，则这是一个原子，找到表示原子的字符串
			top := stack[len(stack)-1]
			str, num := findAtomName(formula[i:])
			top[str] += num
		} else if formula[i] == ')' {
			// 如果这是一个右括号，判定是否有数组在括号之后，如果有，进行乘法运算，然后弹栈，加入下一个栈顶
			num := parseNum(i, len(formula), formula)
			top := stack[len(stack)-1]
			// 弹栈
			stack = stack[:len(stack)-1]
			// 将括号内的原子数乘上num ; 加到上一层的原子数量中
			for atom, v := range top {
				stack[len(stack)-1][atom] += v * num
			}
		}
	}
	// 4.将最后一个栈顶map弹出，遍历加入结构体数组，并对其进行排序，后拼接成为字符串输出
	top := stack[0]
	atomArr := make([]atom, 0, 0)
	for k, v := range top {
		atomArr = append(atomArr, atom{k, v})
	}
	sort.Slice(atomArr, func(i, j int) bool {
		return atomArr[i].name < atomArr[j].name
	})
	// byte数组转str，避免字符串拼接
	res := []byte{}
	for _, v := range atomArr {
		res = append(res, []byte(v.name)...)
		// num 可能大于9， 故不能直接加减运算
		if v.num > 1 {
			res = append(res, strconv.Itoa(v.num)...)

		}
	}
	return string(res)
}

func findAtomName(s string) (str string, num int) {
	left, right := 0, 1
	for right < len(s) && !unicode.IsUpper(rune(s[right])) {
		right++
	}
	if unicode.IsDigit(rune(s[right])) {
		num = int(s[right]) - '0'
	}
	str = s[left:right]
	return
}

func parseNum(i, n int, formula string) (num int) {
	if i == n || !unicode.IsDigit(rune(formula[i])) {
		return 1
	}
	for ; i < n && unicode.IsDigit(rune(formula[i])); i++ {
		num = num*10 + int(formula[i]-'0')
	}
	return
}