#! /usr/bin/python3 
# Last edited on 2025-10-28 16:07:12 by stolfi

import os, sys, re
from random import randrange as abrand, random as rand, seed as inirand
from math import log, exp, sqrt, hypot, inf, nan, pi, floor, ceil, sin, cos

def main():
  p_reset, p_mutate, n_gen, n_out, run = get_options()
  seed = [ "like", "that", "of", "the", "thing", "called", "a", "siren",
           "in", "our", "manufacturing", "towns", "a", "man", "knee~deep",
           "near", "the", "towing", "path", "shouted", "inaudibly", "to",
           "me", "and", "pointed", "looking", "back", "i", "saw", "the",
           "other", "martians", "advancing", "with", "gigantic", "strides",
           "down", "the", "riverbank", "from", "the", "direction", "of",
           "chertsey", ] 
  S = len(seed)
  sys.stderr.write(f"{seed = }\n")
  sys.stderr.write("\n")
  S = len(seed)
  inirand(floor(123456789*sin(exp(1)*run)))
  sys.stderr.write(f"{p_reset = :5.3f}  {p_mutate = :5.3f}\n")
  sys.stderr.write(f"{S = }  {n_gen = }  {n_out = }  {run = }\n")
  
  # Read the tables for the mutate procedure:
  MTB_due, MTB_tre = read_mutate_tables("in/bitrigrams.wct")
  LTB = read_length_table("in/lengths.wct")

  # Tests:
  # test_mutate(seed, MTB_due, MTB_tre)
  # return 0

  # Runs the T&T algoritm for {n_gen} steps, writes the last {n_out} words.
  text, depth, nmuts = tandt(seed, n_gen, p_reset, p_mutate, LTB, MTB_due, MTB_tre)
  assert len(text) == n_gen
  text_clip = text[n_gen-n_out:]
  depth_clip = depth[n_gen-n_out:]
  nmuts_clip = nmuts[n_gen-n_out:]
  sys.stderr.write(f"len(text_clip) = {len(text_clip)}\n")
  assert len(text_clip) == n_out
  Davg = 0; Mavg = 0
  for i in range(len(text_clip)):
    sys.stdout.write(text_clip[i]); sys.stdout.write("\n")
    Davg += depth_clip[i]
    Mavg += nmuts_clip[i]
  Davg /= n_out
  n_mid = n_gen + 1 - (n_out+1)/2
  Dexp = 0.80 + log(n_mid/S); Derr = Davg/Dexp
  sys.stderr.write(f"{Davg = :6.3f}  expected = {Dexp:6.3f}  err = {Derr:6.3f}\n")
  Mavg /= n_out; Mexp = p_mutate * Dexp; Merr = 1 if Mexp == 0 else Mavg/Mexp
  sys.stderr.write(f"{Mavg = :6.3f}  expected = {Mexp:6.3f}  err = {Merr:6.3f}\n")

  return 0
  # ----------------------------------------------------------------------
  
def get_options():
  karg = 0;
  karg += 1; p_reset = float(sys.argv[karg])/1000
  karg += 1; p_mutate = float(sys.argv[karg])/1000
  karg += 1; n_gen = int(sys.argv[karg])
  karg += 1; n_out = int(sys.argv[karg])
  karg += 1; run = int(sys.argv[karg])
  return p_reset, p_mutate, n_gen, n_out, run
  # ----------------------------------------------------------------------
    
def read_mutate_tables(fname):
  with open(fname, "r") as rd:
    nread = 0
    MTB_due = {}
    MTB_tre = {}
    for lin in rd:
      lin = lin.rstrip()
      nread += 1
      if not re.match(r"[ ]*([#]|$)", lin):
        m = re.fullmatch(r" *([0-9]+)[ ]+([^ ]+)[ ]*", lin)
        assert m != None, f"bad table line {nread}: {lin = }"
        ct = int(m.group(1))
        btg = m.group(2)
        if len(btg) == 2:
          # Add to bigram table:
          bg = btg
          if not bg in MTB_due: MTB_due[bg] = []
          MTB_due[bg] = ct
        elif len(btg) == 3:
          # Add to trigram table:
          tg = btg
          bg = btg[0] + btg[2]
          if not bg in MTB_tre: MTB_tre[bg] = []
          MTB_tre[bg].append((btg[1], ct))
    rd.close()
  # Convert {MTB_due} entries to probabilities:
  sum = 0
  for bg in MTB_due.keys(): sum += MTB_due[bg]
  for bg in MTB_due.keys(): MTB_due[bg] /= sum
  
  # Convert each list of {MTB_tre} to probabilities:
  for bg in MTB_tre.keys(): 
    R = MTB_tre[bg]
    n = len(R)
    sum = 0
    for i in range(n): sum += R[i][1]
    for i in range(n): R[i] = (R[i][0], R[i][1]/sum)
  return MTB_due, MTB_tre
  # ----------------------------------------------------------------------
    
