utensil

#include "utensil.h"

#include "cifar10_loader.h"

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

// Standard CIFAR-10 ResNet-8 (He et al.'s CIFAR ResNet family, n=1): a stem
// conv, then 3 residual stages (16/32/64 channels, stride 1/2/2) of one basic
// block each, GAP, and a linear head. Stride/channel-changing blocks use a 1x1
// conv+BN projection shortcut. No data augmentation, so expect lower accuracy
// than a fully-tuned ResNet-8 — this is a scaffold for comparing against the
// same architecture in PyTorch, not a from-scratch SOTA run.

// gradient tensors are borrowed from the optimizer's own grads[]; wired once
// in main() after ut_adam_alloc, then read directly during backward.
typedef struct {
  ut_tensor *c1w, *c1b, *b1w, *b1b, *c2w, *c2b, *b2w, *b2b, *pw, *pb, *pbnw, *pbnb;
} block_grads_t;

typedef struct {
  ut_conv2d conv1, conv2, proj;
  ut_batchnorm2d bn1, bn2, projbn;
  bool has_proj;
  block_grads_t g;
} block_t;

typedef struct {
  ut_conv2d_cache c1, c2, cp;
  ut_batchnorm2d_cache b1, b2, bp;
  ut_tensor *h1bn, *sum;  // kept for backward's two ReLUs
} block_cache_t;

typedef struct {
  ut_conv2d stem_conv;
  ut_batchnorm2d stem_bn;
  block_t blocks[3];
  ut_linear fc;
  ut_tensor *stem_cw, *stem_cb, *stem_bw, *stem_bb, *fc_w, *fc_b;  // borrowed, like block_grads_t
} resnet8_t;

typedef struct {
  ut_conv2d_cache stem_c;
  ut_batchnorm2d_cache stem_bn_c;
  ut_tensor* stem_relu_in;
  block_cache_t block_c[3];
  ut_tensor* pooled;  // kept for the fc layer's backward
  int pool_h, pool_w;
} resnet8_cache_t;

static block_t block_alloc(int in_c, int out_c, int stride, ut_dev dev) {
  block_t b = {0};
  b.conv1 = ut_conv2d_alloc(in_c, out_c, 3, 3, stride, 1, true, dev);
  b.bn1 = ut_batchnorm2d_alloc(out_c, dev);
  b.conv2 = ut_conv2d_alloc(out_c, out_c, 3, 3, 1, 1, true, dev);
  b.bn2 = ut_batchnorm2d_alloc(out_c, dev);
  b.has_proj = in_c != out_c || stride != 1;
  if (b.has_proj) {
    b.proj = ut_conv2d_alloc(in_c, out_c, 1, 1, stride, 0, true, dev);
    b.projbn = ut_batchnorm2d_alloc(out_c, dev);
  }
  return b;
}

static resnet8_t resnet8_alloc(ut_dev dev) {
  resnet8_t m = {0};
  m.stem_conv = ut_conv2d_alloc(3, 16, 3, 3, 1, 1, true, dev);
  m.stem_bn = ut_batchnorm2d_alloc(16, dev);
  m.blocks[0] = block_alloc(16, 16, 1, dev);
  m.blocks[1] = block_alloc(16, 32, 2, dev);
  m.blocks[2] = block_alloc(32, 64, 2, dev);
  m.fc = ut_linear_alloc(64, 10, true, dev);
  return m;
}

// appends this block's param tensors to out[] (for ut_adam_alloc) and returns the count
static int block_params(block_t* b, ut_tensor** out) {
  int n = 0;
  out[n++] = b->conv1.weight, out[n++] = b->conv1.bias;
  out[n++] = b->bn1.weight, out[n++] = b->bn1.bias;
  out[n++] = b->conv2.weight, out[n++] = b->conv2.bias;
  out[n++] = b->bn2.weight, out[n++] = b->bn2.bias;
  if (b->has_proj) {
    out[n++] = b->proj.weight, out[n++] = b->proj.bias;
    out[n++] = b->projbn.weight, out[n++] = b->projbn.bias;
  }
  return n;
}

