/*
 * Double-precision n×n matrix multiply, compiled to WASM with emscripten.
 * Row-major C = A * B. Kernels:
 *   - matmul_naive:      triple loop (sanity baseline)
 *   - matmul_blocked:    blocked + SIMD (the "with optimizations" method)
 *   - matmul_blocked_mt: the blocked+SIMD kernel parallelized over rows with
 *                        pthreads (only in the -DMATMUL_MT / -pthread build)
 */
#include <emscripten.h>
#include <string.h>

#ifdef __wasm_simd128__
#include <wasm_simd128.h>
#endif

EMSCRIPTEN_KEEPALIVE
void matmul_naive(const double* a, const double* b, double* c, int n) {
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++) {
      double sum = 0.0;
      for (int k = 0; k < n; k++) sum += a[i * n + k] * b[k * n + j];
      c[i * n + j] = sum;
    }
  }
}

/* Block size tuned for L1/L2 cache residency of three double blocks. */
#define BS 64

/* Blocked+SIMD kernel over rows [i0, i1) of C. Each row range is independent
 * (disjoint output rows), so threads can own disjoint ranges without locking. */
static void blocked_range(const double* a, const double* b, double* c, int n,
                          int i0, int i1) {
  memset(c + (size_t)i0 * n, 0, (size_t)(i1 - i0) * n * sizeof(double));

  for (int jj = 0; jj < n; jj += BS) {
    int jmax = jj + BS < n ? jj + BS : n;
    for (int kk = 0; kk < n; kk += BS) {
      int kmax = kk + BS < n ? kk + BS : n;
      for (int i = i0; i < i1; i++) {
        const double* arow = a + (size_t)i * n;
        double* crow = c + (size_t)i * n;
        for (int k = kk; k < kmax; k++) {
          double aik = arow[k];
          const double* brow = b + (size_t)k * n;
#ifdef __wasm_simd128__
          v128_t vaik = wasm_f64x2_splat(aik);
          int j = jj;
          for (; j + 1 < jmax; j += 2) {
            v128_t vb = wasm_v128_load(brow + j);
            v128_t vc = wasm_v128_load(crow + j);
            vc = wasm_f64x2_add(vc, wasm_f64x2_mul(vaik, vb));
            wasm_v128_store(crow + j, vc);
          }
          for (; j < jmax; j++) crow[j] += aik * brow[j];
#else
          for (int j = jj; j < jmax; j++) crow[j] += aik * brow[j];
#endif
        }
      }
    }
  }
}

EMSCRIPTEN_KEEPALIVE
void matmul_blocked(const double* a, const double* b, double* c, int n) {
  blocked_range(a, b, c, n, 0, n);
}

#ifdef MATMUL_MT
#include <pthread.h>

typedef struct {
  const double* a;
  const double* b;
  double* c;
  int n, i0, i1;
} mm_task;

static void* mm_worker(void* arg) {
  mm_task* t = (mm_task*) arg;
  blocked_range(t->a, t->b, t->c, t->n, t->i0, t->i1);
  return NULL;
}

/* Same blocked+SIMD kernel, but the rows of C are split across nthreads
 * pthreads (reused from emscripten's pthread pool). */
EMSCRIPTEN_KEEPALIVE
void matmul_blocked_mt(const double* a, const double* b, double* c, int n,
                       int nthreads) {
  if (nthreads < 1) nthreads = 1;
  if (nthreads > n) nthreads = n;

  pthread_t th[nthreads];
  mm_task tasks[nthreads];
  int rows = (n + nthreads - 1) / nthreads;
  int nt = 0;

  for (int i = 0; i < nthreads; i++) {
    int i0 = i * rows;
    if (i0 >= n) break;
    int i1 = i0 + rows;
    if (i1 > n) i1 = n;
    tasks[nt] = (mm_task){ a, b, c, n, i0, i1 };
    if (pthread_create(&th[nt], NULL, mm_worker, &tasks[nt]) == 0) {
      nt++;
    } else {
      /* Pool exhausted — run this chunk inline rather than dropping it. */
      blocked_range(a, b, c, n, i0, i1);
    }
  }

  for (int i = 0; i < nt; i++) pthread_join(th[i], NULL);
}
#endif