def read_length_table(fname):
  with open(fname, "r") as rd:
    nread = 0
    LTB = []
    for lin in rd:
      lin = lin.rstrip()
      nread += 1
      if not re.match(r"[ ]*([#]|$)", lin):
        m = re.fullmatch(r" *([0-9]+)[ ]+([0-9]+)[ ]*", lin)
        assert m != None, f"bad table line {nread}: {lin = }"
        ct = int(m.group(1))
        L = int(m.group(2))
        while len(LTB) <= L: LTB.append(0)
        assert LTB[L] == 0, f"dup length {L = }"
        LTB[L] = ct
    rd.close()
  return LTB
  # ----------------------------------------------------------------------

def tandt(seed, n_gen, p_reset, p_mutate, LTB, MTB_due, MTB_tre):
  S = len(seed)
  text = seed.copy()
  depth = [ 0 ]*S
  nmuts = [ 0 ]*S
  n_trace = 50 # n_gen # 150
  k = 0;
  sys.stderr.write(f"{len(text)},{k}!")
  for i in range(n_gen):
    trace = (i >= 0 and i < n_trace) or i >= n_gen - n_trace
    if i == n_trace: sys.stderr.write(" ... ")
    if p_reset == 0:
      w = "phleghm"
      d = 0
      m = int(floor(1000*p_mutate + 0.5))
      for j in range(m):
        w = mutate(w, LTB, MTB_due, MTB_tre); m += 1
    else:
      if rand() < p_reset:
        nT = len(text)
        # kini = max(0, nT - int(sqrt(nT-S)))
        # klim = min(nT, max(kini + 1, nT-S//2+1)) 
        # kini = 0; klim = nT
        kini = nT//2; klim = nT

        # k = abrand(kini, klim)
        k = max(abrand(kini, klim), abrand(kini, klim))
        # k = max(abrand(kini, klim), abrand(kini, klim), abrand(kini, klim))
        if trace: sys.stderr.write(f"{nT},{k}!")
      w = text[k];
      d = depth[k] + 1;
      m = nmuts[k]
      k += 1;
      if rand() < p_mutate:
        w = mutate(w, LTB, MTB_due, MTB_tre); m += 1
        # if trace: sys.stderr.write("M!")
    if trace: sys.stderr.write(f"{w}.")
    text.append(w)
    depth.append(d)
    nmuts.append(m)
  sys.stderr.write("\n") 
  return text[S:], depth[S:], nmuts[S:]
  # ----------------------------------------------------------------------
  
def txout(text):
  width = 70
  n = 0; s = "  "
  for w in text:
    if n + len(s) + len(w) > width:
      sys.stdout.write("\n"); 
      n = 0; s = "  "
    sys.stdout.write(s);
    sys.stdout.write(w);
    n += len(s) + len(w); s = " "
  sys.stdout.write("\n")
  return 
  # ----------------------------------------------------------------------
  
def mutate(w, LTB, MTB_due, MTB_tre):
  Lw = len(w)
  p_insert, p_change, p_delete = length_change_probs(Lw, LTB)
  if rand() < p_change:
    w = mutate_change(w, MTB_tre)
  elif rand()*(1 - p_change) < p_delete:
    assert len(w) >= 2, "bug p_delete"
    w = mutate_delete(w, MTB_due)
  else:
    w = mutate_insert(w, MTB_tre)
  return w
  # ----------------------------------------------------------------------
  
def mutate_change(w, MTB_tre):
  # Assumes {MTB_tre} is a dict that maps length 2 strings "{x}{z}"
  # to a list of pairs {(y, pr)} where {y} is a single-letter string
  # and {pr} is the count of the trigram "{x}{y}{z}".
  
  # Enumerate all choices for replacements:
  w = "~" + w + "~"
  Lw = len(w)
  C = [None] * (Lw-2)
  
  for iw in range(Lw-2):
    C[iw] = []
    ch = w[iw+1] # Current {w[iw+1]}.
    bg = w[iw] + w[iw+2] # Context of {ch}.
    if bg in MTB_tre:
      R = MTB_tre[bg]
      for ch1, ct1 in R:
        if ch1 != ch: 
          C[iw].append((ch1, ct1,))
  iwr, chr = choose_index_and_letter(C)
  if iwr >= 0:
    w = w[:iwr+1] + chr + w[iwr+2:]    
  else:
    sys.stderr.write(f"!! no suitable replacement for '{w}'\n")
  w = w[1:len(w)-1]
  return w
  # ----------------------------------------------------------------------
  
