from typing import Sequence
from typing import Union
import re


def _parse_problem_sizes(argument: str) -> Sequence[Union[int, Sequence[int]]]:
  """Parse a problem size argument into a possibly nested integer sequence.

  Examples:
  64,128 -> [64, 128]
  32,32,[1,1] -> [32, 32, [1, 1]]
  """
  problem_sizes = []
  while argument:
    # Match size.
    match = re.match(r"""[,]?\d+""", argument)
    if match:
      problem_sizes.append(int(match.group().lstrip(',')))
      argument = argument[match.end():]
      continue
    # Match nested sizes.
    match = re.match(r"""[,]?\[[0-9,]+\]""", argument)
    if match:
      nested = match.group().lstrip(',')[1:-1]
      problem_sizes.append([int(elem) for elem in nested.split(',')])
      argument = argument[match.end():]
      continue
    raise ValueError()
  return problem_sizes