/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
 * example.c — minimal, self-contained SR4 usage demo (no MNIST, no BLAS).
 *
 * Synthetic task: recover a fixed random TERNARY "teacher" weight matrix W*
 * (rows x cols) from noiseless linear regression data. A batch of random inputs
 * X (batch x rows) is mapped by the teacher to targets  Y* = X @ W*. SR4 trains
 * its own ternary W to minimise the mean-squared error  ||X @ W - Y*||^2.
 *
 * It exercises the full library cycle every step:
 *     init  ->  ( materialize -> forward -> backward -> update )  ->  free
 * and prints the loss falling as SR4's ternary weights converge on the teacher.
 *
 * Build:   make            (produces ./example)
 * Run:     ./example
 */

#include "ternary_optim.h"

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

#define ROWS   16      /* input dim  (matrix rows) */
#define COLS   4       /* output dim (matrix cols) */
#define BATCH  64
#define STEPS  4000

/* tiny local xorshift so the demo is self-contained */
static uint64_t g_rng = 12345;
static float frand(void) {
    uint64_t x = g_rng; x ^= x << 13; x ^= x >> 7; x ^= x << 17; g_rng = x;
    return (float)(x >> 11) / (float)(1ULL << 53);
}

int main(void) {
    const int numel = ROWS * COLS;

    /* Fixed ternary teacher W* in {-1,0,+1}. */
    float *W_star = malloc(numel * sizeof(float));
    for (int i = 0; i < numel; i++) {
        float r = frand();
        W_star[i] = (r < 0.33f) ? -1.0f : (r < 0.66f) ? 0.0f : 1.0f;
    }

    /* Create the optimizer for a (ROWS x COLS) weight matrix. */
    ternary_optim_config cfg;
    ternary_optim_config_default(&cfg);
    cfg.seed = 7;
    ternary_optim *opt = ternary_optim_create(ROWS, COLS, &cfg);
    if (!opt) { fprintf(stderr, "create failed\n"); return 1; }

    /* Caller-owned scratch buffers (the loop owns these, not the optimizer). */
    float *W        = malloc(numel * sizeof(float));          /* materialized  */
    float *grad_sum = malloc(numel * sizeof(float));          /* dL/dW summed  */
    float *X        = malloc(BATCH * ROWS * sizeof(float));   /* batch inputs  */
    float *pred     = malloc(BATCH * COLS * sizeof(float));
    float *target   = malloc(BATCH * COLS * sizeof(float));
    float *err      = malloc(BATCH * COLS * sizeof(float));

    printf("SR4 example: recover a %dx%d ternary teacher via linear regression\n",
           ROWS, COLS);
    printf("step |   mse\n");
    printf("-----+---------\n");

    for (uint64_t step = 0; step < STEPS; step++) {
        /* ---- materialize the ternary weights for this forward pass ---- */
        ternary_optim_materialize(opt, W);

        /* ---- draw a fresh batch and build targets from the teacher ---- */
        for (int s = 0; s < BATCH; s++)
            for (int r = 0; r < ROWS; r++)
                X[s * ROWS + r] = frand() * 2.0f - 1.0f;
        for (int s = 0; s < BATCH; s++)
            for (int c = 0; c < COLS; c++) {
                float acc_star = 0.0f, acc_pred = 0.0f;
                for (int r = 0; r < ROWS; r++) {
                    acc_star += X[s * ROWS + r] * W_star[r * COLS + c];
                    acc_pred += X[s * ROWS + r] * W[r * COLS + c];
                }
                target[s * COLS + c] = acc_star;   /* forward (teacher)  */
                pred[s * COLS + c]   = acc_pred;    /* forward (student)  */
            }

        /* ---- loss + backward: err = pred - target,  grad = X^T @ err ---- */
        double mse = 0.0;
        for (int i = 0; i < BATCH * COLS; i++) {
            err[i] = pred[i] - target[i];
            mse += (double)err[i] * err[i];
        }
        mse /= (BATCH * COLS);

        for (int i = 0; i < numel; i++) grad_sum[i] = 0.0f;
        for (int s = 0; s < BATCH; s++)
            for (int r = 0; r < ROWS; r++)
                for (int c = 0; c < COLS; c++)
                    grad_sum[r * COLS + c] += X[s * ROWS + r] * err[s * COLS + c];

        /* ---- SR4 update: integrate the batch gradient into q ---- */
        ternary_optim_step(opt, grad_sum, BATCH, step, -1.0f /* cfg default */);

        if (step % 400 == 0 || step == STEPS - 1)
            printf("%4llu | %.5f\n", (unsigned long long)step, mse);
    }

    /* Report how many ternary weights match the teacher exactly. */
    ternary_optim_materialize(opt, W);
    int match = 0;
    for (int i = 0; i < numel; i++) if (W[i] == W_star[i]) match++;
    printf("\nrecovered %d / %d teacher weights exactly\n", match, numel);

    ternary_optim_free(opt);
    free(W_star); free(W); free(grad_sum); free(X); free(pred); free(target); free(err);
    return 0;
}