/**
 * matmul_test - Standalone test example for matmul_axi_slave IP
 *
 * Tests the 4x4 matrix multiplier IP at 0x1f500000 (axiOut_7).
 * Computes C = A x B using register-level access, then verifies
 * against a software reference matmul.
 *
 * Register map (byte offsets from MATMUL_BASE):
 *   0x000  CTRL       [W] bit0=1 starts computation
 *   0x004  STATUS     [R] bit2=error, bit1=done, bit0=busy
 *   0x008  SRC_BASE   [R/W] source data base address
 *   0x00c  DST_BASE   [R/W] result data base address
 *   0x010  GROUP_NUM  [R/W] number of matrix groups
 *   0x020  A_BASE     [R/W] 16 words of matrix A (row-major)
 *   0x060  B_BASE     [R/W] 16 words of matrix B (row-major)
 *   0x0a0  C_BASE     [R]   48 words of result C (3 words per element)
 */

#include <stdio.h>
#include <stdint.h>

/* ─── BSP globals (required by common.mk / picolibc) ─────────────── */
unsigned long UART_BASE             = 0xbf000000;
unsigned long CONFREG_TIMER_BASE    = 0xbf20f100;
unsigned long CONFREG_CLOCKS_PER_SEC = 50000000L;
unsigned long CORE_CLOCKS_PER_SEC   = 33000000L;

/* ─── MATMUL IP register map ─────────────────────────────────────── */
#define MATMUL_BASE        0x1f500000

#define MATMUL_CTRL        (MATMUL_BASE + 0x000)
#define MATMUL_STATUS      (MATMUL_BASE + 0x004)
#define MATMUL_SRC_BASE    (MATMUL_BASE + 0x008)
#define MATMUL_DST_BASE    (MATMUL_BASE + 0x00c)
#define MATMUL_GROUP_NUM   (MATMUL_BASE + 0x010)
#define MATMUL_A_BASE      (MATMUL_BASE + 0x020)
#define MATMUL_B_BASE      (MATMUL_BASE + 0x060)
#define MATMUL_C_BASE      (MATMUL_BASE + 0x0a0)

/* CTRL bits */
#define CTRL_START          0x1

/* STATUS bits */
#define STATUS_BUSY         0x1
#define STATUS_DONE         0x2
#define STATUS_ERROR        0x4

/* ─── Convenience macros ─────────────────────────────────────────── */
#define REG_WR(addr, val)  (*(volatile uint32_t *)(addr) = (uint32_t)(val))
#define REG_RD(addr)       (*(volatile uint32_t *)(addr))

/* ─── Test data ─────────────────────────────────────────────────────
 * Test group 0:
 *   A = identity (4x4)
 *   B = diagonal {2, 3, 4, 5}
 *   Expected C = B (since identity * X = X)
 *
 * Test group 1:
 *   A = [[1,2,0,0], [0,0,3,0], [0,0,0,4], [0,0,0,0]]
 *   B = [[5,0,0,0], [6,0,0,0], [0,7,0,0], [0,0,8,0]]
 *   Expected C computed by software reference below
 *
 * Each A/B is 16 uint32_t, row-major. Unused slots in sub-4x4 matrices
 * are 0 as specified.
 */

#define N_GROUPS  2

static const uint32_t test_a[N_GROUPS][16] = {
    /* group 0: identity */
    { 1, 0, 0, 0,
      0, 1, 0, 0,
      0, 0, 1, 0,
      0, 0, 0, 1 },
    /* group 1 */
    { 1, 2, 0, 0,
      0, 0, 3, 0,
      0, 0, 0, 4,
      0, 0, 0, 0 },
};

static const uint32_t test_b[N_GROUPS][16] = {
    /* group 0: diagonal */
    { 2, 0, 0, 0,
      0, 3, 0, 0,
      0, 0, 4, 0,
      0, 0, 0, 5 },
    /* group 1 */
    { 5, 0, 0, 0,
      6, 0, 0, 0,
      0, 7, 0, 0,
      0, 0, 8, 0 },
};

/* ─── Software reference matmul ─────────────────────────────────────
 * Returns the 66-bit result element as two 32-bit halves:
 *   words[0] = result[31:0]
 *   words[1] = result[63:32]
 *   words[2] = result[65:64]  (only 2 bits, zero-extended to 32)
 */
static void matmul_ref_u32(
    const uint32_t a[16],
    const uint32_t b[16],
    uint32_t       c_words[48])
{
    int row, col, k;
    for (row = 0; row < 4; row++) {
        for (col = 0; col < 4; col++) {
            uint64_t acc = 0;
            for (k = 0; k < 4; k++) {
                acc += (uint64_t)a[row * 4 + k] * (uint64_t)b[k * 4 + col];
            }
            int idx = (row * 4 + col) * 3;
            c_words[idx + 0] = (uint32_t)(acc & 0xFFFFFFFFULL);
            c_words[idx + 1] = (uint32_t)((acc >> 32) & 0xFFFFFFFFULL);
            c_words[idx + 2] = (uint32_t)((acc >> 64) & 0x3ULL);
        }
    }
}