def mutate_insert(w, MTB_tre):
  # Assumes {MTB_tre} is a dict that maps length 2 strings "{x}{z}"
  # to a list of pairs {(y, pr)} where {y} is a single-letter string
  # and {pr} is the count of the trigram "{x}{y}{z}".
  
  # Enumerate all choices for insertions:
  w = "~" + w + "~"
  Lw = len(w)
  C = [None] * (Lw-1)

  for iw in range(Lw-1):
    C[iw] = []
    bg = w[iw:iw+2] # Context of for insertion.
    if bg in MTB_tre:
      R = MTB_tre[bg]
      for ch1, ct1 in R:
        C[iw].append((ch1, ct1,))
  iwi, chi = choose_index_and_letter(C)
  if iwi >= 0:
    w = w[:iwi+1] + chi + w[iwi+1:]    
  else:
    sys.stderr.write(f"!! no suitable insertion for '{w}'\n")
  w = w[1:len(w)-1]
  return w
  # ----------------------------------------------------------------------
   
def choose_index_and_letter(C):
  # Assumes that each element {C[iw]} is a list of pairs {(y,ct)}
  # where {y} is a letter and {ct} is the occurence count 
  # of the trigram "{x}{y}{z}" for some context chars {x} and {z}.
  nC = len(C)
  sum = 0
  for iw in range(nC):
    for ch1, ct1 in C[iw]:
      sum += ct1
  if sum == 0:
    return -1, "?"
  else:
    # Choose index and letter according to trigram freqs:
    coin = rand()
    cum = 0;
    for iw in range(nC):
      for ch1, ct1 in C[iw]:
        cum += ct1/sum
        if coin <= cum:
          return iw, ch1
  assert False
  # ----------------------------------------------------------------------
  
def mutate_delete(w, MTB_due):
  # Assumes {MTB_due} is a dict that maps length 2 strings "{x}{z}"
  # to the count of the digraph "{x}{z}".
  
  # Enumerate all choices for replacements:
  w = "~" + w + "~"
  Lw = len(w)
  C = [ 0 ] * (Lw-2)

  for iw in range(Lw-2):
    bg = w[iw] + w[iw+2] # Context of {ch}.
    if bg in MTB_due:
      ct = MTB_due[bg]
      C[iw] = ct
  iwd = choose_index(C)
  if iwd >= 0:
    w = w[:iwd+1] + w[iwd+2:]    
  else:
    sys.stderr.write(f"!! no suitable deletion from '{w}'\n")
  w = w[1:len(w)-1]
  return w
  # ----------------------------------------------------------------------

def choose_index(C):
  # Assumes {C} is a list of integers. Selects an index {iw}
  # into {C} with probability proportional to {C[iw]}. 
  # Returns {-1} if all elements of {C} are zero.
  nC = len(C)
  sum = 0
  for iw in range(nC): sum += C[iw]
  if sum == 0:
    return -1
  else:
    # Choose replacement according to trigram freqs:
    coin = rand()
    cum = 0;
    for iw in range(nC):
      cum += C[iw]/sum
      if coin <= cum:
        return iw
  assert False
  # ....................................................................

def test_mutate(seed, MTB_due, MTB_tre):
  sys.stderr.write("testing {mutate_change}, {mutate_delete}  ...\n")
  for sw in seed:
    w = sw; sys.stderr.write(f"{w}")
    for k in range(15):
      for tag, func in (("I", mutate_insert,), ("C", mutate_change,), ("D",mutate_delete),):
        MTB = MTB_due if tag == "D" else MTB_tre
        w = func(w, MTB); sys.stderr.write(f" ({tag}){w}")
        assert re.fullmatch(r"[-a-z']+", w)
    sys.stderr.write("\n")
  return
  # ----------------------------------------------------------------------

def length_change_probs(L, LTB):
  # Assumes {LTB[k]} is the count of tokens with length {k}.
  # Returns {p_insert,p_change,p_delete} proportional
  # to the sum of {LTB[k]} for {k>L}, {k=L}, and {k<L},
  # except that the {k=L} case is boosted by {3×}.
  
  Lmax = len(LTB)
  assert L >= 0, f"bad {L = }"
  sum_lo = 0; sum_eq = 0; sum_hi = 0
  for k in range(Lmax):
    if k < L: sum_lo += LTB[k]
    elif k > L: sum_hi += LTB[k]
    else: sum_eq += 3*LTB[k]
  sum = sum_lo + sum_eq + sum_hi
  assert sum > 0, "bug sum LTB"
  p_delete = 0.0 if L <= 0 else 1.0 if L >= Lmax else sum_lo/sum
  p_insert = 0.0 if L >= Lmax else sum_hi/sum
  p_change = 1.0 - (p_insert + p_delete)
  return p_insert, p_change, p_delete
  # ----------------------------------------------------------------------

main()