#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

// we estimate the probability of a key event of interest given the last few
// key events. we do so by searching the file for occurrences of the last few
// key events and counting how often the key event of interest follows. how many
// is "the last few key events"? it's the longest sequence that occurs at least
// `MIN_MARGINALS` times in the file. then, how large shoud `MIN_MARGINALS`
// be? it's positively correlated with entropy, so we can minimize entropy by
// setting it to 1, in which case entropy goes to 0. well that's not useful.
// this happens because "the last few key events" becomes every key event since
// the beginning of the file, so we overfit to the file and can just look up
// what the next key event ought to be. we can fix this by hiding the key event
// of interest from the file when searching for other occurrences of it. this
// way the test data is not part of the training data. now this brings about
// another difficulty: what if we've identified our sequence of last few key
// events that occurs at least `MIN_MARGINALS` times, but nowhere else in the
// file is that sequence followed by the key event of interest? the calculated
// probability of that key event given the last few key events is zero, thus
// surprisal and entropy both go to infinity. the fix here is straightforward:
// we shift and scale the probability distribution so no event can occur with
// probability less than `MIN_PROB`. this finally ties the loose ends. now the
// optimization of entropy as a function of `MIN_MARGINALS` and `MIN_PROB` gives
// meaningful output. their values below were determined by manual gradient
// descent over numerical differentiation and are a local minimum of entropy.
#define MIN_MARGINALS 16
#define MIN_PROB 0.0008

// wide view on surprisal
// #define SURPRISAL_LO 0.0 // bits
// #define SURPRISAL_HI 6.0 // bits

// focus on high surprisal
// #define SURPRISAL_LO 4.0 // bits
// #define SURPRISAL_HI 8.0 // bits

// focus on low surprisal
#define SURPRISAL_LO 0.0 // bits
#define SURPRISAL_HI 1.5 // bits

#define ALPHABET_SIZE 128 // approximate number of possible key events
#define TAIL_LEN 1024     // number of bytes for the tail entropy printout

int main(void) {
  size_t len = 0, cap = 256;
  char *buf = malloc(cap);
  while (len += fread(buf + len, 1, cap - len, stdin), len == cap)
    buf = realloc(buf, cap *= 2);
  if (!feof(stdin))
    perror("fread"), exit(EXIT_FAILURE);
  buf = realloc(buf, len);

  double information = 0.0, tail_information = 0.0;

  size_t length = 0; // length of the sequence of last few key events
  for (char *p = buf; p < buf + len; p++) {
    size_t marginals, joints;
    do {
      marginals = joints = -1; // hide the key event of interest
      for (char *q = buf + length; q < buf + len; q++)
        if (!memcmp(p - length, q - length, length))
          marginals++, joints += *p == *q;
    } while (marginals < MIN_MARGINALS && length && length--);
    // this key event's joint probability is the next key event's marginal
    // probability, so when it's less than `MIN_MARGINALS` we can decrement
    // the length early
    length++, length -= joints < MIN_MARGINALS;

    double probability = 1.0 * joints / marginals;
    // scale the distribution by `1 - MIN_PROB * ALPHABET_SIZE` so it integrates
    // to that. now shift it up by `MIN_PROB` and it will integrate to 1 again.
    double softened = probability * (1 - MIN_PROB * ALPHABET_SIZE) + MIN_PROB;
    double surprisal = log2(1.0 / softened);
    double clamped = fmax(SURPRISAL_LO, fmin(SURPRISAL_HI, surprisal));
    double lerped = (clamped - SURPRISAL_LO) / (SURPRISAL_HI - SURPRISAL_LO);
    int color = 232 + (int)(15.999 * lerped); // colors 232--247

    fprintf(stdout, "\033[48;5;%dm%c", color, *p);

    information += surprisal;
    if (p >= buf + len - TAIL_LEN)
      tail_information += surprisal;
  }

  double entropy = information / len,
         tail_entropy = tail_information / TAIL_LEN;
  fprintf(stderr, "%f\t%zu\t%f\t%f\n", information, len, entropy, tail_entropy);
}