// mirrors block_params' exact order to wire each gradient pointer to its param
static void block_wire_grads(block_t* b, ut_tensor** g) {
  int n = 0;
  b->g.c1w = g[n++], b->g.c1b = g[n++];
  b->g.b1w = g[n++], b->g.b1b = g[n++];
  b->g.c2w = g[n++], b->g.c2b = g[n++];
  b->g.b2w = g[n++], b->g.b2b = g[n++];
  if (b->has_proj) {
    b->g.pw = g[n++], b->g.pb = g[n++];
    b->g.pbnw = g[n++], b->g.pbnb = g[n++];
  }
}

static ut_tensor* block_forward(block_t* b, ut_tensor* x, bool training, block_cache_t* c) {
  ut_tensor* h1 = ut_conv2d_forward(&b->conv1, x, c ? &c->c1 : NULL);
  ut_tensor* h1bn = ut_batchnorm2d_forward(&b->bn1, h1, training, c ? &c->b1 : NULL);
  ut_free(h1);
  ut_tensor* h1r = ut_relu(h1bn);
  if (c) c->h1bn = h1bn; else ut_free(h1bn);

  ut_tensor* h2 = ut_conv2d_forward(&b->conv2, h1r, c ? &c->c2 : NULL);
  ut_free(h1r);
  ut_tensor* h2bn = ut_batchnorm2d_forward(&b->bn2, h2, training, c ? &c->b2 : NULL);
  ut_free(h2);

  ut_tensor* shortcut;
  if (b->has_proj) {
    ut_tensor* p = ut_conv2d_forward(&b->proj, x, c ? &c->cp : NULL);
    shortcut = ut_batchnorm2d_forward(&b->projbn, p, training, c ? &c->bp : NULL);
    ut_free(p);
  } else {
    shortcut = ut_retain(x);
  }

  ut_tensor* sum = ut_add(h2bn, shortcut);
  ut_free_all(h2bn, shortcut);
  ut_tensor* out = ut_relu(sum);
  if (c) c->sum = sum; else ut_free(sum);
  return out;
}

static ut_tensor* block_backward(block_t* b, block_cache_t* c, ut_tensor* dout) {
  ut_tensor* dsum = ut_relu_backward(dout, c->sum);
  // sum = h2bn + shortcut -> gradient passes unchanged to both branches
  ut_tensor* dh2 = ut_batchnorm2d_backward(&b->bn2, &c->b2, dsum, b->g.b2w, b->g.b2b);
  ut_tensor* dh1r = ut_conv2d_backward(&b->conv2, &c->c2, dh2, b->g.c2w, b->g.c2b);
  ut_free(dh2);
  ut_tensor* dh1bn = ut_relu_backward(dh1r, c->h1bn);
  ut_free(dh1r);
  ut_tensor* dh1 = ut_batchnorm2d_backward(&b->bn1, &c->b1, dh1bn, b->g.b1w, b->g.b1b);
  ut_free(dh1bn);
  ut_tensor* dx_main = ut_conv2d_backward(&b->conv1, &c->c1, dh1, b->g.c1w, b->g.c1b);
  ut_free(dh1);

  ut_tensor* dx;
  if (b->has_proj) {
    ut_tensor* dp = ut_batchnorm2d_backward(&b->projbn, &c->bp, dsum, b->g.pbnw, b->g.pbnb);
    ut_tensor* dx_proj = ut_conv2d_backward(&b->proj, &c->cp, dp, b->g.pw, b->g.pb);
    ut_free(dp);
    dx = ut_add(dx_main, dx_proj);
    ut_free_all(dx_main, dx_proj);
  } else {
    dx = ut_add(dx_main, dsum);
    ut_free(dx_main);
  }
  ut_free(dsum);
  return dx;
}

