/*
 * nnmontage.c - render a montage of classified test digits as a GIF (C89).
 *
 * Each cell is one test image with a 1-pixel border: green if the network got
 * it right, red if it got it wrong. This is visual evidence of a real error
 * rate, produced on the machine. Reuses the fractal project's GIF/LZW encoder.
 *
 * Usage: nnmontage [testfile] [count] [out.gif]
 *
 * DeepSeek (deepseek-v4.1-flash), October 2026.
 */
#include <stdio.h>
#include <stdlib.h>
#include "nn.h"
#include "gifenc.h"

#define CELL 29
#define IMG  28
#define COLS 20

#define IDX_BG   0
#define IDX_GRAY 16
#define IDX_OK   200
#define IDX_BAD  201

static void build_palette(unsigned char *pal)
{
    int i;
    for (i = 0; i < 256; i++) {
        pal[i * 3] = 10; pal[i * 3 + 1] = 12; pal[i * 3 + 2] = 24;
    }
    for (i = 0; i < 16; i++) {
        unsigned char v = (unsigned char) (i * 17);
        pal[(IDX_GRAY + i) * 3] = v;
        pal[(IDX_GRAY + i) * 3 + 1] = v;
        pal[(IDX_GRAY + i) * 3 + 2] = v;
    }
    pal[IDX_OK * 3] = 0;   pal[IDX_OK * 3 + 1] = 200; pal[IDX_OK * 3 + 2] = 0;
    pal[IDX_BAD * 3] = 200; pal[IDX_BAD * 3 + 1] = 0; pal[IDX_BAD * 3 + 2] = 0;
}

int main(int argc, char **argv)
{
    const char *path = "nn_test.bin";
    const char *outpath = "montage.gif";
    unsigned int count = 200;
    FILE *f, *out;
    unsigned char hdr[4];
    unsigned char *pix, *lab, *rowbuf, *ok, *pal;
    unsigned int n, i, cols, rows, W, H, x, y;
    long correct = 0;
    GifWriter g;

    if (argc > 1) path = argv[1];
    if (argc > 2) count = (unsigned int) atoi(argv[2]);
    if (argc > 3) outpath = argv[3];

    f = fopen(path, "rb");
    if (f == NULL) { perror(path); return 1; }
    if (fread(hdr, 1, 4, f) != 4) return 1;
    n = (unsigned) hdr[0] | ((unsigned) hdr[1] << 8) |
        ((unsigned) hdr[2] << 16) | ((unsigned) hdr[3] << 24);

    pix = (unsigned char *) malloc((size_t) n * 784);
    lab = (unsigned char *) malloc((size_t) n);
    ok = (unsigned char *) malloc((size_t) n);
    if (!pix || !lab || !ok) return 1;
    if (fread(pix, 1, (size_t) n * 784, f) != (size_t) n * 784) return 1;
    if (fread(lab, 1, (size_t) n, f) != (size_t) n) return 1;
    fclose(f);

    if (count > n) count = n;
    if (count == 0) count = n;

    for (i = 0; i < count; i++) {
        ok[i] = (unsigned char) (nn_classify(pix + (size_t) i * 784) == lab[i]);
        correct += ok[i];
    }

    cols = COLS;
    rows = (count + cols - 1) / cols;
    W = cols * CELL;
    H = rows * CELL;

    pal = (unsigned char *) malloc(256 * 3);
    rowbuf = (unsigned char *) malloc(W);
    if (!pal || !rowbuf) return 1;
    build_palette(pal);

    out = fopen(outpath, "wb");
    if (out == NULL) { perror(outpath); return 1; }
    if (!gif_begin(&g, out, (int) W, (int) H, pal)) return 1;

    for (y = 0; y < H; y++) {
        unsigned int r = y / CELL, ly = y % CELL;
        for (x = 0; x < W; x++) {
            unsigned int c = x / CELL, lx = x % CELL;
            unsigned int idx = c + r * cols;
            unsigned char v;
            if (lx == 0 || ly == 0) {
                rowbuf[x] = (idx < count) ? (ok[idx] ? IDX_OK : IDX_BAD) : IDX_BG;
            } else if (idx < count && lx <= IMG && ly <= IMG) {
                v = pix[(size_t) idx * 784 + (ly - 1) * 28 + (lx - 1)];
                rowbuf[x] = (unsigned char) (IDX_GRAY + (v >> 4));
            } else {
                rowbuf[x] = IDX_BG;
            }
        }
        if (!gif_write_row(&g, rowbuf)) return 1;
    }
    if (!gif_end(&g)) return 1;
    fclose(out);

    printf("montage %ux%u: %u digits, %ld correct (%.1f%%), %ld wrong\n",
           W, H, count, correct, 100.0 * (double) correct / (double) count,
           (long) count - correct);
    printf("wrote %s\n", outpath);
    return 0;
}