import os
import sys
import argparse

# CONVERTS AN INTEGER INTO ITS BASE-N REPRESENTATION EXPRESSED AS THREE ALIGNED LISTS
# num = \sum_{i = 0}^{k} (c_i)*(b_i)^(e_i)
# THE FIRST LIST (c_0, ..., c_k) CONTAINS THE COEFFICIENTS
# THE SECOND LIST (b_0, ..., b_k) CONTAINS THE BASES (WHICH SHOULD ALL BE THE SAME)
# THE THIRD LIST (e_0, ..., e_k) CONTAINS THE EXPONENTS
def int2repr(num: int, base: int) -> [[int], [int], [int]]:
  assert isinstance(num, int) and isinstance(base, int) and base >= 2

  coefs = []
  bases = []
  exps = []
  k = 0

  while num > 0:
    coefs = [num % base] + coefs
    bases = [base] + bases  # OBVIOUSLY THE BASE IS THE SAME IN EACH TERM, BUT WWHATEVER LOL
    exps = [k] + exps
    num = num // base
    k = k + 1

  return coefs, bases, exps


# CONVERTS FROM " \sum_{i = 0}^{k} (c_i) * (b_i)^(e_i) " TO INTEGER
# CALCULATES THE INTEGER CORRESPONDING TO A STRING REPRESENTATION OF A LINEAR COMBINATION
# THIS IS !!!NOT!!! A FULL PARSE OF THE ABSTRACT SYNTAX TREE FOR AN ARBITRARY ALGEBRAIC EXPRESSION
def str2int(string: str) -> int:
  def _str2int(string: str, idx: int = 0) -> (int, int):
    assert idx < len(string) or len(string) == 0

    if len(string) == 0:
      return 1, -1

    string = string.replace(" ", "")

    addition = lambda x, y: x + y
    multiplication = lambda x, y: x*y
    exponentiation = lambda x, y: x**y

    operation = addition
    result = 0

    while idx < len(string):
      char = string[idx]
      try:
        nit = str(int(char))
        idx += 1
        try:
          while True:
            nit += str(int(string[idx]))
            idx += 1
        except:
          pass
        result = operation(result, int(nit))
      except:
        if char == "+":
          operation = addition
        elif char == "*" or char == "·":
          operation = multiplication
        elif char == "^":
          operation = exponentiation
        elif char == "(":
          subresult, idx = _str2int(string, idx=idx + 1)
          result = operation(result, subresult)
        elif char == ")":
          return result, idx
        else:
          raise AssertionError(char)
        idx += 1

    return result, idx

  return _str2int(string, idx=0)[0]


# COMPUTES A STRING CORRESPONDING TO THE HEREDITARY BASE-N NOTATION FOR A NATURAL NUMBER'S REPRESENTED IN BASE-N
# RECURSIVELY CONVERT (c_0, ..., c_k), (b_0, ..., b_k), (e_0, ..., e_k) TO " \sum_{i = 0}^{k} (c_i) * (b_i)^(e_i) "
def repr2str(coefs: [int], bases: [int], exps: [int], simple: bool = False) -> str:
  assert isinstance(coefs, list) and isinstance(bases, list) and isinstance(exps, list)

  string = ""

  for c, b, e in zip(coefs, bases, exps):
    if c != 0:
      if e != 0:
        e_coefs, e_bases, e_exps = int2repr(e, b)
        substring = repr2str(e_coefs, e_bases, e_exps, simple=simple)
        if simple:
          string += f"{c}·{b}^({substring}) + "
        else:
          string += f"({c}·({b}^({substring}))) + "
      else:
        if simple:
          string += f"{c}*{b}^{e} + "
        else:
          string += f"({c}*({b}^({e}))) + "

  return string.strip().strip("+").strip()


def hereditary(num: int, base: int = 2, simple: bool = True):
  return repr2str(*int2repr(num, base), simple=simple)


def goodstein(num: int, base: int = 2) -> int:
  assert isinstance(num, int) and isinstance(base, int) and base >= 2

  if base == 2:
    coefs, bases, exps = int2repr(num, base)
    string = repr2str(coefs, bases, exps)
    G = str2int(string)
  else:
    prev_G = goodstein(num, base - 1)
    prev_str = hereditary(prev_G, base=base - 1, simple=False)
    new_str = prev_str.replace(str(base - 1), str(base))
    G = str2int(new_str) - 1

  return G


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument("x", type=int)  # INPUT ANCHOR FOR SEQUENCE
  parser.add_argument("-e", "--end",  type=int, default=5)  # POSITION IN SEQUENCE TO STOP AT
  parser.add_argument("-d", "--digits", type=int, default=-1)  # MAX NUMBER OF DIGITS TO CONVERT TO STRING
  parser.add_argument("-v", "--verbose", action="store_true")  # SHOW HEREDITARY BASE-N AND INTERMEDIATE WORK
  args = parser.parse_args()

  separator = "\n" + "-----------------------------------------------------------------------------------"

  if args.digits > -1:
    sys.set_int_max_str_digits(args.digits)

  try:
    if args.verbose:
      for n in range(2, args.end + 1):
        if n == 2:
          prev = goodstein(args.x, base=n)
          Gn = prev
          # print(f"G{n}({x})\t{Gn}\t{hereditary(Gn, base=n)}\n"
          print(f"\t\t{hereditary(Gn, base=n)}"
                + "\n"
                + f"G{n}({args.x})\t{Gn}"
                + (separator if n != args.end else ""))
        else:
          prev = Gn
          Gn = goodstein(args.x, base=n)
          print(f"\t\t{hereditary(prev, base=n - 1)}"
                + "\n"
                + f"\t\t{hereditary(prev, base=n - 1).replace(str(n - 1), str(n))} - 1"
                + "\n"
                + f"\t\t{hereditary(Gn, base=n)}"
                + "\n"
                + f"G{n}({args.x})\t{Gn}"
                + (separator if n != args.end else ""))
    else:
      for n in range(2, args.end + 1):
        if n == 2:
          prev = goodstein(args.x, base=n)
          Gn = prev
          print(f"G{n}({args.x})\t{Gn}")
        else:
          prev = Gn
          Gn = goodstein(args.x, base=n)
          print(f"G{n}({args.x})\t{Gn}")
  except ValueError as err:
    _CODES = {"RED": 31, "GREEN": 32, "YELLOW": 33, "BLUE": 34, "MAGENTA": 35, "CYAN": 36}
    def color(s, name, bold=False, stream=None):
      stream = sys.stdout if stream is None else stream
      if os.environ.get("NO_COLOR") or not hasattr(stream, "isatty") or not stream.isatty():
        return str(s)
      return f"\033[{'1;' if bold else ''}{_CODES[name]}m{s}\033[0m"
    err.add_note(color("Try setting the ", "MAGENTA") +  color("\"--digit\"", "BLUE") + color(" argument if Python complains about the number of digits in an integer string conversion", "MAGENTA"))
    err.add_note(color("Example: ", "MAGENTA") + color("python goodstein.py 41 --end 6 ", "GREEN") + color("--digits 40000", "BLUE"))
    raise

main()