static void block_cache_free(block_t* b, block_cache_t* c) {
  ut_conv2d_cache_free(&c->c1);
  ut_conv2d_cache_free(&c->c2);
  ut_batchnorm2d_cache_free(&c->b1);
  ut_batchnorm2d_cache_free(&c->b2);
  if (b->has_proj) {
    ut_conv2d_cache_free(&c->cp);
    ut_batchnorm2d_cache_free(&c->bp);
  }
  ut_free_all(c->h1bn, c->sum);
}

static void block_free(block_t* b) {
  ut_conv2d_free(&b->conv1);
  ut_conv2d_free(&b->conv2);
  ut_batchnorm2d_free(&b->bn1);
  ut_batchnorm2d_free(&b->bn2);
  if (b->has_proj) {
    ut_conv2d_free(&b->proj);
    ut_batchnorm2d_free(&b->projbn);
  }
}

static ut_tensor* resnet8_forward(resnet8_t* m, ut_tensor* x, bool training, resnet8_cache_t* c) {
  ut_tensor* s = ut_conv2d_forward(&m->stem_conv, x, c ? &c->stem_c : NULL);
  ut_tensor* sbn = ut_batchnorm2d_forward(&m->stem_bn, s, training, c ? &c->stem_bn_c : NULL);
  ut_free(s);
  ut_tensor* h = ut_relu(sbn);
  if (c) c->stem_relu_in = sbn; else ut_free(sbn);

  for (int i = 0; i < 3; i++) {
    ut_tensor* next = block_forward(&m->blocks[i], h, training, c ? &c->block_c[i] : NULL);
    ut_free(h);
    h = next;
  }

  int ph = h->shape.shape[2], pw = h->shape.shape[3];
  ut_tensor* pooled = ut_global_avgpool2d(h);
  ut_free(h);
  if (c) c->pooled = pooled, c->pool_h = ph, c->pool_w = pw;
  ut_tensor* logits = ut_linear_forward(&m->fc, pooled);
  if (!c) ut_free(pooled);
  return logits;
}

static ut_tensor* resnet8_backward(resnet8_t* m, resnet8_cache_t* c, ut_tensor* dlogits) {
  ut_tensor* dpooled = ut_linear_backward(&m->fc, c->pooled, dlogits, m->fc_w, m->fc_b);
  ut_tensor* dh = ut_global_avgpool2d_backward(dpooled, c->pool_h, c->pool_w);
  ut_free(dpooled);

  for (int i = 2; i >= 0; i--) {
    ut_tensor* dprev = block_backward(&m->blocks[i], &c->block_c[i], dh);
    ut_free(dh);
    dh = dprev;
  }

  ut_tensor* dsbn = ut_relu_backward(dh, c->stem_relu_in);
  ut_free(dh);
  ut_tensor* ds = ut_batchnorm2d_backward(&m->stem_bn, &c->stem_bn_c, dsbn, m->stem_bw, m->stem_bb);
  ut_free(dsbn);
  ut_tensor* dx = ut_conv2d_backward(&m->stem_conv, &c->stem_c, ds, m->stem_cw, m->stem_cb);
  ut_free(ds);
  return dx;
}

static void resnet8_cache_free(resnet8_t* m, resnet8_cache_t* c) {
  ut_conv2d_cache_free(&c->stem_c);
  ut_batchnorm2d_cache_free(&c->stem_bn_c);
  ut_free(c->stem_relu_in);
  for (int i = 0; i < 3; i++) block_cache_free(&m->blocks[i], &c->block_c[i]);
  ut_free(c->pooled);
}

static void resnet8_free(resnet8_t* m) {
  ut_conv2d_free(&m->stem_conv);
  ut_batchnorm2d_free(&m->stem_bn);
  for (int i = 0; i < 3; i++) block_free(&m->blocks[i]);
  ut_linear_free(&m->fc);
}

