# Python script to demonstrate almost linear time consistent collatz simulation

import time

from gmpy2 import mpz, bit_mask, log2, log10, floor

# Start: N(4, 4)
# N(2n,m+2)   --> N(3n+3,m)
# N(2n+1,m+1) --> N(3n+4,m)
#
# N(2n+4,0)   --> halt
# N(2n+4,1)   --> N(4,6n+20)
# N(2n+5,0)   --> N(4,6n+22)

def big_str(n):
  if n > 10**100:
    ln = int(log10(n))
    return f"10^{ln} ≡ {n%2} (mod 2)"
  else:
    return str(n)


# Direct computation of t steps
def direct(a, t):
  b = 0
  for _ in range(t):
    b += 2 - (a%2)
    # 2k+r -> 3k + (r+3)
    a += (a>>1) + 3
  return (a, b)

# Accelerated computation of 2**e steps
def accel_pow(a, e):
  if e<8: return direct(a,1<<e)
  t = 1<<(e-1)
  p3t = mpz(3)**t
  m = bit_mask(t)
  (a1, b1) = accel_pow(a&m,e-1)
  a = p3t*(a>>t) + a1
  (a2, b2) = accel_pow(a&m,e-1)
  return (p3t*(a>>t) + a2, b1+b2)


def single(a, b):
  (k, r) = divmod(a,2)
  if r == 0:
    if b >= 2:
      return (3*k+3, b-2)
    if b == 1:
      return (4, 6*k+8)
    if b == 0:
      raise Exception("Halt")
  else:
    if b >= 1:
      return (3*k+4, b-1)
    else:
      return (4, 6*k+10)

def next_reset(a, b):
  """Start from N(a,b), run until the next "reset" (where b increases or halt)."""
  # t = Max steps we can run direct(a, t) without danger of going negative with b.
  while (t := (b-1) // 2) > 0:
    # e = Max pow we can run accel_pow(a, e) without danger of going negative with b.
    e = int(floor(log2(t)))
    (a, bm) = accel_pow(a, e)
    b -= bm
    print(" ", big_str(a), big_str(b), t, e, time.process_time())
  assert b >= 0
  while a != 4:
    (a, b) = single(a, b)
    print(" ", big_str(a), big_str(b), t, e, time.process_time())
  return (a, b)

def sim(a, b):
  a = mpz(a)
  print(big_str(a), big_str(b), time.process_time())
  while True:
    (a, b) = next_reset(a, b)
    print(big_str(a), big_str(b), time.process_time())

# (a, b) = (4, 4)
# for _ in range(12):
#   (a, b) = single(a, b)
#   print(a, b)
# print()

# Start: N(4, 4)
sim(4, 4)
