utensil

#include "utensil.h"

#include "speech_commands_loader.h"

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

// Keyword spotting on Google Speech Commands: a single-layer LSTM reads a
// clip's log-mel spectrogram one frame at a time, and a linear head classifies
// the final hidden state into one of 12 classes (10 core command words plus
// "_unknown_" and "_silence_"). This is the LSTM-primitive analogue of
// mnist_mlp.c -- a small, complete reference to build a real keyword spotter
// (e.g. microwakeword-style single-wakeword detection) from.

#define HIDDEN 64
#define BATCH 32

// x_seq is [T,B,NMEL]; a dataset's samples are stored [n,T,NMEL], so a batch
// has to be transposed on the way in.
static ut_tensor* make_batch(sc_dataset_t* d, const int* idx, int off, int B, int* labels) {
  float* bx = malloc((size_t)SC_T * B * SC_NMEL * sizeof(float));
  for (int i = 0; i < B; i++) {
    int ii = idx[off + i];
    labels[i] = d->labels[ii];
    const float* src = d->feats + (size_t)ii * SC_T * SC_NMEL;
    for (int t = 0; t < SC_T; t++)
      memcpy(bx + ((size_t)t * B + i) * SC_NMEL, src + (size_t)t * SC_NMEL,
             SC_NMEL * sizeof(float));
  }
  ut_tensor* x = ut_from_data(3, (int[]){SC_T, B, SC_NMEL}, bx, UT_CPU);
  free(bx);
  return x;
}

// runs the LSTM over the whole sequence and classifies the final hidden
// state; cache may be NULL for an inference-only pass (see ut_lstm_forward_seq)
static ut_tensor* forward(ut_lstm* lstm, ut_linear* fc, ut_tensor* x, ut_lstm_seq_cache* cache,
                          ut_tensor** h_n_out) {
  ut_tensor *h_n, *c_n;
  ut_tensor* h_seq = ut_lstm_forward_seq(lstm, x, NULL, NULL, cache, &h_n, &c_n);
  ut_free_all(h_seq, c_n);
  ut_tensor* logits = ut_linear_forward(fc, h_n);
  if (h_n_out)
    *h_n_out = h_n;
  else
    ut_free(h_n);
  return logits;
}

static int eval_split(ut_lstm* lstm, ut_linear* fc, sc_dataset_t* d, int cm[SC_NCLASS][SC_NCLASS]) {
  int ok = 0;
  for (int i = 0; i < d->n; i += BATCH) {
    int nb = (i + BATCH <= d->n) ? BATCH : (d->n - i);
    int all_idx[BATCH];
    for (int j = 0; j < nb; j++) all_idx[j] = i + j;
    int bl[BATCH];
    ut_tensor* x = make_batch(d, all_idx, 0, nb, bl);
    ut_tensor* logits = forward(lstm, fc, x, NULL, NULL);
    ut_sync_cpu(logits);
    for (int j = 0; j < nb; j++) {
      int pred = sc_argmax(logits, j);
      if (pred == bl[j]) ok++;
      if (cm) cm[bl[j]][pred]++;
    }
    ut_free_all(x, logits);
  }
  return ok;
}

int main(void) {
  srand(42);
  ut_dev dev = UT_CPU;  // LSTM gate math is CPU-only; see utensil.h

  printf("Loading Speech Commands (this scans the dataset once per split)...\n");
  sc_dataset_t train = sc_load_split("speech_commands", SC_TRAIN);
  sc_dataset_t val = sc_load_split("speech_commands", SC_VAL);
  sc_dataset_t test = sc_load_split("speech_commands", SC_TEST);
  printf("train: %d  val: %d  test: %d\n\n", train.n, val.n, test.n);

  ut_lstm lstm = ut_lstm_alloc(SC_NMEL, HIDDEN, dev);
  ut_linear fc = ut_linear_alloc(HIDDEN, SC_NCLASS, true, dev);
  ut_adam opt = ut_adam_alloc((ut_tensor*[]){lstm.W_ih, lstm.W_hh, lstm.bias, fc.weight, fc.bias},
                              5, 1e-3f, 0.9f, 0.999f, 1e-8f, 0.f);

  int B = BATCH, epochs = 15, batches = train.n / B;
  ut_tensor* grad_logits = ut_alloc(2, (int[]){B, SC_NCLASS}, dev);
  int* idx = malloc((size_t)train.n * sizeof(int));
  for (int i = 0; i < train.n; i++) idx[i] = i;

  for (int ep = 0; ep < epochs; ep++) {
    sc_shuffle(idx, train.n);
    clock_t t0 = clock();
    float loss_sum = 0;
    int correct = 0;

    for (int bi = 0; bi < batches; bi++) {
      int bl[BATCH];
      ut_tensor* x = make_batch(&train, idx, bi * B, B, bl);

      ut_lstm_seq_cache cache;
      ut_tensor* h_n;
      ut_tensor* logits = forward(&lstm, &fc, x, &cache, &h_n);

      float loss = ut_cross_entropy(logits, bl, grad_logits);
      ut_sync_cpu(logits);
      for (int i = 0; i < B; i++)
        if (sc_argmax(logits, i) == bl[i]) correct++;

      // grad w.r.t. h flows in only at the last timestep -- every other
      // slice of grad_h_seq stays zero.
      ut_tensor* dh_n = ut_linear_backward(&fc, h_n, grad_logits, opt.grads[3], opt.grads[4]);
      ut_sync_cpu(dh_n);
      ut_tensor* grad_h_seq = ut_alloc(3, (int[]){SC_T, B, HIDDEN}, dev);
      memcpy(grad_h_seq->data + (size_t)(SC_T - 1) * B * HIDDEN, dh_n->data,
             (size_t)B * HIDDEN * sizeof(float));
      ut_tensor* dx_seq =
          ut_lstm_backward_seq(&lstm, &cache, grad_h_seq, opt.grads[0], opt.grads[1], opt.grads[2]);

      ut_adam_step(&opt, 5.0f);
      loss_sum += loss;

      ut_lstm_seq_cache_free(&cache);
      ut_free_all(x, logits, h_n, dh_n, grad_h_seq, dx_seq);
    }

    float secs = (float)(clock() - t0) / (float)CLOCKS_PER_SEC;
    int val_ok = eval_split(&lstm, &fc, &val, NULL);

    printf("epoch %2d  loss %7.4f  train %5.1f%%  val %5.1f%%  %5.1fs\n", ep + 1,
           loss_sum / (float)batches, 100.f * (float)correct / (float)(batches * B),
           100.f * (float)val_ok / (float)val.n, secs);
  }

  int cm[SC_NCLASS][SC_NCLASS] = {0};
  int test_ok = eval_split(&lstm, &fc, &test, cm);
  printf("\ntest accuracy: %5.1f%%\n\n", 100.f * (float)test_ok / (float)test.n);

  printf("Confusion matrix:\n%14s", "");
  for (int j = 0; j < SC_NCLASS; j++) printf("%8s", SC_CLASSES[j]);
  printf("\n");
  for (int r = 0; r < SC_NCLASS; r++) {
    printf("%14s", SC_CLASSES[r]);
    for (int c = 0; c < SC_NCLASS; c++) printf("%8d", cm[r][c]);
    printf("\n");
  }

  free(idx);
  ut_free(grad_logits);
  ut_adam_free(&opt);
  ut_lstm_free(&lstm);
  ut_linear_free(&fc);
  sc_dataset_free(&train);
  sc_dataset_free(&val);
  sc_dataset_free(&test);
  return 0;
}