#! /usr/bin/python3
# Last edited on 2026-09-30 18:01:01 by stolfi

# Reads from {stdin} a stream of words in EVA or pinyin with the elements marked off by {}.
# Tabulates on {stdout} pairs of elements in the same token or in tokens separated by 
# a fixed number of tokens.

import sys, os, re
from sys import stdin as inp, stderr as err, stdout as out
import html_gen as h
from math import log, hypot

def main(utype:str, wdelta:int, which0:str, which1:str):
  # Reads from {stdin} a file tith tokens separated by spaces and
  # factored into /elems/ (/elements/) by braces {}. Tabulates the
  # frequencies of pairs of elements or words within each token or in
  # different tokens.
   
  assert which0 in { "i", "f", "w", "r" }, f"bad {which0 = !r}"
  assert which1 in { "i", "f", "w", "r" }, f"bad {which1 = !r}"
  
  evoid = "•"

  # If [wdelta} is zero and {which0=initia1} tabulates the pairs of
  # consecutive elements within the same word. Otherwise tabulates some
  # substring of token {w0=t[k]} against some substring of token
  # {w1=t[k+wdelta]}. 
  #
  # The {utype} can be "ps" to indicate that the input tokens are separated
  # piyin syllables, or "wp" to indicate that they are Voynichese tokens 
  # in EVA.
  #
  # The parameters {which0} and {which1} (which should be
  # single-chatacter strings) determine the substring extracted from
  # {wd0} and {wd1}. If {which0} is "i" considers the first
  # "interesting" elem of {w0}; if {which0} is "f" considers the last
  # "interesting" elem. If {which0} is "w", considers the whole of
  # {wd0}. If {which0} is "r" considers all elements after the first
  # "intersting" element of {wd0}. Similarly for {which1} and {wd1}.
  #
  # For pinyin input, the first "interesting" elem is the initial
  # consonant, or {evoid} if the syllable starts with a vowel; the last
  # "interesting" elem is the final consonant ("n", "ng", or "r"), or
  # {evoid} if none. With the "w" selector the whole word is tabulated. With
  # the "r" is selector tabulates whole word minus the initial
  # consonant, or the whole word if it starts with a vowel.
  # 
  # For EVA input, the "interesting" elems are all elems except the
  # "circles" [aoy]. If a token consists entirely of circles (like @y)
  # the first and last "interesting" tokens are {evoid}.
  # With "w" selector the whole word is used. With "r" the whole word
  # minus any prefix circles and the the first non-circle element,
  # or {evoid} if the word has only circles.
  # 
  # Pinyin syllables must have the initial "w" converted to "{u}", or
  # deleted if the next element is "{u}" in any of the tones. Similarly
  # "y" whould have been replaced by "{i}", unless the next element is
  # "{i}" in any of its tones. 
  
  # Syllables that start (resp. end) with a vowel are assumed to have
  # {evoid} as the initial (resp. final) consonant.
  
  write_preamble(utype, wdelta, which0, which1)

  lines = inp.readlines()
  err.write(f"{len(lines):6d} lines read\n")

  ndata = 0 # Num of data lines (excluding blanks, comments).

  npair = 0
  
  # If intial/final pairs only:
  npair_toks_bad = 0 # Num token pairs ignored because of '?'. 
  npair_toks_gud = 0 # Num token pairs considered. 
  
  # If all in-word digraphs only:
  ntoks_bad = 0 # Num tokens ignored because of '?'. 
  ntoks_gud = 0 # Num tokens considered. 
  
  cts_pair = dict()
  cts_e0 = dict()
  cts_e1 = dict()
  
  def tally_pair(el0, el1):
    nonlocal npair, cts_pair, cts_e0, cts_e1
    pair = (el0, el1)
    if pair not in cts_pair: cts_pair[pair] = 0
    cts_pair[pair] += 1
    if pair[0] not in cts_e0: cts_e0[pair[0]] = 0
    cts_e0[pair[0]] += 1
    if pair[1] not in cts_e1: cts_e1[pair[1]] = 0
    cts_e1[pair[1]] += 1
    npair += 1
    return
    # ....................................................................
    
  def split_into_elements(wd):
    wd = re.sub(r"}{", " ", wd)
    wd = re.sub(r"[{}]", "", wd)
    els = wd.split()
    return els
    # ....................................................................

  def get_substring_to_tabulate(els, which):
    # Assumes that {els} is the list of elements of a valid word of type
    # {utype}, without the braces. Returns the substring selected by {which}
    # (first "interesting" elem, or last "intersting" elem, or whole word); 
    # or {evoid} if none.

    e = None
    if utype == "ps":
      e = get_pinyin_substring(els, which)
    elif utype == "wp":
      e = get_eva_substring(els, which)
    else:
      assert False, f"bad {utype = !r}"
    return e
    # ....................................................................

  pinyin_initial_cons = \
    { "ch", "sh", "ts", "zh", 
      "b", "c", "d", "f", "g", "h", "j", "k", "l", "m", "n", "p", 
      "q", "r", "s", "t", "x", "z"
    }
  pinyin_final_cons = { "n", "ng", "r" }

  def get_pinyin_substring(els, which):
    # Assumes that {els} is the list of elements of a valid pinyin
    # syllable, without the braces.
    #
    # If {which} is "i" or "f", returns
    # the initial or final consonant element, respectively, or {evoid}
    # if none. In either case, ignores a final numeric tone [1-5] and/or
    # a disamb code "[.][0-9]". If {which} is "w", returns the whole
    # word, including eventual numeric tone and disamb.
    # If {which} is "r", returns the whole word minus the initial
    # consonant, if any.
    
    assert len(els) >= 1, "empty word"

    e = None
    if which == "i":
      e = els[0];
      if e not in pinyin_initial_cons:
        e = evoid
    elif which == "f":
      k = len(els) - 1;
      while k > 0 and re.fullmatch(r"[.]?[0-9]+", els[k]) is not None: 
        k = k - 1
      if k >= 0 and els[k] in pinyin_final_cons:
        e = els[k]
      else:
        e = evoid
    elif which == "w":
      # Whole word:
      e = "".join(els)
    elif which == "r":
      # Whole word minus initial consonant if any:
      if els[0] in pinyin_initial_cons: els = els[1:]
      e = evoid if len(els) == 0 else "".join(els)
    else:
      assert False, f"invalid {which = !r}"
    return e
    # ....................................................................

  eva_circles = { "a", "o", "y" }

  def get_eva_substring(els, which):
    # Assumes that {els} is the list of elements without the 
    # braces, and the word is valid.
    # 
    # If {which} is "i" or "f", eturns the initial or final element,
    # ignoring the circles, or {evoid} if all are circles.
    # If {which} is "w", returns the whole word, including circles.
    # if {which} is "r" returns word minus leading circles and
    # first non-circle, or {evoid} if it is all circles.

    n = len(els)
    assert n >= 1, "empty word"
    
    e = None
    if which == "i":
      k = 0
      while k < n and els[k] in eva_circles: k = k+1
      e = evoid if k >= n else els[k]
    elif which == "f":
      k = n - 1
      while k > 0 and els[k] in eva_circles: k = k-1
      e = evoid if k < 0 else els[k]
    elif which == "w":
      # Whole word:
      e = "".join(els)
    elif which == "r":
      # Whole word minus leading circles and first non-circle.
      k = 0
      while k < n and els[k] in eva_circles: k = k+1
      k = k + 1;
      e = evoid if k >= n else "".join(els[k:])
    else:
      assert False, f"invalid {which = !r}"
    return e
    # ....................................................................

  def tally_end_elem_pairs (wdA, whichA, wdB, whichB):
    nonlocal npair_toks_gud, npair_toks_bad
    wdA_ok = (re.search(r"[?]", wdA) is None)
    wdB_ok = (re.search(r"[?]", wdB) is None)
    if wdA_ok and wdB_ok:
      elsA = split_into_elements(wdA)
      eA = get_substring_to_tabulate(elsA, whichA)
      elsB = split_into_elements(wdB)
      eB = get_substring_to_tabulate(elsB, whichB)
      tally_pair(eA, eB)
      npair_toks_gud += 1
    else:
      npair_toks_bad += 1
    return
    # ....................................................................
    
  def tally_all_elem_pairs(wd):
    nonlocal ntoks_gud, ntoks_bad
    wd_ok = (re.search(r"[?]", wd) is None)
    if wd_ok:
      els = split_into_elements(wd)
      for i in range(len(els)-1):
        tally_pair(els[i], els[i+1])
      ntoks_gud += 1
    else:
      ntoks_bad += 1
    return
    # ....................................................................
  
  for line in lines:
    line = re.sub(r"[#].*$", "", line)
    if line == "": continue
    ndata += 1
    line = line.strip()
    m = re.fullmatch(r"[<]([a-zA-Z0-9.;]+)[>] *(.*)", line)
    if m is None:
      loc = None; text = line
    else:
      loc = m.group(1); text = m.group(2)
      
    words = text.split()
    n = len(words)
    for k in range(n-abs(wdelta)):
      wd0 = words[k]
      if wdelta == 0 and which0 == which1:
        tally_all_elem_pairs(wd0)
      else:
        wd1 = words[k + abs(wdelta)]
        if wdelta > 0:
          tally_end_elem_pairs(wd0, which0, wd1, which1)
        else:
          tally_end_elem_pairs(wd1, which1, wd0, which0)

  err.write(f"{ndata:6d} data lines processed\n")
  if wdelta == 0 and which0 == which1:
    err.write(f"{ntoks_bad:6d} invalid tokens, ignored\n")
    err.write(f"{ntoks_gud:6d} good tokens considered\n")
  else:
    err.write(f"{npair_toks_bad:6d} invalid token pairs, ignored\n")
    err.write(f"{npair_toks_gud:6d} good token pairs considered\n")
  err.write(f"{npair:6d} item pairs talied\n")
  
  # Get lits of elements for rows and columns:
  min_occs = 0
  
  e0_list = list(cts_e0.keys()) 
  e0_list.sort(key = lambda e: cts_e0[e], reverse = True); 
  n0 = len(e0_list)
  n0_max = 10000
  if which0 == "w" or which0 == "r":
    e0_list = e0_list[:n0_max]
    e0_list = [ e for e in e0_list if cts_e0[e] >= min_occs ]
  trunc0 = len(e0_list) < n0 # Is {e0_list} whole or truncated?
  n0 = len(e0_list)
  if trunc0: out.write(f"# Rows limited to the {n0} most common entries.\n")
  
  e1_list = list(cts_e1.keys());
  e1_list.sort(key = lambda e: cts_e1[e], reverse = True); 
  n1 = len(e1_list)
  n1_max = 50
  if which1 == "w" or which1 == "r":
    e1_list = e1_list[:n1_max]
    e1_list = [ e for e in e1_list if cts_e1[e] >= min_occs ]
  trunc1 = len(e1_list) < n1 # Is {e1_list} whole or truncated?
  n1 = len(e1_list)
  if trunc0: out.write(f"# Columns limited to the {n1} most common entries.\n")
 
  e0_wd = 3
  for e0 in e0_list: e0_wd = max(e0_wd, len(e0))
 
  # Fill holes in table with zeros:
  for e0 in e0_list:
    for e1 in e1_list:
      if (e0,e1) not in cts_pair: cts_pair[(e0,e1)] = 0

  def print_header_row():
    xe0 = " " * e0_wd
    xtotct = "Tot"
    xtotpc = ""
    out.write(f"{xe0}  {xtotct:>5s} {xtotpc:>4s}")
    for e1 in e1_list:
      xe1 = '•' if e1 is None or e1 == '' else e1
      xpc1 = " "*4
      out.write(f" {xe1:>5s} {xpc1:>4s}")
    out.write("\n")
    return
    # ....................................................................
  
  def print_dashes_row():
    xe0 = "-" * e0_wd
    xtotct = "-"*5
    xtotpc = "-"*4
    out.write(f"{xe0}  {xtotct:>5s} {xtotpc:>4s}")
    for e1 in e1_list:
      xct1 = "-"*5
      xpc1 = "-"*4
      out.write(f" {xct1:>5s} {xpc1:>4s}")
    out.write("\n")
    return
    # ....................................................................
  
  def print_counts_row(e0, totct, cts):
    n = len(cts)
    assert len(e1_list) == n
    xe0 = '•' if e0 == '' else e0
    xe0 = xe0.ljust(e0_wd)
    xtotct = f"{totct:5d}"
    xtotpc = " "*4 if totct == npair else f"{100*totct/npair:.1f}"
    out.write(f"{xe0}  {xtotct:>5s} {xtotpc:>4s}")
    sum_row = 0
    for i in range(n):
      e1 = e1_list[i]; assert e1 is not None
      ct1 = cts[i]; assert ct1 is not None
      xct1 = f"{ct1:5d}"
      xpc1 = f"{100*ct1/totct:.1f}"
      sum_row += ct1
      out.write(f" {xct1:>5s} {xpc1:>4s}")
    out.write("\n")
    assert (sum_row <= totct if trunc1 else sum_row == totct)
    return
    # ....................................................................

  # Printout
  print_header_row()
  print_dashes_row()
  
  sum_tot = 0
  for e0 in e0_list:
    totct = cts_e0[e0]
    assert totct is not None
    cts = [ cts_pair[(e0,e1)] for e1 in e1_list ]
    print_counts_row(e0, totct, cts);
    sum_tot += totct
  assert (sum_tot <= npair if trunc0 else sum_tot == npair)
  
  print_dashes_row()
  
  totct = npair
  cts = [ cts_e1[e1] for e1 in e1_list ]
  print_counts_row("Tot", totct, cts);

  out.flush()
  
  return
  # ----------------------------------------------------------------------