/* ─── Write 16 words of matrix data to the IP ────────────────────── */
static void write_matrix(uint32_t base_addr, const uint32_t data[16])
{
    int i;
    for (i = 0; i < 16; i++) {
        REG_WR(base_addr + i * 4, data[i]);
    }
}

/* ─── Read 48 words of result C from the IP ──────────────────────── */
static void read_result(uint32_t c_words[48])
{
    int i;
    for (i = 0; i < 48; i++) {
        c_words[i] = REG_RD(MATMUL_C_BASE + i * 4);
    }
}

/* ─── Wait until IP is done, return 0 on success, -1 on error ────── */
static int wait_done(void)
{
    while (1) {
        uint32_t status = REG_RD(MATMUL_STATUS);
        if (status & STATUS_ERROR) {
            printf("[ERROR] matmul IP reports error (status=0x%08lx)\n",
                   (unsigned long)status);
            return -1;
        }
        if (status & STATUS_DONE) {
            return 0;
        }
    }
}

/* ─── Verify group result against expected C words ───────────────── */
static int verify_group(int group, const uint32_t got[48],
                        const uint32_t expected[48])
{
    int i, errors = 0;
    for (i = 0; i < 48; i++) {
        if (got[i] != expected[i]) {
            if (errors < 8) {
                int elem = i / 3;
                int sub  = i % 3; /* 0=low, 1=high, 2=overflow */
                const char *sub_name[] = {"low", "high", "ovf"};
                printf("  group%d C[%d][%d].%s: got 0x%08lx exp 0x%08lx\n",
                       group, elem / 4, elem % 4, sub_name[sub],
                       (unsigned long)got[i], (unsigned long)expected[i]);
            }
            errors++;
        }
    }
    if (errors > 8) {
        printf("  ... and %d more mismatches\n", errors - 8);
    }
    return errors;
}

/* ─── Main ────────────────────────────────────────────────────────── */
int main(void)
{
    int g, pass_count = 0, fail_count = 0;

    printf("\n===== MATMUL IP Test =====\n");
    printf("Base addr  : 0x%08lx\n", (unsigned long)MATMUL_BASE);
    printf("Groups     : %d\n", N_GROUPS);
    printf("===========================\n");

    /* Write group count */
    REG_WR(MATMUL_GROUP_NUM, N_GROUPS);

    /* For this register-level example, SRC_BASE / DST_BASE are
     * informational only (no DMA).  Set them to 0. */
    REG_WR(MATMUL_SRC_BASE, 0);
    REG_WR(MATMUL_DST_BASE, 0);

    /* ── Run each test group ──────────────────────────────────── */
    for (g = 0; g < N_GROUPS; g++) {
        uint32_t expected_c[48];
        uint32_t got_c[48];
        int errs;

        printf("\n--- Group %d ---\n", g);
        printf("Loading A[16] and B[16] into IP...\n");

        write_matrix(MATMUL_A_BASE, test_a[g]);
        write_matrix(MATMUL_B_BASE, test_b[g]);

        /* Clear any stale done/error flags */
        REG_WR(MATMUL_STATUS, STATUS_DONE | STATUS_ERROR);

        printf("Starting computation (CTRL=0x%x)...\n", CTRL_START);
        REG_WR(MATMUL_CTRL, CTRL_START);

        printf("Waiting for done...\n");
        if (wait_done() != 0) {
            fail_count++;
            continue;
        }

        printf("Reading C[48] from IP...\n");
        read_result(got_c);

        printf("Computing software reference...\n");
        matmul_ref_u32(test_a[g], test_b[g], expected_c);

        printf("Verifying...\n");
        errs = verify_group(g, got_c, expected_c);

        if (errs == 0) {
            printf("Group %d: PASS\n", g);
            pass_count++;
        } else {
            printf("Group %d: FAIL (%d mismatches)\n", g, errs);
            fail_count++;
        }
    }

    /* ── Final report ─────────────────────────────────────────── */
    printf("\n============================\n");
    printf("RESULTS: %d passed, %d failed\n", pass_count, fail_count);
    if (fail_count == 0) {
        printf("ALL TESTS PASSED\n");
    } else {
        printf("SOME TESTS FAILED\n");
    }
    printf("============================\n");

    /* Signal completion to testbench via UART */
    printf("MATMUL_DONE\n");

    return (fail_count == 0) ? 0 : 1;
}