#! /usr/bin/gawk -f 
# Last edited on 2026-07-11 06:56:53 by stolfi

BEGIN {
  abort = -1;
  usage = ( \
    "compute-cond-entropy < INFILE > OUTFILE" \
  );
  
  # Expects the input to have records of the form 
  # 
  #   COUNT EVENT SYMB
  # 
  # representing COUNT occurrences of SYMB after the specified EVENT.
  # The records must be sorted by EVENT then SYMB.
  # 
  # Outputs for each EVENT one line with format
  # 
  #   NUMSYMB TOTCOUNT ENTROPY EVENT
  # 
  # where NUMSYMB is the number of distinct SYMBs for the event,
  # TOTCOUNT is the sum of all COUNT for the EVENT, and ENTROPY is the
  # conditional entropy of the next SYMB given that EVENT. Assumes that
  # all events are mutually exclusive, and expects all SYMB for each
  # EVENT are distinct.
  #
  # Also writes to {stderr} the total {SYMB} count for all events
  # and the average next-symbol conditional entropy.
  
  global_numEvents = 0      # Number of distinct {EVENT}s seen.
  global_numPairs = 0       # Number of distincy {EVENT}-{SYMB} pairs read.
  global_totCt = 0          # Sum of {TOTCOUNT} for all events.
  global_totCtEntropy = 0   # Sum of {TOTCOUNT*ENTROPY} for all events.
  
  start_event("")
}

function start_event(event)
{
  curEvent = event;            # Current EVENT.
  prevSymb = "";               # Previosu SYMB of current EVENT.
  curEvent_TotCt = 0;          # Total COUNT for the current event.
  curEvent_nSymb = 0;          # Count of symbols for this event.
  split("", curEvent_SymbCt);  # Counts of each SYMB for the current event.
}

(abort >= 0) { exit abort; }

/./ {
  if (NF != 3) { data_error("bad field count"); }
  count = $1;  event = $2;  symb = $3;
  if (event != curEvent)
    { if (event == "") { prog_error("bug event"); }
      if (event < curEvent) { data_error("not sorted by EVENT"); }
      totalize_event_entropy();
      start_event(event);
    }
  if (symb == "" ) { prog_error("bug symb"); }
  if (symb < prevSymb) { data_error("not sorted by SYMB"); }
  if (count !~ /^[+]*[.0-9]*[0-9][.0-9]*$/) { data_error("bad count format"); }
  curEvent_SymbCt[curEvent_nSymb] = count; 
  curEvent_nSymb += 1;
  curEvent_TotCt += count;
  global_numPairs += 1
  prevSymb = symb;
  next;
}

END {
  totalize_event_entropy();
  printf "%6d distinct event-symbol pairs\n", global_numPairs > "/dev/stderr"
  printf "%6d distinct events\n", global_numEvents > "/dev/stderr"
  printf "%6d symbols in pairs\n", global_totCt > "/dev/stderr"
  avg_entropy = global_totCtEntropy/global_totCt
  printf "%6.4f average next-symbol entropy\n", avg_entropy > "/dev/stderr"
}

function totalize_event_entropy(   i, entropy)
{
  if (curEvent != "")
    { for (i = 0; i < curEvent_nSymb; i++) { curEvent_SymbCt[i] /= curEvent_TotCt; }
      entropy = 0;
      for (i = 0; i < curEvent_nSymb; i++)
        { if (curEvent_SymbCt[i] > 0) 
            { entropy += -curEvent_SymbCt[i]*log(curEvent_SymbCt[i]); }
        }
      entropy /= log(2.0);
      printf "%7d %7d %8.4f %s\n", curEvent_nSymb, curEvent_TotCt, entropy, curEvent;
      global_numEvents += 1
      global_totCt += curEvent_TotCt
      global_totCtEntropy += curEvent_TotCt*entropy
    }
}

function data_error(msg)
{ 
  printf "line %d: %s\n", NR, msg >> "/dev/stderr";
  abort = 1;
  exit 1;
}

function arg_error(msg)
{ 
  printf "%s\n", msg >> "/dev/stderr";
  abort = 1;
  exit 1;
}