// The tfjs half — skipped gracefully where the optional native dep is absent.
import test from "node:test";
import assert from "node:assert/strict";
import { mlpInit, mlpForward } from "../dist/forward.js";
import { netToTf, tfToNet, fitNet } from "../dist/tf.js";

let tf = null;
try {
  tf = await import("@tensorflow/tfjs-node");
} catch {
  /* optional peer dep not installed — tests below skip */
}

test("netToTf → tfToNet round-trips the row-major transpose", { skip: !tf }, () => {
  const net = mlpInit([4, 8, 1], (() => {
    let s = 42;
    return () => (s = (s * 16807) % 2147483647) / 2147483647;
  })());
  const back = tfToNet(tf, netToTf(tf, net, 0.001), net.sizes);
  assert.deepEqual(back.sizes, net.sizes);
  for (let l = 0; l < net.layers.length; l++) {
    for (let i = 0; i < net.layers[l].w.length; i++) {
      assert.ok(Math.abs(back.layers[l].w[i] - net.layers[l].w[i]) < 1e-6, `w[${l}][${i}]`);
    }
  }
});

test("fitNet reduces MSE on a toy regression", { skip: !tf }, async () => {
  // learn y = x0 - x1 from scratch
  const xs = [];
  const ys = [];
  for (let i = 0; i < 200; i++) {
    const a = (i % 20) / 10 - 1;
    const b = ((i * 7) % 20) / 10 - 1;
    xs.push([a, b]);
    ys.push(a - b);
  }
  let net = mlpInit([2, 8, 1], (() => {
    let s = 7;
    return () => (s = (s * 16807) % 2147483647) / 2147483647;
  })());
  const mse = (n) => xs.reduce((acc, x, i) => acc + (mlpForward(n, x)[0] - ys[i]) ** 2, 0) / xs.length;
  const before = mse(net);
  net = await fitNet(net, xs, ys, { lr: 0.01, epochs: 30 });
  const after = mse(net);
  assert.ok(after < before / 2, `MSE ${before.toFixed(3)} → ${after.toFixed(3)}`);
});