def write_preamble(utype, wdelta, which0, which1):
  now = h.get_current_date()
  out.write("# \n")
  out.write(f"# Created on {now}")
  out.write(f" with {wdelta = }, {which0 = !r}, {which1 = !r}.\n")

  if utype == "ps":
    xsyll = " (syllable)"
    xeltype = "phonetic" 
    xelclass = "consonant"
    xparag = "recipe"
  elif utype == "wp":
    xsyll = ""
    xeltype = "CMC"
    xelclass = "non-circle elem"
    xparag = "parag"
  else:
    assert False, f"bad {utype = !r}"
  if wdelta == 0 and which0 == which1:
    out.write(f"# Each row is one {xeltype} element of a token{xsyll}.\n")
    out.write(f"# Each column is the next {xeltype} element of the same token.\n")
  else:
    for i in 0, 1:
      if i == 0:
        xrorc = "row"
        xdelta = ""
      else:
        xrorc = "column"
        xdelta = "" if wdelta == 0 else f"{wdelta:+2d}"

      which = ( which0, which1 )[i]
      xchop = ""
      if which == "i": 
        xel = f"the initial {xelclass} of "
      elif which == "f":
        xel = f"the final {xelclass} of "
      elif which == "w":
        xel = "" # The whole token.
      elif which == "r":
        xel = "" # The whole token.
        if utype == "ps":
          xchop = " minus the initial consonant"
        elif utype == "wp":
          xchop = " minus leading circles and first non-circle"
        else:
          assert False
      else:
        assert False, f"bad which{i} = '{which[i]}'"

      out.write(f"# Each {xrorc} is {xel}a token{xsyll} t[k{xdelta}]{xchop}.\n")

  out.write(f"# Considers only tokens in the same {xparag}e.\n")
  out.write("# \n")
  out.write("# The first column is the row total and its percentage over total pairs.\n")
  out.write("# The last row is the column total and its percentage over total pairs.\n")
  out.write("# Each of the other entries is the count for the pair and its percentage over the row total.\n")
  out.write("# \n")
  return
  # ----------------------------------------------------------------------
  
ia = 1
utype = sys.argv[ia]; ia += 1
wdelta = int(sys.argv[ia]); ia += 1
which0 = sys.argv[ia]; ia += 1
which1 = sys.argv[ia]; ia += 1
main(utype, wdelta, which0, which1)