let words = [...new Set(phrase.split(" "))];

let bigramCounts = Array(words.length)

  .fill()

  .map(() => Array(words.length).fill(0));

let currentIndex = 0;

let bigram = "";

let fontSize = 24;

let isAnimationFinished = false;

let wordCounts = Array(words.length).fill(0);

​

function setup() {

  createCanvas(1000, 700);

  textAlign(CENTER, CENTER);

  textSize(fontSize);

  textFont("Arial, Helvetica, sans-serif");

}

​

function draw() {

  background(255);

​

  // Display the title

  textSize(fontSize * 1.2);

  fill(0);

  text("Bigram Learning Animation", width / 2, 50);

​

  // Display the matrix

  let matrixX = 100;

  let matrixY = height / 2 - (words.length * 40) / 2;

  let cellSize = 50;

  textSize(fontSize * 0.6);

​

  // Display matrix column headers

  fill(0);

  for (let i = 0; i < words.length; i++) {

    text(words[i], matrixX + (i + 1) * cellSize + cellSize / 2, matrixY - 30);

  }

​

  // Display matrix row headers

  for (let i = 0; i < words.length; i++) {

    push();

    translate(matrixX - 30, matrixY + (i + 1) * cellSize - cellSize / 2);

    rotate(-PI / 2);

    text(words[i], 0, 0);

    pop();

  }

​

  // Display matrix cells

  for (let i = 0; i < words.length; i++) {

    for (let j = 0; j < words.length; j++) {

      let count = bigramCounts[i][j];

      let probability = wordCounts[i] > 0 ? count / wordCounts[i] : 0;

      fill(255);

      stroke(0);

      rect(

        matrixX + (j + 1) * cellSize,

        matrixY + i * cellSize,

        cellSize,

        cellSize

      );

      fill(0);

      text(

        probability.toFixed(2),

        matrixX + (j + 1) * cellSize + cellSize / 2,

        matrixY + i * cellSize + cellSize / 2

      );

    }

  }

​

  // Display the words passing on the right

  let wordsX = matrixX + (words.length + 2) * cellSize;

  let wordsY = height / 2 - 50;

  textSize(fontSize);

  fill(0);

  text(

    phrase

      .split(" ")

      .slice(0, currentIndex + 1)

      .join(" "),

    wordsX,

    wordsY,

    300,

    100

  );

​

  // Update the bigram every 1 second

  if (frameCount % 60 === 0) {

    if (currentIndex < phrase.split(" ").length - 1) {

      bigram =

        phrase.split(" ")[currentIndex] +

        " " +

        phrase.split(" ")[currentIndex + 1];

      let row = words.indexOf(phrase.split(" ")[currentIndex]);

      let col = words.indexOf(phrase.split(" ")[currentIndex + 1]);

      bigramCounts[row][col]++;

      wordCounts[row]++;

      currentIndex++;

    } else {

      isAnimationFinished = true;

      noLoop(); // Stop the animation when reaching the end

    }

  }

  if (isAnimationFinished) {

    fill(0);

    textSize(fontSize);

    text("Animation finished. Click to restart.", width / 2, height - 50);

  }

  // Display the current bigram below the words passing

  fill(0);

  textSize(fontSize * 0.8);

  text("Current Bigram: " + bigram, wordsX + 150, wordsY + 200);

}

​

function resetAnimation() {

  currentIndex = 0;

  bigram = "";

  bigramCounts = Array(words.length)

    .fill()

    .map(() => Array(words.length).fill(0));

  isAnimationFinished = false;

  wordCounts = Array(words.length).fill(0);

  loop();

}

function mousePressed() {

  if (isAnimationFinished) {

    resetAnimation();

  }

}