int main(void) {
  srand(42);

  ut_dev dev = UT_METAL;  // flip to UT_CPU to compare against a (CPU-only) PyTorch script

  printf("Loading CIFAR-10…\n");
  const char* train_files[5] = {"cifar/data_batch_1.bin", "cifar/data_batch_2.bin",
                                "cifar/data_batch_3.bin", "cifar/data_batch_4.bin",
                                "cifar/data_batch_5.bin"};
  const char* test_files[1] = {"cifar/test_batch.bin"};
  cifar10_t train = cifar10_load(train_files, 5);
  cifar10_t test = cifar10_load(test_files, 1);
  printf("train: %d  test: %d\n\n", train.n, test.n);

  resnet8_t m = resnet8_alloc(dev);

  ut_tensor* params[64];
  int np = 0;
  params[np++] = m.stem_conv.weight, params[np++] = m.stem_conv.bias;
  params[np++] = m.stem_bn.weight, params[np++] = m.stem_bn.bias;
  int block_off[3];
  for (int i = 0; i < 3; i++) {
    block_off[i] = np;
    np += block_params(&m.blocks[i], params + np);
  }
  int fc_off = np;
  params[np++] = m.fc.weight, params[np++] = m.fc.bias;

  ut_adam opt = ut_adam_alloc(params, np, 1e-3f, 0.9f, 0.999f, 1e-8f, 1e-4f);

  // wire each layer's gradient pointers to the optimizer's own grads[], in the
  // same order the params[] above was built
  m.stem_cw = opt.grads[0], m.stem_cb = opt.grads[1];
  m.stem_bw = opt.grads[2], m.stem_bb = opt.grads[3];
  for (int i = 0; i < 3; i++) block_wire_grads(&m.blocks[i], opt.grads + block_off[i]);
  m.fc_w = opt.grads[fc_off], m.fc_b = opt.grads[fc_off + 1];

  int B = 64, epochs = 20, batches = train.n / B;
  ut_tensor* grad_logits = ut_alloc(2, (int[]){B, 10}, 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++) {
    cifar10_shuffle(idx, train.n);
    clock_t t0 = clock();
    float loss_sum = 0;
    int correct = 0;

    for (int bi = 0; bi < batches; bi++) {
      float bx[B * CIFAR10_PIX];
      int bl[B];
      for (int i = 0; i < B; i++) {
        int ii = idx[bi * B + i];
        memcpy(bx + i * CIFAR10_PIX, train.imgs + ii * CIFAR10_PIX, CIFAR10_PIX * sizeof(float));
        bl[i] = train.labels[ii];
      }
      ut_tensor* x = ut_from_data(4, (int[]){B, 3, 32, 32}, bx, dev);

      resnet8_cache_t c;
      ut_tensor* logits = resnet8_forward(&m, x, true, &c);
      float loss = ut_cross_entropy(logits, bl, grad_logits);
      ut_sync_cpu(logits);
      for (int i = 0; i < B; i++)
        if (cifar10_argmax(logits, i) == bl[i]) correct++;

      ut_tensor* dx = resnet8_backward(&m, &c, grad_logits);
      ut_adam_step(&opt, 5.0f);
      loss_sum += loss;

      resnet8_cache_free(&m, &c);
      ut_free_all(x, logits, dx);
    }

    float secs = (float)(clock() - t0) / (float)CLOCKS_PER_SEC;

    int test_ok = 0;
    for (int i = 0; i < test.n; i += B) {
      int nb = (i + B <= test.n) ? B : (test.n - i);
      ut_tensor* tx = ut_from_data(4, (int[]){nb, 3, 32, 32}, test.imgs + i * CIFAR10_PIX, dev);
      ut_tensor* tl = resnet8_forward(&m, tx, false, NULL);
      ut_sync_cpu(tl);
      for (int j = 0; j < nb; j++)
        if (cifar10_argmax(tl, j) == test.labels[i + j]) test_ok++;
      ut_free_all(tx, tl);
    }

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

  free(idx);
  ut_free(grad_logits);
  ut_adam_free(&opt);
  resnet8_free(&m);
  cifar10_free(&train);
  cifar10_free(&test);
  return 0;
}